import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from sklearn.preprocessing import StandardScaler
from ultralytics import YOLO

# Load models
lstm_model = load_model('suspicious_activity_model.h5')
yolo_model = YOLO('yolov8n-pose.pt')  # Ensure this model supports keypoint detection
scaler = StandardScaler()

# Function to extract keypoints from a frame
def extract_keypoints(frame):
    results = yolo_model(frame, verbose=False)
    for r in results:
        if r.keypoints is not None and len(r.keypoints) > 0:
            keypoints = r.keypoints.xyn.tolist()[0]  # Use the first person's keypoints
            flattened_keypoints = [kp for keypoint in keypoints for kp in keypoint[:2]]  # Flatten x, y values
            return flattened_keypoints
    return None  # Return None if no keypoints are detected

# Function to process each frame
def process_frame(frame):
    results = yolo_model(frame, verbose=False)

    for box in results[0].boxes:
        cls = int(box.cls[0])  # Class ID
        confidence = float(box.conf[0])

        if cls == 0 and confidence > 0.5:  # Detect persons only
            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates

            # Extract ROI for classification
            roi = frame[y1:y2, x1:x2]
            if roi.size > 0:
                keypoints = extract_keypoints(roi)
                if keypoints is not None and len(keypoints) > 0:
                    keypoints_scaled = scaler.fit_transform([keypoints])
                    keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))

                    prediction = (lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]

                    color = (0, 0, 255) if prediction == 1 else (0, 255, 0)
                    label = 'Suspicious' if prediction == 1 else 'Normal'
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    return frame

# Gradio video streaming function
def video_processing(video_frame):
    frame = cv2.cvtColor(video_frame, cv2.COLOR_BGR2RGB)  # Convert to RGB
    processed_frame = process_frame(frame)
    return processed_frame

# Launch Gradio app
gr.Interface(
    fn=video_processing,
    inputs=gr.Video(streaming=True),  # Correct the Video component
    outputs="video",
    live=True,
    title="Suspicious Activity Detection"
).launch(debug=True)