|
import streamlit as st |
|
import cv2 |
|
import numpy as np |
|
from ultralytics import YOLO |
|
|
|
|
|
model = YOLO('yolov5s.pt') |
|
|
|
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] |
|
|
|
|
|
for det in detections.boxes.data: |
|
class_id = int(det[5]) |
|
if class_id == 0: |
|
count += 1 |
|
|
|
cap.release() |
|
return count |
|
|
|
|
|
st.title("Person Detection in Video") |
|
st.write("Upload a video file to count the number of times a person appears.") |
|
|
|
|
|
video_file = st.file_uploader("Choose a video file", type=["mp4", "avi", "mov"]) |
|
|
|
if video_file is not None: |
|
|
|
with open("temp_video.mp4", "wb") as f: |
|
f.write(video_file.getbuffer()) |
|
|
|
st.video(video_file) |
|
|
|
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}") |
|
|