Spaces:
Sleeping
Sleeping
File size: 1,481 Bytes
88e2ff7 90b7de5 88e2ff7 90b7de5 88e2ff7 90b7de5 88e2ff7 e1662bf 88e2ff7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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}")
|