import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np
from huggingface_hub import from_pretrained_keras

# Load the models
model1 = from_pretrained_keras("arsath-sm/face_classification_model1")
model2 = from_pretrained_keras("arsath-sm/face_classification_model2")

# Preprocess the image
def preprocess_image(img):
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img = img / 255.0
    return img

# Make predictions
def predict(img):
    preprocessed_img = preprocess_image(img)
    
    prediction1 = model1.predict(preprocessed_img)[0][0]
    prediction2 = model2.predict(preprocessed_img)[0][0]
    
    result1 = "Real" if prediction1 > 0.5 else "Fake"
    result2 = "Real" if prediction2 > 0.5 else "Fake"
    
    confidence1 = prediction1 if result1 == "Real" else 1 - prediction1
    confidence2 = prediction2 if result2 == "Real" else 1 - prediction2
    
    return {
        "Model 1 Prediction": f"{result1} (Confidence: {confidence1:.2f})",
        "Model 2 Prediction": f"{result2} (Confidence: {confidence2:.2f})"
    }

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs={
        "Model 1 Prediction": gr.Textbox(),
        "Model 2 Prediction": gr.Textbox()
    },
    title="Real vs AI Face Classification",
    description="Upload an image to classify whether it's a real face or an AI-generated face using two different models."
)

# Launch the app
iface.launch()