import gradio as gr
import numpy as np
import cv2
from keras.models import load_model

# Load the trained model
model = load_model("bullying_detection_model.keras")

# Constants (should match your training)
IMAGE_HEIGHT, IMAGE_WIDTH = 64, 64
SEQUENCE_LENGTH = 16
CLASSES_LIST = ["NonBullying", "Bullying"]

# Helper to preprocess frames
def preprocess_frames(frames):
    processed = [cv2.resize(f, (IMAGE_HEIGHT, IMAGE_WIDTH))/255.0 for f in frames]
    return np.array(processed)

# Gradio function: receives a video, returns prediction
def predict_bullying(video):
    cap = cv2.VideoCapture(video)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)
    cap.release()
    if len(frames) < SEQUENCE_LENGTH:
        return "Video too short"
    # Sample SEQUENCE_LENGTH evenly spaced frames
    idxs = np.linspace(0, len(frames)-1, SEQUENCE_LENGTH).astype(int)
    sampled_frames = [frames[i] for i in idxs]
    input_frames = preprocess_frames(sampled_frames)
    input_frames = np.expand_dims(input_frames, axis=0)
    preds = model.predict(input_frames)[0]
    pred_idx = np.argmax(preds)
    pred_class = CLASSES_LIST[pred_idx]
    confidence = float(preds[pred_idx])
    return f"Prediction: {pred_class} (Confidence: {confidence:.2f})"

iface = gr.Interface(
    fn=predict_bullying,
    inputs=gr.Video(sources=["webcam", "upload"], label="Input Video"),
    outputs="text",
    title="Bullying Detection in Video",
    description="Upload a video or use your webcam. The model will predict if bullying is present."
)

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