import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO

# Load the YOLO model
model = YOLO('yolov5s.pt')  # Use 'yolov5s.pt' or any YOLO model of your choice

def count_people(video_file):
    count = 0
    cap = cv2.VideoCapture(video_file)

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

        results = model(frame)
        detections = results[0]  # Access the first result
        
        # Count people detected (class ID for person is usually 0)
        for det in detections.boxes.data:  # Access the boxes
            class_id = int(det[5])  # Class ID is the 6th element
            if class_id == 0:  # Check if class ID is 0 (person)
                count += 1

    cap.release()
    return count

# Streamlit app layout
st.title("Person Detection in Video")
st.write("Upload a video file to count the number of times a person appears.")

# File uploader for video files
video_file = st.file_uploader("Choose a video file", type=["mp4", "avi", "mov"])

if video_file is not None:
    # Save the uploaded video to a temporary location
    with open("temp_video.mp4", "wb") as f:
        f.write(video_file.getbuffer())

    st.video(video_file)  # Display the video

    if st.button("Count People"):
        with st.spinner("Counting..."):
            print("model loaded")
            count = count_people("temp_video.mp4")
        st.success(f"Total number of people detected: {count}")