#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 14 11:44:05 2021

@author: valentinahrtonova
"""
from numpy import expand_dims, asarray
from numpy.random import randn
from matplotlib import pyplot as plt
import cv2

from keras.models import load_model

def generate_latent_points(latent_dim, n_samples):
	# generate points in the latent space
	x_input = randn(latent_dim * n_samples)
	# reshape into a batch of inputs for the network
	x_input = x_input.reshape(n_samples, latent_dim)
	return x_input

# GENERATE MASKS

model = load_model('DCGAN_melanoma_model_010.h5')

latent_dim = 100
n_samples = 1

melanoma_masks = list ()

for i in range(10):
    
    x_input = generate_latent_points(latent_dim, n_samples)
    examples = model.predict(x_input)
    examples = (examples + 1) / 2.0
    examples = examples[0,:,:,:]
    
    _, image_b = cv2.threshold(src=examples, thresh=0.5, maxval=255, type=cv2.THRESH_BINARY)
    
    melanoma_masks.append(image_b)

melanoma_masks = asarray(melanoma_masks)

# GENERATE MELANOMA SAMPLES

model2 = load_model('pix2pix_melanoma_model_200.h5')

X1 = (melanoma_masks - 127.5) / 127.5

for i in range(X1.shape[0]):
    
    plt.subplot(1,2,1)
    plt.imshow(X1[i])
    plt.title('Mask generated by DCGAN')
    plt.axis('off')

    X = X1[i]
    X = expand_dims(X, axis=0)
    
    fakeB = model2.predict(X)
    
    fakeB = fakeB[0,:,:,:]
    fakeB = (fakeB + 1) / 2.0
    
    plt.subplot(1,2,2)
    plt.imshow(fakeB)
    plt.title('Output generated by pix2pix')
    plt.axis('off')
    
    filename = 'demo_output/generated_%02d.png' % (i)
    plt.savefig(filename)


    