import torch
import numpy as np
import gradio as gr
from PIL import Image
import multiprocessing
import tensorflow as tf
from RealESRGAN import RealESRGAN
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications import InceptionV3


# function to return model
def create_model():
    inp_shape = (200,200,3)
    base_model = InceptionV3(input_shape=inp_shape, include_top=False, weights='imagenet')
    x = layers.Flatten()(base_model.output)
    x = layers.Dense(256, activation='relu')(x)
    x = layers.Dropout(0.5)(x)

    output = layers.Dense(8, activation='softmax')(x)

    clf_model = Model(inputs=base_model.input, outputs=output)

    clf_model.compile(optimizer=Adam(learning_rate=0.0001),
                      loss='categorical_crossentropy',
                      metrics=['accuracy'])
    return clf_model

clf_model = create_model()
clf_model.load_weights('modelac90.weights.h5')

# input image preprocessing and prediction
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
esrgan_model = RealESRGAN(device, scale=4)
esrgan_model.load_weights('weights/RealESRGAN_x4.pth')
enhancement_model = esrgan_model

def classify_logo(inp_image):
    pil_img = Image.fromarray(inp_image).resize((200,200), resample=0)
    image = np.array(pil_img).astype(np.float16)/255.0
    new_img = np.expand_dims(image, axis=0)
    predictions = clf_model.predict(new_img)
    labels = ['Adidas Fake', 'Adidas Real', 'Allen Solly Fake', 'Allen Solly Real',
              'Puma Fake', 'Puma Real', 'Us Polo Fake', 'Us Polo Real']
    pred_dict = {}
    for i in range(len(labels)):
        pred_dict[labels[i]] = predictions[0][i]
    return pred_dict

enhaced_images_pipe = []
def enhance_image(input_image):
    try:
        enhanced_image = enhancement_model.predict(input_image).convert('RGB')
        global enhanced_images_pipe
        enhaced_images_pipe.append(np.array(enhanced_image))
        return
    except Exception as e:
        return None

def fake_logo_detection(input_image, flag):
    print("Input image shape => ", input_image.shape)
    print("flag => ", flag)
    if flag == "Enhance Image before prediction (Uses Real ESRGAN 4X)":
        process = multiprocessing.Process(target=enhance_image, args=(input_image,))
        process.start()
        process.join(500)
        if process.is_alive():
            process.terminate()
            process.join()
            return "Enhancement model prediction taking too long time⌛. Please! try without enhancement."
        else:
            global enhanced_images_pipe
            if len(enhaced_images_pipe) != 0:
                input_image = enhaced_images_pipe[-1]
                arr = enhaced_images_pipe.pop()
            else:
                return "An error occurred while enhancing the image 😞."
    return classify_logo(input_image)

# --------------gradio interface-------------------------
    
input_choices = gr.Dropdown(["No enhancement (prefer for Medium-High resolution images)", "Enhance Image before prediction (Uses Real ESRGAN 4X)"],
                            value="No enhancement (prefer for Medium-High resolution images)", label="Choose Option")
input_image = gr.Image(sources=['upload'], label="Input Image")

iface = gr.Interface(
    fn = fake_logo_detection,
    inputs = [input_image, input_choices],
    outputs=gr.Label(num_top_classes=2, label="Confidence percentages"),
    title="Fake Logo Detection",
    description="Upload an image to check if it contains a fake or original logo."
)

if __name__ == '__main__':
    iface.launch(debug=False)