import gradio as gr
import numpy as np
import os
import zipfile
from transformers import pipeline

with zipfile.ZipFile("examples.zip","r") as zip_ref:
    zip_ref.extractall(".")

EXAMPLES_FOLDER = 'examples'
examples_names = os.listdir(EXAMPLES_FOLDER)
examples = []
for example_name in examples_names:
    example_path = os.path.join(EXAMPLES_FOLDER, example_name)
    label = example_name.split('_')[0]
    example = {
        'path': example_path,
        'label': label
    }
    examples.append(example)
np.random.shuffle(examples)

# Load the Hugging Face model
model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection", device="cuda")

# Define the prediction function
def classify_image(image):
    results = model(image)
    return {result['label']: result['score'] for result in results}

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),  # Accepts images in PIL format
    outputs=gr.Label(num_top_classes=2),  # Displays top two classifications with probabilities
    title="Deepfake vs Real Image Detection",
    description="Upload an image to determine whether it's real or deepfake.",
    examples=[[examples[i]["path"], examples[i]["label"]] for i in range(10)]
)

# Launch the Gradio app
interface.launch()