import gradio as gr
import cv2
from PIL import Image
from yolov3 import YOLO  # Update this import based on your YOLO implementation
from sort import Sort

# Initialize YOLO and SORT
yolo = YOLO()
tracker = Sort()

def detect_people(frame):
    try:
        # Add YOLO detection logic here
        # Return the detected bounding boxes
        return []
    except Exception as e:
        print(f"Detection Error: {e}")
        return []

def track_people(detections):
    try:
        # Update SORT tracker with the detections
        return tracker.update(detections)
    except Exception as e:
        print(f"Tracking Error: {e}")
        return []

def count_people(video_file):
    cap = cv2.VideoCapture(video_file)
    unique_people = {}

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        detections = detect_people(frame)
        tracked_objects = track_people(detections)

        for obj in tracked_objects:
            obj_id = int(obj[4])
            if obj_id not in unique_people:
                x1, y1, x2, y2 = map(int, obj[:4])
                person_img = frame[y1:y2, x1:x2]
                unique_people[obj_id] = person_img

    cap.release()

    output_images = [Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) for img in unique_people.values()]
    return len(unique_people), output_images

iface = gr.Interface(
    fn=count_people,
    inputs=gr.inputs.Video(),
    outputs=[gr.outputs.Textbox(label="Number of Unique People"), gr.outputs.Gallery(label="Unique People Images")],
    title="Unique People Counter",
    description="Upload a video to count and view unique people."
)

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