CheckVideo / app.py
ish717's picture
Added file
88e2ff7 verified
raw
history blame
1.33 kB
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}")