import cv2
import tempfile
import gradio as gr
from ultralytics import YOLO
import pandas as pd

def process_video(video_file):
    # Define colors for each class (8 classes)
    colors = [
        (255, 0, 0),   # Class 0 - Blue
        (50, 205, 50),   # Class 1 - Green
        (0, 0, 255),   # Class 2 - Red
        (255, 255, 0), # Class 3 - Cyan
        (255, 0, 255), # Class 4 - Magenta
        (255, 140, 0), # Class 5 - Orange
        (128, 0, 128), # Class 6 - Purple
        (0, 128, 128)  # Class 7 - Teal
    ]

    # Define class names (example names, replace with actual class names if available)
    class_names = ['Hymenoptera', 'Mantodea', 'Odonata', 'Orthoptera', 'Coleoptera', 'Lepidoptera', 'Hemiptera']

    # Load the YOLOv8 model
    model = YOLO("insect_detection4.pt")

    # Open the video file
    cap = cv2.VideoCapture(video_file)

    # Prepare DataFrame for storing detection data
    columns = ["frame", "id", "class", "x", "y", "w", "h"]
    df = pd.DataFrame(columns=columns)

    frame_id = 0

    # Loop through the video frames
    while cap.isOpened():
        # Read a frame from the video
        success, frame = cap.read()

        if success:
            frame_id += 1

            # Run YOLOv8 tracking on the frame, persisting tracks between frames
            results = model.track(frame, persist=True, tracker="insect_tracker.yaml")

            for result in results:
                boxes = result.boxes.cpu().numpy()
                confidences = boxes.conf
                class_ids = boxes.cls

                for i, box in enumerate(boxes):
                    class_id = int(class_ids[i])
                    confidence = confidences[i]

                    # Append detection data to DataFrame
                    new_row = pd.DataFrame({
                        "frame": [frame_id],
                        "id": [box.id],
                        "class": [int(box.cls[0])],
                        "x": [box.xywh[0][0]],
                        "y": [box.xywh[0][1]],
                        "w": [box.xywh[0][2]],
                        "h": [box.xywh[0][3]]
                    })
                    df = pd.concat([df, new_row], ignore_index=True)

        else:
            break

    # Release the video capture
    cap.release()

    # Save DataFrame to CSV
    csv_path = tempfile.mktemp(suffix=".csv")
    df.to_csv(csv_path, index=False)

    return df, csv_path

# Create a Gradio interface
inputs = gr.Video(label="Input Video")
outputs = [
    gr.DataFrame(label="Detection Data"),
    gr.File(label="Download CSV")
]

gr.Interface(fn=process_video, inputs=inputs, outputs=outputs).launch()