import gradio as gr
from model.load_model import load_model
from utils.video_utils import extract_frames
from utils.face_utils import extract_faces
from predictor.predict import predict_faces
from utils.gradcam import get_gradcam, get_conv_layers
import numpy as np
from PIL import Image
from tqdm import tqdm

model = load_model()
conv_layer_names = get_conv_layers(model)  # populate dropdown choices

def deepfake_app(video, selected_layer, progress=gr.Progress(track_tqdm=True)):
    frames = extract_frames(video)
    frames = list(frames)

    faces = extract_faces(frames)
    faces = list(faces)

    if not faces:
        return "No face detected", None

    predictions = predict_faces(model, faces)
    predictions = list(predictions)

    avg_score = np.mean(predictions)
    label = "FAKE" if avg_score > 0.5 else "REAL"

    max_idx = np.argmax(predictions)
    cam_image = get_gradcam(model, faces[max_idx], selected_layer)
    cam_image = Image.fromarray(cam_image)

    return label, cam_image

gr.Interface(
    fn=deepfake_app,
    inputs=[gr.Video(label="Upload a Video"),
            gr.Dropdown(choices=conv_layer_names, label="Model target layer", value=conv_layer_names[-1])
        ],
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Image(label="RAI Explainability")
        ],
    title="Deepfake video detection system",
    description="Upload a video, and the model will predict whether its a deepfake content, with RAI features."
).launch()