datascientist22's picture
Update app.py
f167ce4 verified
import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from tempfile import NamedTemporaryFile
import os
# Initialize YOLOv8 model
model = YOLO("yolov8n.pt")
# Streamlit app title and creator information
st.markdown("Created by: [Engr. Hamesh Raj](https://www.linkedin.com/in/datascientisthameshraj/)")
st.title("🎥 YOLOv8 Object Detection on Videos")
# Sidebar for video upload
st.sidebar.header("Upload Video")
uploaded_video = st.sidebar.file_uploader("Choose a video...", type=["mp4", "mov", "avi", "mkv"])
if uploaded_video is not None:
# Save the uploaded video to a temporary file
temp_video = NamedTemporaryFile(delete=False)
temp_video.write(uploaded_video.read())
video_path = temp_video.name
# Display the uploaded video
st.sidebar.video(uploaded_video)
# Submit button to process the video
if st.sidebar.button("Submit"):
st.subheader("Processing Video...")
# Open the video file
cap = cv2.VideoCapture(video_path)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
# Create a temporary file to save the output video
temp_output_video = NamedTemporaryFile(delete=False, suffix='.mp4')
output_video_path = temp_output_video.name
# Define codec and create VideoWriter object
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
# Process each frame of the video
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Perform object detection
results = model(frame)
# Draw bounding boxes on the frame
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = box.conf[0]
cls = box.cls[0]
label = f'{model.names[int(cls)]} {conf:.2f}'
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
out.write(frame)
cap.release()
out.release()
# Display the processed video
st.subheader("Processed Video")
st.video(output_video_path)
# Download button for the processed video
with open(output_video_path, "rb") as file:
st.download_button(
label="Download Processed Video",
data=file,
file_name="processed_video.mp4",
mime="video/mp4"
)
# Clean up temporary files
os.remove(video_path)
os.remove(output_video_path)