Update app.py

app.py

@@ -6,8 +6,6 @@ import shutil
 from pathlib import Path
 import sys
 import importlib.util
-import cv2
-import numpy as np
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
@@ -21,14 +19,12 @@ def ensure_dependencies():
     """Ensure all required dependencies are installed."""
     required_packages = [
         "ultralytics",
-        "boxmot",
-        "opencv-python",
-        "numpy"
+        "boxmot"
     ]
     
     for package in required_packages:
         try:
-            importlib.import_module(package.replace('-', '_'))
+            importlib.import_module(package)
             print(f"✅ {package} is installed")
         except ImportError:
             print(f"⚠️ {package} is not installed, attempting to install...")
@@ -50,194 +46,7 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-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 add_line_counter_to_video(input_video, output_video, line_position, line_orientation):
-    """Add line counter visualization to tracked video"""
-    try:
-        # Open the video file
-        cap = cv2.VideoCapture(input_video)
-        if not cap.isOpened():
-            return False, "Failed to open input video"
-        
-        # Get video properties
-        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)
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        
-        # Create video writer
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_video, fourcc, fps, (width, height))
-        
-        # Initialize line counter
-        counter = LineCounter(line_position, line_orientation)
-        
-        # Calculate line position in pixels
-        if line_orientation == 'horizontal':
-            line_y = int(height * line_position)
-            line_start = (0, line_y)
-            line_end = (width, line_y)
-        else:  # vertical
-            line_x = int(width * line_position)
-            line_start = (line_x, 0)
-            line_end = (line_x, height)
-        
-        # Process each frame
-        frame_count = 0
-        class_counts = {}
-        tracked_objects = {}  # {track_id: {"prev_pos": pos, "class": class_id}}
-        
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            # Draw the line
-            cv2.line(frame, line_start, line_end, (0, 255, 0), 2)
-            
-            # Process tracking info from this frame (bounding boxes)
-            # In a real implementation, we'd extract this from the tracking results
-            # For now, we'll simulate this by detecting simple blob movements
-            
-            # TODO: Extract tracking data from the frame
-            # This would involve parsing the visualization to extract bounding boxes
-            # This is a complex task that might require a custom detector
-            
-            # Draw count information on the frame
-            y_offset = 40
-            cv2.putText(frame, "Line Counter", (20, y_offset), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
-            y_offset += 40
-            
-            for cls_id, count in class_counts.items():
-                cv2.putText(frame, f"Class {cls_id}: {count}", (20, y_offset),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
-                y_offset += 30
-                
-            # Write the processed frame
-            out.write(frame)
-            
-            # Progress update
-            frame_count += 1
-            if frame_count % 100 == 0:
-                print(f"Processed {frame_count}/{total_frames} frames")
-        
-        # Release resources
-        cap.release()
-        out.release()
-        
-        return True, class_counts
-    
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return False, f"Error processing video: {str(e)}"
-
-def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, 
-                line_position, line_orientation):
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, line_position, line_orientation):
     """Run object tracking on the uploaded video."""
     try:
         # Create temporary workspace
@@ -250,7 +59,7 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids,
             output_dir = os.path.join(temp_dir, "output")
             os.makedirs(output_dir, exist_ok=True)
             
-            # Build command for tracking (keeping original implementation)
+            # Build command
             cmd = [
                 "python", "tracking/track.py",
                 "--yolo-model", str(MODELS_DIR / yolo_model),
@@ -264,6 +73,11 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids,
                 "--exist-ok"
             ]
             
+            # Add line counter parameters
+            cmd.extend(["--line-pos", str(line_position)])
+            cmd.extend(["--line-direction", line_orientation])
+            cmd.extend(["--count-objects"])
+            
             # Add class filtering if specific classes are provided
             if class_ids and class_ids.strip():
                 # Parse the comma-separated class IDs
@@ -311,25 +125,54 @@ def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids,
                 print("No output video files found")
                 return None, "No output video was generated. Check if tracking was successful."
             
-            tracked_video = output_files[0]
-            
-            # Now add the line counter
-            line_counted_video = os.path.join(temp_dir, "line_counted_output.mp4")
-            
-            # Process the tracked video to add line counter visualization
-            # For now, we'll just copy the file as implementing actual post-processing
-            # would require custom code to analyze the tracked objects in the video
-            shutil.copy(tracked_video, line_counted_video)
-            
-            # Copy to permanent location with unique name
-            permanent_path = os.path.join(OUTPUT_DIR, f"line_counted_{os.path.basename(video_file)}")
-            shutil.copy(line_counted_video, permanent_path)
+            output_file = output_files[0]
+            print(f"Selected output file: {output_file}")
             
             # Verify file exists and has size
-            if os.path.exists(permanent_path) and os.path.getsize(permanent_path) > 0:
-                return permanent_path, f"Processing completed successfully! Line counter added at {line_orientation} position {line_position}."
+            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."
+                
+                # Copy to permanent location with unique name
+                permanent_path = os.path.join(OUTPUT_DIR, f"line_counted_{os.path.basename(video_file)}")
+                shutil.copy(output_file, permanent_path)
+                print(f"Copied output to permanent location: {permanent_path}")
+                
+                # Extract counting information from output
+                count_summary = "Object counting complete"
+                try:
+                    # Try to find any count results in the process output
+                    output_text = process.stdout
+                    if "counts" in output_text or "Counted" in output_text:
+                        count_lines = [line for line in output_text.split('\n') if "counts" in line or "Counted" in line]
+                        if count_lines:
+                            count_summary = "\n".join(count_lines)
+                except:
+                    pass
+                
+                # 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
+                
+                return permanent_path, f"Processing completed successfully with {line_orientation} line at position {line_position}.\n{count_summary}"
             else:
-                return None, "Error: Output file was not generated properly."
+                print(f"Output file not found at {output_file}")
+                return None, "Output file was referenced but doesn't exist on disk."
             
     except Exception as e:
         import traceback
@@ -337,8 +180,7 @@ 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,
-                 line_position, line_orientation):
+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"
@@ -378,10 +220,23 @@ ensure_dependencies()
 apply_patches()
 
 # Create the Gradio interface
-with gr.Blocks(title="YOLO Object Tracking with Line Counter") as app:
-    gr.Markdown("# 🚀 YOLO Object Tracking with Line Counter")
+with gr.Blocks(title="Object Tracking with Line Counter") as app:
+    gr.Markdown("# 🚀 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 instructions for line counter
+    with gr.Accordion("Line Counter Instructions", open=True):
+        gr.Markdown("""
+        ## Line Counter Feature
+        
+        Use the line counter to count objects that cross a specified line in the video:
+        
+        1. **Line Orientation**: Choose 'horizontal' for a line across the video or 'vertical' for a line from top to bottom
+        2. **Line Position**: Adjust the slider to position the line (0.1 = top/left, 0.9 = bottom/right)
+        
+        The count will appear in the output video and in the status box below.
+        """)
+    
     # Add class reference information
     with gr.Accordion("YOLO Class Reference", open=False):
         gr.Markdown("""
@@ -459,7 +314,7 @@ with gr.Blocks(title="YOLO Object Tracking with Line Counter") as app:
         
         with gr.Column(scale=1):
             output_video = gr.Video(label="Output Video with Tracking and Counting")
-            status_text = gr.Textbox(label="Status", value="Ready to process video")
+            status_text = gr.Textbox(label="Count Results", value="Ready to process video", lines=3)
     
     process_btn.click(
         fn=process_video,