import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import img_to_array

# Load the pre-trained model
model = load_model('deepfake_detection_mobilenet_model.h5')

# Define the function for prediction
def predict_image(img):
    # Resize and preprocess the input image
    x = cv2.resize(img, (224, 224))
    x = img_to_array(x) / 255.0  # Normalize the image
    x = np.expand_dims(x, axis=0)  # Add batch dimension
    
    # Predict with the model
    prediction = (model.predict(x) > 0.5).astype("int32")[0][0]
    
    # Return result based on the prediction
    if prediction == 1:
        return "Fake Image"
    else:
        return "Real Image"

# Define the Gradio Interface
description_html = """
<p>Upload a face image to check if it's real or morphed with deepfake</p>
"""

custom_css = """
div {background-color: whitesmoke;}
"""

# Create the Gradio app interface
gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy", label="Upload Face Image"),
    outputs=gr.Textbox(label="Prediction"),
    title="Deepfake Image Detection",
    description=description_html,
    allow_flagging='never'
).launch()