# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="3PturmDN40NY"
import numpy as np
import os

from time import time
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

from keras.preprocessing.image import load_img, array_to_img 
from keras.models import Model
from keras.layers import Conv2D, SeparableConv2D
from keras.layers import BatchNormalization, Activation, advanced_activations
from keras.layers import Input, MaxPooling2D, Add
from keras.layers import Conv2DTranspose, UpSampling2D

# + id="tLtNuCH7l2rx"
TRAIN_MODE = 1
EPOCH = 1000
BATCH_SIZE = 32

# + id="r25FCPYv41Kd"
  with open('/content/drive/MyDrive/IDEC/arrays.npy', 'rb') as f:
      X_data = np.load(f)
      Y_data = np.load(f)

# + colab={"base_uri": "https://localhost:8080/"} id="bf3al1tr45EE" outputId="75636429-9c67-4db6-b86d-ccb9f8070107"
X_train, X_test, Y_train, Y_test = train_test_split(X_data,
                                                    Y_data,
                                                    test_size = 0.2)

print("Train Input Data :", X_train.shape)
print("Train Output Data :", Y_train.shape)
print("Test Input Data :", X_test.shape)
print("Test Output Data :", Y_test.shape)

# + colab={"base_uri": "https://localhost:8080/", "height": 449} id="KOFvl5al5EW0" outputId="ef65fe73-21c1-45f1-b950-0b903b6aa05c"
# Image show after Scaling
flg, spot = plt.subplots(1,2,figsize=(15,10)) 

# Color Image
img = array_to_img(X_data[1154])
spot[0].imshow(img)

# Edge Image
img = array_to_img(Y_data[1154])
spot[1].imshow(img)


# + [markdown] id="x5fQo4xo69Sd"
# Design Network

# + id="tKMPgs7X5Ewd"
# U-Net Encoder
def build_encoder():
    inputs = Input(shape=(160, 160, 3))

    x = Conv2D(filters=32,
               kernel_size=3,
               strides=2,
               padding='same')(inputs)

    x = BatchNormalization()(x) # Normalizatoin : mean = 0, deviation = 1
    x = Activation('relu')(x)
    
    jump = x

    for filters in [64, 128, 256]:
        #x = advanced_activations.LeakyReLU(alpha=0.2)(x)

        # Conv1
        x = Activation('relu')(x)
        x = SeparableConv2D(filters,
                            kernel_size=3,
                            padding='same')(x)
        x = BatchNormalization()(x)

        # Conv2
        x = Activation('relu')(x)
        x = SeparableConv2D(filters,
                            kernel_size=3,
                            padding='same')(x)
        x = BatchNormalization()(x)

        # Pooling
        x = MaxPooling2D(pool_size=3,
                         strides=2,
                         padding='same')(x)
        
        # Residual
        residual = Conv2D(filters,
                          kernel_size=1,
                          strides=2,
                          padding='same')(jump)

        x = Add()([x, residual])
        jump = x

    return inputs, x



# + id="JMb-WpfR7ERb"
# U-Net Decoder
def build_decoder(inputs, x):
    # Residual
    jump = x

    # De-Conv
    for filters in [256, 128, 64, 32]:
        # Conv1
        x = Activation('relu')(x)
        x = Conv2DTranspose(filters,
                            kernel_size=3,
                            padding='same')(x)
        x = BatchNormalization()(x)
    
        # Conv2
        x = Activation('relu')(x)
        x = Conv2DTranspose(filters,
                            kernel_size=3,
                            padding='same')(x)
        x = BatchNormalization()(x)
        x = UpSampling2D(size=2)(x)

        # Residual
        residual = UpSampling2D(size=2)(jump)
        residual = Conv2D(filters,
                          kernel_size=1,
                          padding='same')(residual)

        x = Add()([x, residual])
        jump = x
        
    outputs = Conv2D(filters=3,
                     kernel_size=3,
                     activation='softmax',
                     padding='same')(x)

    model = Model(inputs, outputs)

    return model


# + id="TOt41x0LDjOR"
inputs, link = build_encoder()
model = build_decoder(inputs, link)

# + [markdown] id="b3ovK_gdoN7_"
# Train Model

# + colab={"base_uri": "https://localhost:8080/"} id="Swza4Pa0fmVG" outputId="15cc9d0b-1781-4b4d-819d-e6e1772770fd"
# Train Model
if TRAIN_MODE:
    model.compile(optimizer='rmsprop',
                  loss='sparse_categorical_crossentropy')
    
    print("Start Train\n")
    begin = time()
    model.fit(X_train, Y_train, BATCH_SIZE, EPOCH, verbose=1)
    end = time()

    print("Learning TIme : {:.2f}".format(end-begin))
    model.save_weights('segment.h5')
else:
    model.load_weights('segment.h5')

# + [markdown] id="oevjOqGHopKu"
# Test Model

# + colab={"base_uri": "https://localhost:8080/", "height": 333} id="dLXTw40OoEMe" outputId="3dc0b764-8dc9-4180-c0df-0715474117f8"
which = 232

fig, spot = plt.subplots(1, 3, figsize=(15,8))

img = array_to_img(X_test[which])
spot[0].imshow(img)

img = array_to_img(Y_test[which])
spot[1].imshow(img)

pred = model.predict(X_test, verbose=1)
mask = np.argmax(pred[which], axis=2)
mask = np.expand_dims(mask, axis=2)
img = array_to_img(mask)
spot[2].imshow(img)

StarcoderdataJupyterScriptsDedupFiltered