Update app.py

app.py

@@ -6,6 +6,9 @@ import shutil
 from pathlib import Path
 import sys
 import importlib.util
+import cv2
+import numpy as np
+from ultralytics.utils.plotting import Annotator, colors
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
@@ -19,12 +22,14 @@ def ensure_dependencies():
     """Ensure all required dependencies are installed."""
     required_packages = [
         "ultralytics",
-        "boxmot"
+        "boxmot",
+        "opencv-python",
+        "numpy"
     ]
     
     for package in required_packages:
         try:
-            importlib.import_module(package)
+            importlib.import_module(package.replace('-', '_'))
             print(f"✅ {package} is installed")
         except ImportError:
             print(f"⚠️ {package} is not installed, attempting to install...")
@@ -46,8 +51,219 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
-    """Run object tracking on the uploaded video."""
+class LineCounter:
+    """Count objects crossing a line"""
+    def __init__(self, line_position=0.5, line_orientation='horizontal'):
+        """
+        Initialize a line counter
+        
+        Args:
+            line_position: float between 0 and 1, position of line (default: middle)
+            line_orientation: 'horizontal' or 'vertical'
+        """
+        self.line_position = line_position
+        self.line_orientation = line_orientation
+        self.counts = {}  # Track counts by class
+        self.crossed_ids = set()  # Track IDs that have crossed the line
+        self.prev_positions = {}  # Store previous positions of tracked objects
+        
+    def update(self, bboxes, identities, clss, frame_shape):
+        """
+        Update counter with new detections
+        
+        Args:
+            bboxes: list of bounding boxes [x1, y1, x2, y2]
+            identities: list of track IDs
+            clss: list of class IDs
+            frame_shape: tuple of (height, width)
+        
+        Returns:
+            count_info: dict of counts by class
+        """
+        if not len(bboxes):
+            return self.counts, []
+            
+        height, width = frame_shape[:2]
+        
+        # Calculate line position in pixels
+        if self.line_orientation == 'horizontal':
+            line_pos = int(height * self.line_position)
+            start_point = (0, line_pos)
+            end_point = (width, line_pos)
+        else:  # vertical
+            line_pos = int(width * self.line_position)
+            start_point = (line_pos, 0)
+            end_point = (line_pos, height)
+        
+        # Store line info for drawing
+        line_info = {
+            'start': start_point,
+            'end': end_point,
+            'orientation': self.line_orientation
+        }
+        
+        # Process each detection
+        for i, (bbox, track_id, cls) in enumerate(zip(bboxes, identities, clss)):
+            x1, y1, x2, y2 = bbox
+            
+            # Use center point of bbox to determine position
+            if self.line_orientation == 'horizontal':
+                center_pos = (y1 + y2) / 2
+                crossed_now = False
+                
+                # Check if crossed the line
+                if track_id in self.prev_positions:
+                    prev_pos = self.prev_positions[track_id]
+                    # Crossed from above to below
+                    if prev_pos < line_pos and center_pos >= line_pos:
+                        crossed_now = True
+                    # Crossed from below to above
+                    elif prev_pos >= line_pos and center_pos < line_pos:
+                        crossed_now = True
+                
+                # Store current position for next frame
+                self.prev_positions[track_id] = center_pos
+                
+            else:  # vertical line
+                center_pos = (x1 + x2) / 2
+                crossed_now = False
+                
+                # Check if crossed the line
+                if track_id in self.prev_positions:
+                    prev_pos = self.prev_positions[track_id]
+                    # Crossed from left to right
+                    if prev_pos < line_pos and center_pos >= line_pos:
+                        crossed_now = True
+                    # Crossed from right to left
+                    elif prev_pos >= line_pos and center_pos < line_pos:
+                        crossed_now = True
+                
+                # Store current position for next frame
+                self.prev_positions[track_id] = center_pos
+            
+            # Count if crossed and not counted before
+            if crossed_now and track_id not in self.crossed_ids:
+                self.crossed_ids.add(track_id)
+                cls_id = int(cls)
+                self.counts[cls_id] = self.counts.get(cls_id, 0) + 1
+        
+        return self.counts, line_info
+    
+    def reset(self):
+        """Reset all counts and tracking info"""
+        self.counts = {}
+        self.crossed_ids = set()
+        self.prev_positions = {}
+
+def process_video_with_counter(input_path, output_path, model_path, reid_model, tracking_method, 
+                              selected_classes, conf_threshold, line_position, line_orientation):
+    """Process video with the line counter"""
+    # Import here to avoid import errors if dependencies are missing
+    from ultralytics import YOLO
+    
+    # Load the model
+    model = YOLO(model_path)
+    
+    # Prepare classes filter
+    classes = None
+    if selected_classes and selected_classes.strip():
+        try:
+            classes = [int(c.strip()) for c in selected_classes.split(",") if c.strip()]
+        except ValueError:
+            print("Invalid class IDs, using all classes")
+    
+    # Initialize video capture and get video info
+    cap = cv2.VideoCapture(input_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)
+    
+    # Initialize video writer
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    
+    # Initialize line counter
+    counter = LineCounter(line_position=line_position, line_orientation=line_orientation)
+    
+    # Track with YOLO
+    results = model.track(
+        source=input_path,
+        conf=conf_threshold,
+        classes=classes,
+        tracker=tracking_method,
+        save=False,
+        stream=True,
+        verbose=False
+    )
+    
+    # Process each frame
+    for i, result in enumerate(results):
+        frame = result.orig_img
+        
+        # Skip if no detections or tracking info
+        if result.boxes.id is None:
+            annotator = Annotator(frame)
+            # Draw the line
+            if line_orientation == 'horizontal':
+                line_y = int(height * line_position)
+                annotator.line((0, line_y), (width, line_y), color=(0, 255, 0), thickness=2)
+            else:
+                line_x = int(width * line_position)
+                annotator.line((line_x, 0), (line_x, height), color=(0, 255, 0), thickness=2)
+                
+            # Add count text
+            count_text = "Count: 0"
+            annotator.text((20, 40), count_text, color=(0, 0, 255), thickness=2)
+            
+            out.write(frame)
+            continue
+        
+        # Get boxes, track IDs and classes
+        boxes = result.boxes.xyxy.cpu().numpy()
+        track_ids = result.boxes.id.cpu().numpy().astype(int)
+        classes = result.boxes.cls.cpu().numpy().astype(int)
+
+        # Update counter
+        counts, line_info = counter.update(boxes, track_ids, classes, frame.shape)
+        
+        # Start drawing on frame
+        annotator = Annotator(frame)
+        
+        # Draw tracking results
+        for box, track_id, cls in zip(boxes, track_ids, classes):
+            x1, y1, x2, y2 = map(int, box)
+            id = int(track_id)
+            color = colors(id % 10)
+            # Draw bbox
+            annotator.box_label([x1, y1, x2, y2], f'{id} {model.names[int(cls)]}', color=color)
+        
+        # Draw the counting line
+        line_start, line_end = line_info['start'], line_info['end']
+        annotator.line(line_start, line_end, color=(0, 255, 0), thickness=2)
+        
+        # Draw counts for each class
+        y_offset = 40
+        for cls_id, count in counts.items():
+            cls_name = model.names.get(cls_id, f"Class {cls_id}")
+            count_text = f"{cls_name}: {count}"
+            annotator.text((20, y_offset), count_text, color=(0, 0, 255), thickness=2)
+            y_offset += 30
+        
+        # Write the frame
+        out.write(frame)
+        
+        # Print progress
+        if i % 100 == 0:
+            print(f"Processed {i} frames")
+    
+    # Release resources
+    cap.release()
+    out.release()
+    return counts
+
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, 
+                line_position, line_orientation):
+    """Run object tracking with line counting on the uploaded video."""
     try:
         # Create temporary workspace
         with tempfile.TemporaryDirectory() as temp_dir:
@@ -55,107 +271,45 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids,
             input_path = os.path.join(temp_dir, "input_video.mp4")
             shutil.copy(video_file, input_path)
             
-            # Prepare output directory
-            output_dir = os.path.join(temp_dir, "output")
-            os.makedirs(output_dir, exist_ok=True)
-            
-            # Build command
-            cmd = [
-                "python", "tracking/track.py",
-                "--yolo-model", str(MODELS_DIR / yolo_model),
-                "--reid-model", str(MODELS_DIR / reid_model),
-                "--tracking-method", tracking_method,
-                "--source", input_path,
-                "--conf", str(conf_threshold),
-                "--save",
-                "--project", output_dir,
-                "--name", "track",
-                "--exist-ok"
-            ]
+            # Prepare output file
+            output_path = os.path.join(temp_dir, "output_video.mp4")
             
-            # Add class filtering if specific classes are provided
-            if class_ids and class_ids.strip():
-                # Parse the comma-separated class IDs
-                try:
-                    # Split by comma and convert to integers to validate
-                    class_list = [int(c.strip()) for c in class_ids.split(",") if c.strip()]
-                    # Add each class ID as a separate argument
-                    if class_list:
-                        cmd.append("--classes")
-                        cmd.extend(str(c) for c in class_list)
-                except ValueError:
-                    return None, "Invalid class IDs. Please enter comma-separated numbers (e.g., '0,1,2')."
+            # Get full model path
+            model_path = str(MODELS_DIR / yolo_model)
             
-            # Special handling for OcSort
-            if tracking_method == "ocsort":
-                cmd.append("--per-class")
+            print(f"Processing video with counter. Model: {model_path}, Line: {line_orientation} at {line_position}")
             
-            # Execute tracking with error handling
-            print(f"Executing command: {' '.join(cmd)}")
-            process = subprocess.run(
-                cmd,
-                capture_output=True,
-                text=True
+            # Process the video with our counter function
+            counts = process_video_with_counter(
+                input_path=input_path,
+                output_path=output_path,
+                model_path=model_path,
+                reid_model=reid_model,
+                tracking_method=tracking_method,
+                selected_classes=class_ids,
+                conf_threshold=conf_threshold,
+                line_position=line_position,
+                line_orientation=line_orientation
             )
             
-            # Check for errors in output
-            if process.returncode != 0:
-                error_message = process.stderr or process.stdout
-                print(f"Process failed with return code {process.returncode}")
-                print(f"Error: {error_message}")
-                return None, f"Error in tracking process: {error_message}"
-            
-            print(f"Process completed with return code {process.returncode}")
-            
-            # Find output video
-            output_files = []
-            for root, _, files in os.walk(output_dir):
-                for file in files:
-                    if file.lower().endswith((".mp4", ".avi", ".mov")):
-                        output_files.append(os.path.join(root, file))
-            
-            print(f"Found output files: {output_files}")
-            
-            if not output_files:
-                print("No output video files found")
-                return None, "No output video was generated. Check if tracking was successful."
-            
-            output_file = output_files[0]
-            print(f"Selected output file: {output_file}")
-            
-            # Verify file exists and has size
-            if os.path.exists(output_file):
-                file_size = os.path.getsize(output_file)
-                print(f"Output file exists with size: {file_size} bytes")
-                
-                if file_size == 0:
-                    return None, "Output video was generated but has zero size."
-                
+            # Check if output file exists and has size
+            if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                 # Copy to permanent location with unique name
-                permanent_path = os.path.join(OUTPUT_DIR, f"output_{os.path.basename(video_file)}")
-                shutil.copy(output_file, permanent_path)
+                permanent_path = os.path.join(OUTPUT_DIR, f"counted_{os.path.basename(video_file)}")
+                shutil.copy(output_path, permanent_path)
                 print(f"Copied output to permanent location: {permanent_path}")
                 
-                # Ensure the file is in MP4 format for better compatibility with Gradio
-                if not permanent_path.lower().endswith('.mp4'):
-                    mp4_path = os.path.splitext(permanent_path)[0] + '.mp4'
-                    try:
-                        print(f"Converting to MP4 format: {mp4_path}")
-                        subprocess.run([
-                            'ffmpeg', '-i', permanent_path, 
-                            '-c:v', 'libx264', '-preset', 'fast',
-                            '-c:a', 'aac', mp4_path
-                        ], check=True, capture_output=True)
-                        os.remove(permanent_path)  # Remove the original file
-                        permanent_path = mp4_path
-                    except Exception as e:
-                        print(f"Failed to convert to MP4: {str(e)}")
-                        # Continue with original file if conversion fails
+                # Format counts for display
+                count_message = "Objects counted:\n"
+                if counts:
+                    for cls_id, count in counts.items():
+                        count_message += f"Class {cls_id}: {count}\n"
+                else:
+                    count_message += "No objects crossed the line"
                 
-                return permanent_path, "Processing completed successfully!"
+                return permanent_path, count_message
             else:
-                print(f"Output file not found at {output_file}")
-                return None, "Output file was referenced but doesn't exist on disk."
+                return None, "Error: Output video was not generated properly."
             
     except Exception as e:
         import traceback
@@ -163,13 +317,15 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids,
         return None, f"Error: {str(e)}"
 
 # Define the Gradio interface
-def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
+def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids, conf_threshold,
+                 line_position, line_orientation):
     # Validate inputs
     if not video_path:
         return None, "Please upload a video file"
     
     print(f"Processing video: {video_path}")
     print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, classes={class_ids}, conf={conf_threshold}")
+    print(f"Line counter: {line_orientation} at position {line_position}")
     
     output_path, status = run_tracking(
         video_path,
@@ -177,7 +333,9 @@ def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids
         reid_model,
         tracking_method,
         class_ids,
-        conf_threshold
+        conf_threshold,
+        line_position,
+        line_orientation
     )
     
     if output_path:
@@ -193,15 +351,16 @@ def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids
 yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
 reid_models = ["osnet_x0_25_msmt17.pt"]
 tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
+line_orientations = ["horizontal", "vertical"]
 
 # Ensure dependencies and apply patches at startup
 ensure_dependencies()
 apply_patches()
 
 # Create the Gradio interface
-with gr.Blocks(title="YOLO Object Tracking") as app:
-    gr.Markdown("# 🚀 YOLO Object Tracking")
-    gr.Markdown("Upload a video file to detect and track objects. Processing may take a few minutes depending on video length.")
+with gr.Blocks(title="YOLO Object Tracking with Line Counter") as app:
+    gr.Markdown("# 🚀 YOLO Object Tracking with Line Counter")
+    gr.Markdown("Upload a video file to detect, track and count objects crossing a line. Processing may take a few minutes depending on video length.")
     
     # Add class reference information
     with gr.Accordion("YOLO Class Reference", open=False):
@@ -261,15 +420,31 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                     label="Confidence Threshold"
                 )
                 
+                # Line counter settings
+                gr.Markdown("### Line Counter Settings")
+                line_orientation = gr.Dropdown(
+                    choices=line_orientations,
+                    value="horizontal",
+                    label="Line Orientation"
+                )
+                line_position = gr.Slider(
+                    minimum=0.1,
+                    maximum=0.9,
+                    value=0.5,
+                    step=0.05,
+                    label="Line Position (0.1 = top/left, 0.9 = bottom/right)"
+                )
+                
             process_btn = gr.Button("Process Video", variant="primary")
         
         with gr.Column(scale=1):
-            output_video = gr.Video(label="Output Video with Tracking")
-            status_text = gr.Textbox(label="Status", value="Ready to process video")
+            output_video = gr.Video(label="Output Video with Tracking and Counting")
+            status_text = gr.Textbox(label="Count Results", value="Ready to process video")
     
     process_btn.click(
         fn=process_video,
-        inputs=[input_video, yolo_model, reid_model, tracking_method, class_ids, conf_threshold],
+        inputs=[input_video, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, 
+                line_position, line_orientation],
         outputs=[output_video, status_text]
     )