imagesearch4.py

# -*- coding: utf-8 -*-
"""ImageSearch4.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1mwXBNOTWdYHb7NfKeWBwpdSEJNIe4xkg
"""

# pip install faiss-gpu
!nvidia-smi

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import load_model
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import plot_model
from tensorflow.keras import backend as K
from tensorflow.keras.preprocessing import image
from tensorflow.keras.callbacks import ModelCheckpoint

import pickle
import numpy as np
import matplotlib.pyplot as plt
from imutils import build_montages
import os

from skimage.io import imread, imshow
from skimage.color import rgb2gray, gray2rgb
from skimage.transform import resize

from google.colab.patches import cv2_imshow
import cv2

# import matplotlib
# matplotlib.use('Agg')

from google.colab import drive
drive.mount('/content/drive')

print('[INFO] Tensorflow version ' + tf.__version__)

base_dir='/content/drive/My Drive/ImageSearch4/'
cifar10model= base_dir + 'cifar10.h5'
cifar10histplot= base_dir +  'cifar10histplot.png'
cifar10modelplot= base_dir + 'cifar10modelplot.png'
cifar10index= base_dir + 'cifar10index.pickle'
cifar10chptweights= base_dir + 'cifar10weights.h5'

LATENT_DIM= 1024
WIDTH, HEIGTH, DEPTH= (32, 32, 3)

EPOCHS= 150
INIT_LR= 1e-3
BATCH_SIZE= 32
sample= 10
classes= 10

print('[INFO] loading cifar10 dataset...')
((trainX, trainY), (testX, testY))= cifar10.load_data()
trainX= trainX.astype('float32') / 255.0
trainX= trainX.reshape(trainX.shape)
print('Xtrain shape: ', trainX.shape)

testX= testX.astype('float32') / 255.0
testX= testX.reshape(testX.shape)
print('Xtest shape: ', testX.shape)

def build2(width, heigth, depth, latentDim=16):
    inputs= layers.Input(shape=(width, heigth, depth), name='Inputs')
    x= layers.Conv2D(32, (3, 3), strides=1, padding='same', activation='relu')(inputs)
    x= layers.Conv2D(32, (3, 3), strides=2, padding='same', activation='relu')(inputs)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)
    x= layers.Conv2D(64, (3, 3), strides=1, padding='same', activation='relu')(x)
    x= layers.Conv2D(64, (3, 3), strides=2, padding='same', activation='relu')(x)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)
    x= layers.Conv2D(128, (3, 3), strides=1, padding='same', activation='relu')(x)
    x= layers.Conv2D(128, (3, 3), strides=2, padding='same', activation='relu')(x)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)

    volumeSize= K.int_shape(x)
    x= layers.Flatten()(x)
    latent= layers.Dense(latentDim, name='Encoded')(x)

    x= layers.Dense(np.prod(volumeSize[1:]))(latent)
    x= layers.Reshape((volumeSize[1], volumeSize[2], volumeSize[3]))(x)

    x= layers.Conv2DTranspose(128, (3, 3), strides=2, padding='same', activation='relu')(x)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)
    x= layers.Conv2DTranspose(64, (3, 3), strides=2, padding='same', activation='relu')(x)
    x= layers.Conv2DTranspose(32, (3, 3), strides=2, padding='same', activation='relu')(x)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)

    x= layers.Conv2DTranspose(depth, (3, 3), strides=1, padding='same')(x)
    outputs= layers.Activation('sigmoid', name='Decoded')(x)

    x= layers.Dense(128, activation='relu')(latent)
    x= layers.BatchNormalization()(x)
    x= layers.Dropout(0.25)(x)
    Classifier= layers.Dense(classes, activation='softmax', name='Classifier')(x)


    model= Model(inputs=inputs, outputs=[outputs, Classifier])

    return model
    
# ===============================

def visualize_predictions(decoded, gt, samples=10):
    outputs=None
    for i in range(0, samples):
        original= (gt[i] * 255).astype('uint8')
        recon= (decoded[i] * 255).astype('uint8')
        output= np.hstack([original, recon])
        
        if outputs is None:
            outputs= output
        else:
            outputs= np.vstack([outputs, output])
            
    return outputs

print('[INFO] building autoencoder...')
# tpu= tf.distribute.cluster_resolver.TPUClusterResolver()
# print('[INFO] Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
# tf.config.experimental_connect_to_cluster(tpu)
# tf.tpu.experimental.initialize_tpu_system(tpu)
# strategy= tf.distribute.experimental.TPUStrategy(tpu)
# print('[INFO] REPLICAS: ', strategy.num_replicas_in_sync)

# with strategy.scope():
autoencoder= build2(WIDTH, HEIGTH, DEPTH, latentDim=LATENT_DIM)
plot_model(autoencoder, to_file=cifar10modelplot,
           show_shapes=True)
autoencoder.summary()

# opt= Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
autoencoder.compile(loss=['mse', 'sparse_categorical_crossentropy'], optimizer='adam',
                    metrics=['accuracy'])

checkpointer= ModelCheckpoint(filepath=cifar10chptweights,
                               verbose=0, save_best_only=True)
H= autoencoder.fit(trainX, [trainX,trainY],
                    validation_split=0.1,
					epochs=EPOCHS*4,
					batch_size=BATCH_SIZE,
					verbose=2,
					callbacks=[checkpointer])

print('[INFO] saving autoencoder...')
autoencoder.save(cifar10model, save_format='h5')

print('[INFO] making predictions...')
decoded= autoencoder.predict(testX)[0]
vis= visualize_predictions(decoded, testX)
cv2_imshow(vis)

print('[INFO] makeing autoencoder plot ...')
N= np.arange(0, EPOCHS)
plt.style.use('ggplot')
fig= plt.figure(figsize=(8,10))
axs= fig.subplots(nrows=2, ncols=1)
axs[0].plot(N, H.history['Decoded_loss'], label='Decoded_loss', color='g')
axs[0].plot(N, H.history['val_Decoded_loss'], label='val_Decoded_loss', color='y')
axs[1].plot(N, H.history['Classifier_accuracy'], label='Classifier_accuracy', color='tab:pink')
axs[1].plot(N, H.history['val_Classifier_accuracy'], label='val_Classifier_accuracy', color='tab:brown')
axs[1].set_xlabel('Epoch #')
axs[1].set_ylabel('Loss/Accuracy')
axs[0].legend(loc='upper right')
axs[1].legend(loc='upper right')

# print('[INFO] loading cifar10 training split...')
# ((trainX, _), (testX, _))= cifar10.load_data()

# trainX= trainX.astype('float32') / 255.0

print('[INFO] loading autoencoder model...')
autoencoder= load_model(cifar10model)

encoder= Model(inputs=autoencoder.input,
                outputs=autoencoder.get_layer('Encoded').output)

print('[INFO] encoding images...')
features= encoder.predict(trainX)

indexes= list(range(0, trainX.shape[0]))
data= {'indexes': indexes, 'features': features}

print('[INFO] saving index...')
f= open(cifar10index, 'wb')
f.write(pickle.dumps(data))
f.close()

def euclidean(a, b):
    diff= np.linalg.norm(a - b)
    return diff

def perform_search(queryFeatures, index, maxResults=64):
	results= []

	for i in range(0, len(index['features'])):
		d= euclidean(queryFeatures, index['features'][i])
		results.append((d, i))

	results= sorted(results)[:maxResults]

	return results

# print('[INFO] loading cifar10 dataset...')
# ((trainX, _), (testX, _))= cifar10.load_data()

# trainX= trainX.astype('float32') / 255.0
# testX= testX.astype('float32') / 255.0

print('[INFO] loading autoencoder and index...')
autoencoder= load_model(cifar10model)
index= pickle.loads(open(cifar10index, 'rb').read())

encoder= Model(inputs=autoencoder.input,
                outputs=autoencoder.get_layer('Encoded').output)

print('[INFO] encoding testing images...')
features= encoder.predict(testX)

queryIdxs= list(range(0, testX.shape[0]))
queryIdxs= np.random.choice(queryIdxs, size=sample,
                             replace=False)

for i in queryIdxs:
	queryFeatures= features[i]
	results= perform_search(queryFeatures, index, 64)
	images= []

	for (d, j) in results:
		image= (trainX[j] * 255).astype('uint8')
		images.append(image)

	query= (testX[i] * 255).astype('uint8')
	cv2_imshow(cv2.resize(query, (64,64)))

	montage= build_montages(images, (64, 64), (3, 3))[0]
	cv2_imshow(montage)