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

# Load the YOLO model
model = YOLO('yolov5s.pt')  # You can use a different model if needed

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.pred[0]  # Get predictions

        # Count people detected (class ID for person is usually 0)
        for det in detections:
            if det[5] == 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..."):
            count = count_people("temp_video.mp4")
        st.success(f"Total number of people detected: {count}")