# ---
# 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="5sdf6Jn0Gj-a" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34.0} outputId="3441d973-4809-4805-8a7c-f6336214c62c"
from google.colab import drive
from PIL import Image
import os
from keras import layers
from keras.models import Model, Sequential
from keras.layers import Dense, Conv2D, Flatten, BatchNormalization, Layer, UpSampling2D, Activation, Input
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
import numpy as np
from torchvision import transforms
import matplotlib.pyplot as plt
from keras import backend as K
from keras import losses
from keras import optimizers
import tensorflow as tf

# + id="HE-7VsJJKxmZ" colab_type="code" colab={}
drive.mount('/content/gdrive')


# + id="ugDFI5dIGvCk" colab_type="code" colab={}
# !tar xvf "../content/gdrive/My Drive/glimpse_dataset/coco.tar"

# + id="StcWty5UG697" colab_type="code" colab={}
def dataset_generator(batch_size):
  heatmaps_names = os.listdir("coco/heatmaps/")
  resize_240 = transforms.Resize(240)
  resize_30 = transforms.Resize(30)
  
  last_index = 0
  
  while True:
    heatmaps = []
    images = []
  
    for heatmap_name in heatmaps_names[last_index:last_index+batch_size]:
      heatmap = Image.open("coco/heatmaps/" + heatmap_name)
      heatmap = resize_30(heatmap)
      heatmaps.append(img_to_array(heatmap) / 255)

      image = Image.open("coco/images/" + heatmap_name).convert('RGB')   
      image = resize_240(image)
      images.append(img_to_array(image) / 255)

    last_index += batch_size 
    
    if last_index >= len(heatmaps_names):
      last_index = 0
      
    yield np.array(images), np.array(heatmaps)


# + id="21ZVHq1IIzfn" colab_type="code" outputId="dafea71c-4cac-4ebf-ae0d-fed2cdc49136" colab={"base_uri": "https://localhost:8080/", "height": 34.0}
images, heatmaps = next(iter(dataset_generator(8)))
images.shape, heatmaps.shape


# + id="Q0Vl0qUAyGFd" colab_type="code" colab={}
def residual_block(output, filters):
  out = output
  
  output = Conv2D(filters, kernel_size=3, padding='same', strides=1)(output)
  output = BatchNormalization()(output)
  output = Activation('relu')(output)

  output = Conv2D(filters, kernel_size=3, padding='same', strides=1)(output)
  output = BatchNormalization()(output)
  
  output = layers.add([out, output])

  return output


# + id="Ll9g15UYtFEE" colab_type="code" outputId="f71afe22-18bd-4768-97b2-14a792951307" colab={"base_uri": "https://localhost:8080/", "height": 1227.0}
visible = Input(shape=(240, 320, 3))

output = Conv2D(8, kernel_size=9, padding='same', strides=4)(visible)
output = BatchNormalization()(output)
output = Activation('relu')(output)

output = Conv2D(16, kernel_size=3, padding='same', strides=2)(output)
output = BatchNormalization()(output)
output = Activation('relu')(output)

output = Conv2D(32, kernel_size=3, padding='same', strides=2)(output)
output = BatchNormalization()(output)
output = Activation('relu')(output)

output = residual_block(output, 32)
output = Conv2D(64, kernel_size=1, padding='same', strides=1, bias=False)(output)
output = residual_block(output, 64)

output = Conv2D(16, kernel_size=3, padding='same', strides=1)(output)
output = BatchNormalization()(output)
output = Activation('relu')(output)
  
output = UpSampling2D(2)(output)

output = Conv2D(1, kernel_size=9, activation='sigmoid', padding='same')(output)

model = Model(inputs=visible, outputs=output)
model.compile(optimizer=optimizers.Adam(lr=0.001), loss='mse', metrics=['accuracy'])

model.summary()

# + id="SToz3cwXLZzZ" colab_type="code" outputId="b546e231-b0c7-4365-90ad-d68e12cc6cd9" colab={"base_uri": "https://localhost:8080/", "height": 258.0}
report = model.fit_generator(dataset_generator(8), steps_per_epoch=1875, epochs=5)

# + id="PidN5UOOLfgI" colab_type="code" outputId="0c06d8fa-38e9-45fe-c8e1-e3c22cd08d78" colab={"base_uri": "https://localhost:8080/", "height": 376.0}
plt.plot(report.history['loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

# + id="KGCbMNnfOCZ-" colab_type="code" colab={}
data_gen = dataset_generator(1)

# + id="sZBwawJGNmn6" colab_type="code" outputId="759d9952-c2b6-449b-eed4-bd436653f9ac" colab={"base_uri": "https://localhost:8080/", "height": 1010.0}
images, heatmaps = next(iter(data_gen))


plt.imshow(images[0])
plt.show()

true = K.softmax(heatmaps[0].reshape(-1))
true = K.eval(true).reshape(30,40) 
plt.imshow(true, cmap='Greys_r')
plt.show()

output = K.softmax(model.predict(images)[0].reshape(-1))
output = K.eval(output).reshape(30,40) 
plt.imshow(output, cmap='Greys_r')
plt.show()

# + id="lzws7QsfQCE1" colab_type="code" colab={}
model.save('gdrive/My Drive/k_saliency_model.h5') 

# + id="d2eoiVFd1JUk" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered