from keras.applications.vgg19 import VGG19
import keras.backend as K
from keras.models import Model


    
class VGG_MSE(object):

    def __init__(self, image_shape):
        
        self.image_shape = image_shape

    # computes VGG loss or content loss
    # def vgg_mse(self, y_true, y_pred):
    def loss(self, y_true, y_pred):
        print(y_true.shape, y_pred.shape)
        vgg19 = VGG19(include_top=False, weights='imagenet', input_shape=self.image_shape)
        vgg19.trainable = False
        # Make trainable as False
        for l in vgg19.layers:
            l.trainable = False
        model = Model(inputs=vgg19.input, outputs=vgg19.get_layer('block5_conv4').output)
        # model.build(input_shape=(10, self.image_shape))
        model.trainable = False
    
        # return K.mean(K.square(model(y_true) - model(y_pred)))
        return K.mean(K.square(model(y_true) - model(y_pred))) + (K.mean(K.square(y_true - y_pred)))*0.5
