Ramakrishna1999/Object_Tracking_App

1import streamlit as st
2import cv2
3import os
4from ultralytics import YOLO
5
6# Streamlit config
7st.set_page_config(page_title="YOLOv8s Video Tracking", layout="wide")
8
9st.title("🎥 YOLOv8s Object Tracking App")
10st.write("Upload a video and track objects with unique IDs in real-time using YOLOv8s (best speed + accuracy balance).")
11
12# Sidebar tracker option
13tracker_choice = st.sidebar.selectbox("Choose Tracker", ["bytetrack.yaml"])
14
15# Load YOLOv8s model (fixed)
16@st.cache_resource
17def load_model():
18    return YOLO("yolov8s.pt")
19
20model = load_model()
21
22# Upload video
23uploaded_file = st.file_uploader("📂 Upload a Video", type=["mp4", "avi", "mov", "mkv", "wmv"])
24
25if uploaded_file is not None:
26    os.makedirs("uploads", exist_ok=True)
27    os.makedirs("outputs", exist_ok=True)
28
29    video_path = os.path.join("uploads", uploaded_file.name)
30    output_path = os.path.join("outputs", f"tracked_{uploaded_file.name}")
31
32    with open(video_path, "wb") as f:
33        f.write(uploaded_file.getbuffer())
34
35    # Two video panels (side-by-side)
36    col1, col2 = st.columns(2)
37    col1.subheader("📹 Original Video")
38    col2.subheader("🎯 Tracked Video")
39
40    st.write("🔄 Tracking in progress...")
41
42    # Placeholders for frames
43    orig_placeholder = col1.empty()
44    track_placeholder = col2.empty()
45
46    # Video writer for saving output
47    cap = cv2.VideoCapture(video_path)
48    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
49    out = cv2.VideoWriter(output_path, fourcc, cap.get(cv2.CAP_PROP_FPS),
50                          (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
51                           int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
52
53    # YOLO tracking (conf=0.6, stride=2 for speed)
54    results = model.track(
55        source=video_path,
56        conf=0.6,
57        tracker=tracker_choice,
58        stream=True,
59        persist=True,
60        vid_stride=2
61    )
62
63    for idx, r in enumerate(results):
64        frame = r.orig_img
65        tracked_frame = r.plot()
66
67        # Save tracked frame
68        out.write(tracked_frame)
69
70        # Convert BGR → RGB
71        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
72        tracked_rgb = cv2.cvtColor(tracked_frame, cv2.COLOR_BGR2RGB)
73
74        # Update UI every 3rd frame for speed
75        if idx % 3 == 0:
76            orig_placeholder.image(frame_rgb, channels="RGB", use_container_width=True)
77            track_placeholder.image(tracked_rgb, channels="RGB", use_container_width=True)
78
79    cap.release()
80    out.release()
81
82    st.success("✅ Tracking Complete!")
83    with open(output_path, "rb") as f:
84        st.download_button("⬇️ Download Tracked Video", f, file_name=f"tracked_{uploaded_file.name}")
85