import gradio as gr
import numpy as np
from PIL import Image
import cv2
from inference import FakeImageDetector

# Initialize detector
detector = FakeImageDetector()

def predict_image(image):
    """
    Predict if image is real or fake
    """
    try:
        # Convert Gradio image to PIL
        if isinstance(image, np.ndarray):
            image_pil = Image.fromarray(image)
        else:
            image_pil = image
        
        # Make prediction
        result = detector.predict(image_pil, return_confidence=True)
        
        if result is None:
            return "Error processing image", None
        
        # Prepare output
        label = result['prediction']
        confidence = result['confidence']
        
        # Create visualization
        display_img = np.array(image_pil)
        display_img = cv2.resize(display_img, (400, 400))
        
        # Add text overlay
        color = (0, 255, 0) if label == 'Real' else (255, 0, 0)  # Green for real, red for fake
        cv2.putText(display_img, f"{label}: {confidence:.1f}%", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
        
        return f"Prediction: {label} (Confidence: {confidence:.1f}%)", display_img
    
    except Exception as e:
        return f"Error: {str(e)}", None

# Create Gradio interface
title = "Fake Image Detector"
description = "Upload an image to detect if it's real or AI-generated/fake"
examples = [["real_0.jpg"], ["fake_0.jpg"]]  # Add your example images

iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Image(label="Visualization", type="numpy")
    ],
    title=title,
    description=description,
    examples=examples,
    theme="soft"
)

if __name__ == "__main__":
    iface.launch(share=True)