Update app.py

app.py

@@ -1,74 +1,87 @@
-# streamlit_yolo_video.py
-
-import streamlit as st
-import cv2
-import numpy as np
-from ultralytics import YOLO
-import tempfile
-
-# -------------------------------
-# Streamlit App Title
-# -------------------------------
-st.title("YOLOv8 Object Detection on Video")
-
-# -------------------------------
-# Load YOLO Detection Model
-# -------------------------------
-@st.cache_resource
-def load_model(model_name="yolov8n.pt"):
-    return YOLO(model_name)
-
-model = load_model()
-
-# -------------------------------
-# Upload Video
-# -------------------------------
-uploaded_file = st.file_uploader("Upload a video", type=["mp4", "mov", "avi", "mkv"])
-
-if uploaded_file is not None:
-    # Save uploaded video temporarily
-    tfile = tempfile.NamedTemporaryFile(delete=False)
-    tfile.write(uploaded_file.read())
-    video_path = tfile.name
-
-    # Open video
-    cap = cv2.VideoCapture(video_path)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-
-    # Output video writer
-    output_path = "output_video.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-
-    stframe = st.empty()
-    st.info("Processing video... This may take some time depending on length.")
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # YOLO expects RGB
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = model(frame_rgb)
-
-        # Annotated frame
-        annotated_frame = results[0].plot()
-        annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
-
-        # Write frame to output video
-        out.write(annotated_frame_bgr)
-
-        # Display in Streamlit
-        stframe.image(annotated_frame, channels="RGB", use_column_width=True)
-
-    cap.release()
-    out.release()
-
-    st.success("Video processing completed!")
-
-    # Download button
-    with open(output_path, "rb") as f:
-        st.download_button("Download Annotated Video", f, "annotated_video.mp4")
+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