from fastai.vision.all import *
import gradio as gr
from PIL import Image, ImageChops, ImageEnhance
import cv2

#Converting the images to ELA format for prediction
def convert_to_ela_image(img_array, quality):
    resaved_filename = 'temp.jpg'
    img = Image.fromarray(img_array)  
    im = img.convert('RGB')
    im.save(resaved_filename, 'JPEG', quality=quality)
    resaved_im = Image.open(resaved_filename)
    
    ela_im = ImageChops.difference(im, resaved_im)
    
    extrema = ela_im.getextrema()
    max_diff = max([ex[1] for ex in extrema])
    if max_diff == 0:
        max_diff = 1
    scale = 255.0 / max_diff
    
    ela_im = ImageEnhance.Brightness(ela_im).enhance(scale)
    ela_im = ela_im.resize((256, 256))
    return ela_im

#Importing the model
learn = load_learner('resnet34_ela_images.pkl')

categories = ('Fake', 'Real')

#Defining the function to get the prediction values
def classify_image(img):
    im = convert_to_ela_image(img, 90)
    pred,idx,probs = learn.predict(im)
    return dict(zip(categories, map(float,probs)))

image = gr.Image()
label = gr.Label()
examples = ['lion_f.png', 'boat_r.jpg', 'potato_f.jpg']

intf = gr.Interface(fn=classify_image, inputs=image, outputs = label, examples = examples)
intf.launch(inline=False)