main.py

"""
Vehicle Detection, Tracking, Counting, and Speed Estimation System
===================================================================

A comprehensive computer vision pipeline for analyzing traffic videos,
detecting vehicles, tracking their movement, counting them, and estimating
their speeds using YOLO object detection and perspective transformation.

Authors:
    - Abhay Gupta (0205CC221005)
    - Aditi Lakhera (0205CC221011)
    - Balraj Patel (0205CC221049)
    - Bhumika Patel (0205CC221050)

Technical Approach:
    - YOLO for real-time object detection
    - ByteTrack for multi-object tracking
    - Perspective transformation for speed calculation
    - Line zones for vehicle counting
"""

import sys
import logging
from pathlib import Path
from typing import Dict, Optional, Callable
from time import time

import cv2
import numpy as np
import supervision as sv
from ultralytics import YOLO

from src import FrameAnnotator, VehicleSpeedEstimator, PerspectiveTransformer
from src.exceptions import (
    VideoProcessingError,
    ModelLoadError,
    ConfigurationError
)
from config import VehicleDetectionConfig

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class VehicleDetectionPipeline:
    """
    Main pipeline for vehicle detection, tracking, counting, and speed estimation.
    
    This class orchestrates the entire processing workflow, from loading the model
    to processing each frame and generating the output video.
    """
    
    def __init__(self, config: VehicleDetectionConfig):
        """
        Initialize the detection pipeline.
        
        Args:
            config: Configuration object with all parameters
            
        Raises:
            ModelLoadError: If model cannot be loaded
            ConfigurationError: If configuration is invalid
        """
        self.config = config
        self.model = None
        self.tracker = None
        self.line_zone = None
        self.speed_estimator = None
        self.annotator = None
        self.video_info = None
        
        logger.info(f"Initializing pipeline with config: {config}")
        self._initialize_components()
    
    def _initialize_components(self) -> None:
        """Initialize all pipeline components."""
        try:
            # Load YOLO model
            logger.info(f"Loading YOLO model: {self.config.model_path}")
            self.model = YOLO(self.config.model_path)
            self.model.conf = self.config.confidence_threshold
            self.model.iou = self.config.iou_threshold
            logger.info("Model loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise ModelLoadError(f"Could not load YOLO model from {self.config.model_path}: {e}")
    
    def _setup_video_components(self, video_path: str) -> None:
        """
        Set up video-specific components.
        
        Args:
            video_path: Path to input video
            
        Raises:
            VideoProcessingError: If video cannot be opened
        """
        try:
            # Get video information
            self.video_info = sv.VideoInfo.from_video_path(video_path)
            logger.info(f"Video info: {self.video_info.width}x{self.video_info.height} @ {self.video_info.fps}fps")
            
            # Initialize ByteTrack tracker
            self.tracker = sv.ByteTrack(
                frame_rate=self.video_info.fps,
                track_activation_threshold=self.config.confidence_threshold
            )
            logger.info("Tracker initialized")
            
            # Set up counting line zone
            line_start = sv.Point(
                x=self.config.line_offset,
                y=self.config.line_y
            )
            line_end = sv.Point(
                x=self.video_info.width - self.config.line_offset,
                y=self.config.line_y
            )
            
            self.line_zone = sv.LineZone(
                start=line_start,
                end=line_end,
                triggering_anchors=(sv.Position.BOTTOM_CENTER,)
            )
            logger.info(f"Line zone created at y={self.config.line_y}")
            
            # Initialize perspective transformer
            source_pts = np.array(self.config.source_points, dtype=np.float32)
            target_pts = np.array(self.config.target_points, dtype=np.float32)
            
            transformer = PerspectiveTransformer(
                source_points=source_pts,
                target_points=target_pts
            )
            logger.info("Perspective transformer initialized")
            
            # Initialize speed estimator
            self.speed_estimator = VehicleSpeedEstimator(
                fps=self.video_info.fps,
                transformer=transformer,
                history_duration=self.config.speed_history_seconds,
                speed_unit=self.config.speed_unit
            )
            logger.info("Speed estimator initialized")
            
            # Initialize frame annotator
            self.annotator = FrameAnnotator(
                video_resolution=(self.video_info.width, self.video_info.height),
                show_boxes=self.config.enable_boxes,
                show_labels=self.config.enable_labels,
                show_traces=self.config.enable_traces,
                show_line_zones=self.config.enable_line_zones,
                trace_length=self.config.trace_length,
                zone_polygon=source_pts
            )
            logger.info("Frame annotator initialized")
            
        except Exception as e:
            logger.error(f"Failed to setup video components: {e}")
            raise VideoProcessingError(f"Error setting up video processing: {e}")
    
    def _process_single_frame(self, frame: np.ndarray) -> tuple:
        """
        Process a single video frame.
        
        Args:
            frame: Input video frame
            
        Returns:
            Tuple of (annotated_frame, detections)
        """
        # Run YOLO detection
        results = self.model(frame, verbose=False)[0]
        detections = sv.Detections.from_ultralytics(results)
        
        # Update tracker
        detections = self.tracker.update_with_detections(detections)
        
        # Trigger line zone counting
        self.line_zone.trigger(detections)
        
        # Estimate speeds
        detections = self.speed_estimator.estimate(detections)
        
        # Generate labels
        labels = self._create_labels(detections)
        
        # Annotate frame
        annotated_frame = self.annotator.draw_annotations(
            frame=frame,
            detections=detections,
            labels=labels,
            line_zones=[self.line_zone]
        )
        
        return annotated_frame, detections
    
    def _create_labels(self, detections: sv.Detections) -> list:
        """
        Create display labels for detected vehicles.
        
        Args:
            detections: Detection results
            
        Returns:
            List of label strings
        """
        labels = []
        
        if not hasattr(detections, 'tracker_id') or detections.tracker_id is None:
            return labels
        
        for idx, tracker_id in enumerate(detections.tracker_id):
            # Get class name
            class_name = "Vehicle"
            if "class_name" in detections.data:
                class_name = detections.data["class_name"][idx]
            
            # Get speed
            speed_text = ""
            if "speed" in detections.data:
                speed = detections.data["speed"][idx]
                if speed > 0:
                    speed_text = f" {speed:.0f}{self.config.speed_unit}"
            
            # Create label
            label = f"{class_name} #{tracker_id}{speed_text}"
            labels.append(label)
        
        return labels
    
    def process_video(
        self,
        progress_callback: Optional[Callable[[float], None]] = None
    ) -> Dict:
        """
        Process the entire video.
        
        Args:
            progress_callback: Optional callback for progress updates
            
        Returns:
            Dictionary with processing statistics
            
        Raises:
            VideoProcessingError: If video processing fails
        """
        start_time = time()
        
        try:
            # Validate input video
            if not Path(self.config.input_video).exists():
                raise VideoProcessingError(f"Input video not found: {self.config.input_video}")
            
            # Setup components
            self._setup_video_components(self.config.input_video)
            
            # Create output directory if needed
            output_path = Path(self.config.output_video)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Initialize statistics
            frame_count = 0
            total_frames = self.video_info.total_frames or 0
            all_speeds = []
            
            # Setup display window if enabled (disabled in headless environments like HF Spaces)
            if self.config.display_enabled:
                try:
                    cv2.namedWindow(self.config.window_name, cv2.WINDOW_NORMAL)
                    cv2.resizeWindow(
                        self.config.window_name,
                        self.video_info.width,
                        self.video_info.height
                    )
                except Exception as e:
                    logger.warning(f"Could not create display window (headless environment?): {e}")
                    self.config.display_enabled = False
            
            # Process video
            logger.info("Starting video processing...")
            frame_generator = sv.get_video_frames_generator(self.config.input_video)
            
            with sv.VideoSink(self.config.output_video, self.video_info) as sink:
                for frame in frame_generator:
                    try:
                        # Process frame
                        annotated_frame, detections = self._process_single_frame(frame)
                        
                        # Collect speed statistics
                        if "speed" in detections.data:
                            speeds = detections.data["speed"]
                            all_speeds.extend([s for s in speeds if s > 0])
                        
                        # Write to output
                        sink.write_frame(annotated_frame)
                        
                        # Display if enabled
                        if self.config.display_enabled:
                            cv2.imshow(self.config.window_name, annotated_frame)
                            
                            # Check for quit
                            if cv2.waitKey(1) & 0xFF == ord('q'):
                                logger.info("Processing interrupted by user")
                                break
                            
                            # Check if window was closed
                            if cv2.getWindowProperty(
                                self.config.window_name,
                                cv2.WND_PROP_VISIBLE
                            ) < 1:
                                logger.info("Window closed by user")
                                break
                        
                        # Update progress
                        frame_count += 1
                        if progress_callback and total_frames > 0:
                            progress = frame_count / total_frames
                            progress_callback(progress)
                    
                    except Exception as e:
                        logger.warning(f"Error processing frame {frame_count}: {e}")
                        continue
            
            # Cleanup
            if self.config.display_enabled:
                cv2.destroyAllWindows()
            
            # Calculate statistics
            processing_time = time() - start_time
            stats = {
                'total_count': self.line_zone.in_count + self.line_zone.out_count,
                'in_count': self.line_zone.in_count,
                'out_count': self.line_zone.out_count,
                'avg_speed': np.mean(all_speeds) if all_speeds else 0.0,
                'max_speed': np.max(all_speeds) if all_speeds else 0.0,
                'min_speed': np.min(all_speeds) if all_speeds else 0.0,
                'frames_processed': frame_count,
                'processing_time': processing_time,
                'fps': frame_count / processing_time if processing_time > 0 else 0
            }
            
            logger.info(f"Processing complete: {frame_count} frames in {processing_time:.2f}s")
            logger.info(f"Vehicles counted: {stats['total_count']} (In: {stats['in_count']}, Out: {stats['out_count']})")
            
            return stats
            
        except Exception as e:
            logger.error(f"Video processing failed: {e}", exc_info=True)
            raise VideoProcessingError(f"Failed to process video: {e}")


def process_video(
    config: VehicleDetectionConfig,
    progress_callback: Optional[Callable[[float], None]] = None
) -> Dict:
    """
    Convenience function to process a video with given configuration.
    
    Args:
        config: Configuration object
        progress_callback: Optional progress callback
        
    Returns:
        Processing statistics dictionary
    """
    pipeline = VehicleDetectionPipeline(config)
    return pipeline.process_video(progress_callback)


def main():
    """Main entry point for CLI usage."""
    try:
        logger.info("=" * 60)
        logger.info("Vehicle Speed Estimation & Counting System")
        logger.info("=" * 60)
        
        # Load configuration
        config = VehicleDetectionConfig()
        logger.info(f"Configuration: {config}")
        
        # Process video
        stats = process_video(config)
        
        # Display results
        print("\n" + "=" * 60)
        print("PROCESSING RESULTS")
        print("=" * 60)
        print(f"Output saved to: {config.output_video}")
        print(f"\nVehicle Count:")
        print(f"  Total: {stats['total_count']}")
        print(f"  In: {stats['in_count']}")
        print(f"  Out: {stats['out_count']}")
        print(f"\nSpeed Statistics ({config.speed_unit}):")
        print(f"  Average: {stats['avg_speed']:.1f}")
        print(f"  Maximum: {stats['max_speed']:.1f}")
        print(f"  Minimum: {stats['min_speed']:.1f}")
        print(f"\nProcessing Info:")
        print(f"  Frames: {stats['frames_processed']}")
        print(f"  Time: {stats['processing_time']:.2f}s")
        print(f"  FPS: {stats['fps']:.1f}")
        print("=" * 60)
        
        return 0
        
    except KeyboardInterrupt:
        logger.info("Processing interrupted by user")
        return 1
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        print(f"\n❌ Error: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())