import gradio as gr
import xgboost as xgb
import numpy as np
from PIL import Image
from image_functions import extract_features_for_ml

# Load XGBoost model (sklearn wrapper)
model = xgb.XGBClassifier()
model.load_model("XGBoost_Artifact_OpenFake.json")

def predict(image):
    """
    Classify image as real or fake.
    
    Args:
        image: PIL Image from Gradio
        threshold: Classification threshold
    
    Returns:
        Dictionary with probabilities for Gradio Label component
    """
    try:
        # Convert PIL Image to numpy array (RGB)
        img_array = np.array(image.convert("RGB"))
        
        # Extract features using your pipeline
        features = extract_features_for_ml(img_array)
        
        # Handle NaN values
        features = np.nan_to_num(features, nan=0.0)
        
        # Reshape for prediction
        X = features.reshape(1, -1)
        
        # Get prediction probabilities
        y_proba = model.predict_proba(X)[0]
        
        return {
            "Real": float(y_proba[0]),
            "Fake": float(y_proba[1])
        }
        
    except Exception as e:
        print(f"Error during prediction: {e}")
        return {"Error": 1.0}

# Create Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
    ],
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="Image Fake Detector",
    description="Upload an image to classify it as real or fake using image statistics with XGBoost."
)

demo.launch(show_error=True)