Spaces:
Runtime error
Runtime error
File size: 2,529 Bytes
a612e50 | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import streamlit as st
import cv2
import tempfile
from ultralytics import YOLO
import imageio
import os
# ---------------- YOLO Model ----------------
st.title("🎥 YOLOv8 Object Tracking on Video")
model = YOLO("yolov8n.pt") # Pretrained YOLOv8n model
# ---------------- Function to Process Video ----------------
def detect_objects_in_video(video_path):
cap = cv2.VideoCapture(video_path)
fps = int(cap.get(cv2.CAP_PROP_FPS)) or 25
# Temporary output file
temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
output_path = temp_output.name
temp_output.close()
writer = imageio.get_writer(output_path, fps=fps, codec="libx264")
stframe = st.empty()
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
processed_frames = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame, imgsz=1280, conf=0.25)
annotated_frame = results[0].plot()
writer.append_data(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
# Progress bar
processed_frames += 1
stframe.text(f"Processing frame {processed_frames}/{total_frames}...")
cap.release()
writer.close()
# Read video as bytes for Streamlit
with open(output_path, "rb") as f:
video_bytes = f.read()
try:
os.remove(output_path)
except PermissionError:
pass
return video_bytes
# ---------------- Streamlit UI ----------------
st.write("Upload a video and see detections in MP4 format.")
uploaded_file = st.file_uploader("Upload a video", type=["mp4", "mov", "avi"])
if uploaded_file is not None:
# Save uploaded video to temp file
temp_input = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
temp_input.write(uploaded_file.read())
temp_input.close()
st.subheader("📥 Original Video")
st.video(temp_input.name) # Display original video
st.write("🔄 Processing video with YOLOv8... please wait")
video_bytes = detect_objects_in_video(temp_input.name)
st.subheader("✅ Processed Video")
st.video(video_bytes) # Display processed video
st.download_button(
label="⬇️ Download Processed Video",
data=video_bytes,
file_name="processed_output.mp4",
mime="video/mp4"
)
# Cleanup
try:
os.remove(temp_input.name)
except PermissionError:
pass
|