import gradio as gr
import torch
from model import CNNLSTMClassifier
from utils import extract_frames

model = CNNLSTMClassifier()
model.load_state_dict(torch.load("lbw_classifier.pt", map_location='cpu'))
model.eval()

classes = ["Not LBW", "LBW"]

def predict(video):
    frames = extract_frames(video)
    with torch.no_grad():
        output = model(frames)
        pred = torch.argmax(output, dim=1).item()
        prob = torch.softmax(output, dim=1)[0][pred].item()
    return f"Prediction: {classes[pred]} (Confidence: {prob:.2%})"

iface = gr.Interface(
    fn=predict,
    inputs=gr.Video(type="filepath"),
    outputs=gr.Text(),
    title="Smart LBW Classifier",
    description="Upload a cricket video. The AI model will predict whether it's an LBW or not."
)

iface.launch()