
# Code Author: Kunal Suthar
# Preprocessing file for performing Out-painting for SML project

from segmentation_models import Unet, Nestnet, Xnet

import numpy as np
from keras import backend as K
from keras.models import Model
from keras.layers import Flatten, Dense, Dropout
from keras.layers import Conv2D
from keras.applications.resnet50 import ResNet50, preprocess_input
from keras.optimizers import Adam
from keras.preprocessing.image import ImageDataGenerator
import os
import sys
from PIL import Image
from skimage.color import rgb2gray
import numpy as np
from matplotlib import pyplot as plt
import cv2

DATASET_PATH='/home/kunal/Feature_Learning_for_Image_Classification/Temp_Dataset'
IMAGE_SIZE    = (256, 256)
CROP_LENGTH   = 836
NUM_CLASSES   = 2
BATCH_SIZE    = 54  # try reducing batch size or freeze more layers if your GPU runs out of memory
FREEZE_LAYERS = 2  # freeze the first this many layers for training
NUM_EPOCHS    = 20
WEIGHTS_FINAL = 'model-cropped-final.h5'


def random_crop(img, random_crop_size):
    # Note: image_data_format is 'channel_last'

    height, width = img.shape[0], img.shape[1]
    dy, dx = random_crop_size
    x = np.random.randint(0, width - dx + 1)
    y = np.random.randint(0, height - dy + 1)
    img=img[y:(y+dy), x:(x+dx), :]
    img=img/255
    img = cv2.resize(img,(224,224))
    

    return img


def crop(img, random_crop_size):
    # Note: image_data_format is 'channel_last'
    # assert img.shape[2] == 3
    
    
    height, width = img.shape[0], img.shape[1]
    dy0, dx0 = 836,836
    x0 = 94
    y0 = 45
    img=img[y0:(y0+dy0), x0:(x0+dx0), :]
    img=img/255
    img = cv2.resize(img,(224,224))
    
    return img

def crop_generator(batches, crop_length):#224
    """Take as input a Keras ImageGen (Iterator) and generate random
    crops from the image batches generated by the original iterator.
    """
    while True:
        batch_x= next(batches)
        
        batch_crops_inp = np.zeros((batch_x.shape[0], 224, 224,3))#224
        batch_crops_tar = np.zeros((batch_x.shape[0], 224, 224,3))
        for i in range(batch_x.shape[0]):
            batch_crops_inp[i] = random_crop(batch_x[i], (crop_length, crop_length))
            batch_crops_tar[i] = crop(batch_x[i], (crop_length, crop_length))
        
        yield (batch_crops_inp,batch_crops_tar)


def main():
	
	train_datagen = ImageDataGenerator(rescale=1/255)
	train_batches = train_datagen.flow_from_directory(DATASET_PATH,
                                                  target_size=(1024,1024),
                                                  shuffle=True,
                                                  class_mode=None,
                                                  batch_size=BATCH_SIZE)

	valid_datagen = ImageDataGenerator(rescale=1/255)
	valid_batches = valid_datagen.flow_from_directory(DATASET_PATH ,
    	                                              target_size=(1024,1024),
    	                                              shuffle=False,
        	                                          class_mode=None,
        	                                          batch_size=BATCH_SIZE)
	
	train_crops = crop_generator(train_batches, CROP_LENGTH) #224
	valid_crops = crop_generator(valid_batches, CROP_LENGTH)


	batch_x_random_crop, batch_y_targeted_crop = next(train_crops)
	valid_x, valid_y = next(valid_crops)
	
	in_painted_x= in_painting_mask(batch_x_random_crop,batch_y_targeted_crop)
	valid_in_x=in_painting_mask(valid_x,valid_y)

	batch_x_random_crop=rgb2gray(batch_x_random_crop)
	batch_x_random_crop=np.reshape(batch_x_random_crop,(batch_x_random_crop.shape[0],224,224,1))
	
	valid_x=rgb2gray(valid_x)
	valid_x=np.reshape(valid_x,(valid_x.shape[0],224,224,1))

	
	model = Unet(backbone_name='resnet18', encoder_weights='imagenet', decoder_block_type='transpose') # build U-Net
	model.compile(optimizer='Adam', loss='mean_squared_error')
	model.summary()
	model.fit(x=in_painted_x,y=batch_x_random_crop,validation_data= (valid_in_x,valid_x),validation_steps=5, steps_per_epoch=5,epochs=1)

	


def in_painting_mask(batch_x,batch_y):

	
	#Creating random mask dimensions for inpainting
	in_paint_x=np.zeros((batch_x.shape[0],224,224,3))
	
	width=224
	height=224

	for i in range(0,batch_x.shape[0]):
		mask=np.zeros((224,224,3))

		#random size of mask
		dx= np.random.randint(0,224+1)
		dy= np.random.randint(0,224+1)

		#random location of mask
		x = np.random.randint(0, 224 + 1)
		y = np.random.randint(0, 224 + 1)

		x_plus_dx=min(x+dx,224)
		y_plus_dy=min(y+dy,224)

		mask[x:(x_plus_dx),y:(y_plus_dy),:]=1
		
		
		in_paint_x[i]=np.multiply(batch_x[i],mask)

		

	return in_paint_x










if __name__=="__main__":
	main()




