Create app.py

app.py

@@ -0,0 +1,154 @@
+import streamlit as st
+import cv2
+import tempfile
+import torch
+import numpy as np
+from ultralytics import YOLO
+from deep_sort_realtime.deepsort_tracker import DeepSort
+import warnings
+import os
+import platform
+
+# Suppress ScriptRunContext warnings from threads
+warnings.filterwarnings("ignore", message=".*missing ScriptRunContext.*")
+
+# Check if running in headless environment
+IS_HEADLESS = platform.system() == 'Linux'
+
+# Initialize YOLO model
+@st.cache_resource
+def load_yolo_model():
+    try:
+        return YOLO("best.pt")
+    except Exception as e:
+        st.error(f"Error loading YOLO model: {str(e)}")
+        return None
+
+# Main app
+st.title("📦 Inventory Management")
+
+# Settings
+st.sidebar.header("Settings")
+CONF_THRESHOLD = st.sidebar.slider(
+    "Confidence Threshold",
+    min_value=0.0,
+    max_value=1.0,
+    value=0.4,
+    help="Higher values mean more confident detections but might miss objects"
+)
+
+FRAME_SKIP = st.sidebar.slider(
+    "Frame Skip",
+    min_value=0,
+    max_value=10,
+    value=2,
+    help="Process every Nth frame (higher values = faster processing but may miss objects)"
+)
+
+# Load YOLO model
+yolo_model = load_yolo_model()
+
+if yolo_model is None:
+    st.error("Failed to load YOLO model. Please check if the model file exists.")
+    st.stop()
+
+# File uploader
+uploaded_file = st.sidebar.file_uploader(
+    "Upload Video", 
+    type=["mp4", "avi", "mov"],
+    help="Supported formats: MP4, AVI, MOV"
+)
+
+if uploaded_file is not None:
+    try:
+        tfile = tempfile.NamedTemporaryFile(delete=False)
+        tfile.write(uploaded_file.read())
+        video_path = tfile.name
+
+        if st.sidebar.button("Start Processing"):
+            tracker = DeepSort(
+                embedder="mobilenet",
+                embedder_gpu=torch.cuda.is_available(),
+                max_age=30  # Increase max_age for longer tracking retention
+            )
+
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                st.error("Error opening video file")
+                st.stop()
+
+            # Get video properties for processing
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            
+            counted_objects = set()
+            frame_placeholder = st.empty()
+            status_text = st.sidebar.empty()
+            progress_bar = st.progress(0)
+            
+            # Counter for frame skipping
+            frame_counter = 0
+            
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                frame_counter += 1
+                current_position = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                progress = current_position / frame_count
+                progress_bar.progress(progress)
+                
+                # Skip frames based on user setting
+                if FRAME_SKIP > 0 and frame_counter % (FRAME_SKIP + 1) != 0:
+                    continue
+
+                try:
+                    # Resize frame for faster processing (if needed)
+                    # h, w = frame.shape[:2]
+                    # if w > 1280:  # Only resize if the frame is large
+                    #     frame = cv2.resize(frame, (1280, int(h * 1280 / w)))
+                    
+                    results = yolo_model(frame, verbose=False)  # Turn off verbose output for speed
+                    detections = []
+                    for result in results:
+                        for box in result.boxes.data.tolist():
+                            x1, y1, x2, y2, score, class_id = box
+                            if score > CONF_THRESHOLD:
+                                detections.append([[x1, y1, x2 - x1, y2 - y1], score, int(class_id)])
+
+                    tracks = tracker.update_tracks(detections, frame=frame)
+                    for track in tracks:
+                        if not track.is_confirmed():
+                            continue
+                        track_id = track.track_id
+                        ltrb = track.to_ltrb()
+                        x1, y1, x2, y2 = map(int, ltrb)
+                        counted_objects.add(track_id)
+                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                        cv2.putText(frame, f"ID: {track_id}", (x1, y1-10), 
+                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+                    cv2.putText(frame, f"Total Objects: {len(counted_objects)}", 
+                                (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    frame_placeholder.image(frame_rgb, channels="RGB", use_column_width=True)
+                    status_text.info(f"Processing... Current count: {len(counted_objects)}")
+                    
+                    # Remove sleep to maximize performance
+                    # time.sleep(0.01)
+
+                except Exception as e:
+                    st.error(f"Error processing frame: {str(e)}")
+                    continue
+
+            cap.release()
+            progress_bar.progress(1.0)
+            st.sidebar.success(f"Final count: {len(counted_objects)} objects")
+            st.balloons()
+
+    except Exception as e:
+        st.error(f"Error processing video: {str(e)}")
+    finally:
+        if 'tfile' in locals():
+            tfile.close()