processor.py

"""
UrbanFlow Main Video Processor Module
======================================
This module orchestrates the entire traffic monitoring pipeline.

It integrates all modules:
- Vehicle Detection (YOLO)
- Vehicle Tracking (Centroid-based)
- Lane Detection (OpenCV)
- Traffic Density Estimation
- Adaptive Signal Control
- Violation Detection and Logging

Author: UrbanFlow Team
Version: 1.0.0
"""

import os
import cv2
import numpy as np
from typing import Optional, Callable, Tuple, List, Dict
from tqdm import tqdm

# Import UrbanFlow modules
from config import (
    VEHICLE_CONFIG, DENSITY_CONFIG, LANE_CONFIG, VIOLATION_CONFIG,
    TRACKING_CONFIG, OUTPUT_CONFIG,
    get_signal_timing, get_density_label, ensure_directories
)
from utils.vehicle_detector import VehicleDetector, check_plate_visibility, draw_plate_region
from utils.lane_detector import LaneDetector, check_lane_violation, draw_violation_marker
from utils.tracker import VehicleTracker, draw_rash_driving_marker
from utils.violation_logger import ViolationLogger


class VideoProcessor:
    """
    Main video processing pipeline for UrbanFlow.
    
    This class orchestrates the entire traffic monitoring workflow:
    1. Reads video frames
    2. Detects vehicles using YOLO
    3. Tracks vehicles across frames
    4. Detects lanes using OpenCV
    5. Estimates traffic density
    6. Controls adaptive traffic signal
    7. Detects violations (lane, rash driving, plate visibility)
    8. Logs violations to CSV
    9. Generates output video with overlays
    
    Attributes:
        video_path: Path to input video file
        output_path: Path to output video file
        detector: VehicleDetector instance
        tracker: VehicleTracker instance
        lane_detector: LaneDetector instance
        logger: ViolationLogger instance
        signal_timer: Current signal countdown value
    """
    
    def __init__(
        self,
        video_path: str,
        output_path: str = "urbanflow_output.mp4",
        model_path: str = None,
        frame_skip: int = 1,
        confidence_threshold: float = None
    ):
        """
        Initialize the Video Processor.
        
        Args:
            video_path: Path to input MP4 video file
            output_path: Path for saving processed output video
            model_path: Path to YOLO model weights (default from config)
            frame_skip: Process every Nth frame for speed optimization
            confidence_threshold: Detection confidence threshold
        """
        self.video_path = video_path
        self.output_path = output_path
        self.frame_skip = max(1, frame_skip)
        
        # Initialize video capture
        self.cap = cv2.VideoCapture(video_path)
        if not self.cap.isOpened():
            raise ValueError(f"Cannot open video file: {video_path}")
        
        # Get video properties
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        print(f"[VideoProcessor] Video loaded: {self.width}x{self.height} @ {self.fps}fps")
        print(f"[VideoProcessor] Total frames: {self.total_frames}")
        
        # Initialize modules
        model = model_path or VEHICLE_CONFIG.model_path
        conf_thresh = confidence_threshold or VEHICLE_CONFIG.confidence_threshold
        
        self.detector = VehicleDetector(
            model_path=model,
            confidence_threshold=conf_thresh,
            classes_to_detect=VEHICLE_CONFIG.classes_to_detect
        )
        
        self.tracker = VehicleTracker(
            max_distance=TRACKING_CONFIG.max_tracking_distance,
            max_disappeared=TRACKING_CONFIG.max_disappeared,
            speed_smoothing=TRACKING_CONFIG.speed_smoothing_frames
        )
        
        self.lane_detector = LaneDetector(
            canny_low=LANE_CONFIG.canny_low_threshold,
            canny_high=LANE_CONFIG.canny_high_threshold,
            hough_rho=LANE_CONFIG.hough_rho,
            hough_theta=LANE_CONFIG.hough_theta,
            hough_threshold=LANE_CONFIG.hough_threshold,
            hough_min_line_length=LANE_CONFIG.hough_min_line_length,
            hough_max_line_gap=LANE_CONFIG.hough_max_line_gap
        )
        
        # Set lane detector ROI based on frame size
        self.lane_detector.set_roi_vertices(self.height, self.width)
        
        # Initialize violation logger
        self.logger = ViolationLogger(log_file="violations.csv")
        
        # Initialize video writer
        fourcc = cv2.VideoWriter_fourcc(*OUTPUT_CONFIG.output_codec)
        self.writer = cv2.VideoWriter(
            output_path, fourcc, self.fps // self.frame_skip, (self.width, self.height)
        )
        
        # Adaptive signal state
        self.signal_timer = 0
        self.current_density = "LOW"
        self.target_green_time = DENSITY_CONFIG.low_green_time
        self.current_vehicle_count = 0
        
        # Processing state
        self.frame_count = 0
        self.processed_count = 0
        
    def process_video(
        self,
        progress_callback: Optional[Callable[[int, int, np.ndarray], None]] = None
    ) -> Tuple[str, str]:
        """
        Process the entire video.
        
        Args:
            progress_callback: Optional callback function(current_frame, total_frames, current_frame_data)
            
        Returns:
            Tuple of (csv_log_path, output_video_path)
        """
        print("[VideoProcessor] Starting video processing...")
        
        # Reset tracker for new video
        self.tracker.reset()
        
        # Progress bar for console
        pbar = tqdm(total=self.total_frames, desc="Processing", unit="frames")
        
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            
            self.frame_count += 1
            
            # Skip frames if needed
            if self.frame_count % self.frame_skip != 0:
                pbar.update(1)
                continue
            
            self.processed_count += 1
            
            # Process this frame
            processed_frame = self._process_frame(frame)
            
            # Write to output video
            self.writer.write(processed_frame)
            
            # Update progress
            pbar.update(self.frame_skip)
            
            # Call progress callback if provided (for UI updates)
            if progress_callback and self.frame_count % 5 == 0:
                progress_callback(
                    self.frame_count, 
                    self.total_frames, 
                    processed_frame,
                    density=self.current_density,
                    signal_time=self.target_green_time,
                    vehicle_count=self.current_vehicle_count
                )
        
        pbar.close()
        
        # Release resources
        self.cap.release()
        self.writer.release()
        
        # Save violations
        csv_path = self.logger.save_to_csv()
        
        # Print summary
        self.logger.print_summary()
        
        print(f"[VideoProcessor] Processing complete!")
        print(f"[VideoProcessor] Output video: {self.output_path}")
        print(f"[VideoProcessor] Violation log: {csv_path}")
        
        return csv_path, self.output_path
    
    def _process_frame(self, frame: np.ndarray) -> np.ndarray:
        """
        Process a single video frame.
        
        Args:
            frame: Input video frame (BGR format)
            
        Returns:
            Processed frame with all overlays
        """
        output_frame = frame.copy()
        
        # =================================================================
        # 1. VEHICLE DETECTION (YOLO)
        # =================================================================
        detections = self.detector.detect(frame)
        
        # Get vehicle counts
        vehicle_counts = self.detector.get_vehicle_count(detections)
        total_vehicles = vehicle_counts["total"]
        
        # =================================================================
        # 2. VEHICLE TRACKING
        # =================================================================
        tracks = self.tracker.update(detections)
        
        # =================================================================
        # 3. TRAFFIC DENSITY ESTIMATION & ADAPTIVE SIGNAL
        # =================================================================
        # Calculate density and signal time using new rules
        density_level = get_density_label(total_vehicles)
        green_time = get_signal_timing(total_vehicles)
        
        # Update class state variables
        self.current_density = density_level
        self.target_green_time = green_time
        self.current_vehicle_count = total_vehicles
        
        # Print signal state to console
        if self.frame_count % self.fps < self.frame_skip:
            print(f"Vehicle Count: {total_vehicles} | Density: {density_level} | Green Signal Time: {green_time} seconds")
        
        # =================================================================
        # 4. LANE DETECTION
        # =================================================================
        lane_overlay, lane_lines, lane_boundaries = self.lane_detector.detect_lanes(frame)
        output_frame = self.lane_detector.draw_lanes(output_frame, lane_overlay)
        
        # Get lane boundaries from detection
        left_boundary, right_boundary = lane_boundaries
        output_frame = self.lane_detector.draw_lane_boundaries(
            output_frame, left_boundary, right_boundary
        )
        
        # =================================================================
        # 5. PROCESS EACH VEHICLE
        # =================================================================
        for track_id, track_info in tracks.items():
            detection = track_info["detection"]
            bbox = detection["bbox"]
            class_name = detection["class_name"]
            confidence = detection["confidence"]
            center_x, center_y = track_info["centroid"]
            speed = track_info["smoothed_speed"]
            
            # Draw vehicle bounding box
            color = OUTPUT_CONFIG.class_colors.get(
                class_name, OUTPUT_CONFIG.class_colors["default"]
            )
            x1, y1, x2, y2 = bbox
            cv2.rectangle(output_frame, (x1, y1), (x2, y2), color, 2)
            
            # Draw label
            label = f"{class_name} {confidence:.2f}"
            cv2.putText(output_frame, label, (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
            
            # -----------------------------------------------------
            # 5a. NUMBER PLATE VISIBILITY CHECK
            # -----------------------------------------------------
            plate_region = self.detector.get_plate_region(detection)
            if plate_region:
                is_visible, brightness = check_plate_visibility(
                    frame, plate_region, VIOLATION_CONFIG.plate_brightness_threshold
                )
                
                # Draw plate region
                output_frame = draw_plate_region(output_frame, plate_region, is_visible)
                
                # Log violation if plate not visible
                if not is_visible:
                    self.logger.log_plate_not_visible(
                        frame_number=self.frame_count,
                        vehicle_id=track_id,
                        vehicle_class=class_name,
                        confidence=confidence,
                        bbox=bbox,
                        brightness=brightness,
                        brightness_threshold=VIOLATION_CONFIG.plate_brightness_threshold
                    )
            
            # -----------------------------------------------------
            # 5b. RASH DRIVING DETECTION
            # -----------------------------------------------------
            is_rash = self.tracker.check_rash_driving(
                track_id,
                speed_threshold=VIOLATION_CONFIG.rash_speed_threshold,
                consecutive_frames=VIOLATION_CONFIG.rash_frames_threshold
            )
            
            if is_rash:
                # Draw rash driving marker
                output_frame = draw_rash_driving_marker(output_frame, bbox)
                
                # Log violation
                self.logger.log_rash_driving(
                    frame_number=self.frame_count,
                    vehicle_id=track_id,
                    vehicle_class=class_name,
                    confidence=confidence,
                    bbox=bbox,
                    speed=speed,
                    speed_threshold=VIOLATION_CONFIG.rash_speed_threshold
                )
            
            # -----------------------------------------------------
            # 5c. LANE VIOLATION DETECTION (using centroid with tolerance)
            # -----------------------------------------------------
            is_violation, boundary = check_lane_violation(
                center_x, center_y, left_boundary, right_boundary, tolerance=15
            )
            
            if is_violation:
                print(f"Lane Violation Detected: Vehicle {track_id} crossed {boundary} boundary at frame {self.frame_count}")
                
                # Draw violation marker
                output_frame = draw_violation_marker(
                    output_frame, bbox, "Lane Violation"
                )
                
                # Log violation
                self.logger.log_lane_violation(
                    frame_number=self.frame_count,
                    vehicle_id=track_id,
                    vehicle_class=class_name,
                    confidence=confidence,
                    bbox=bbox,
                    boundary_crossed=boundary
                )
            
            # -----------------------------------------------------
            # 5d. DRAW TRACKING INFO
            # -----------------------------------------------------
            output_frame = self.tracker.draw_tracks(
                output_frame,
                {track_id: track_info},
                show_id=True,
                show_speed=True
            )
        
        # =================================================================
        # 6. DRAW HUD OVERLAYS
        # =================================================================
        output_frame = self._draw_hud(output_frame, total_vehicles, density_level, green_time)
        
        return output_frame
    
    def _draw_hud(self, frame: np.ndarray, vehicle_count: int, density_level: str, green_time: int) -> np.ndarray:
        """
        Draw heads-up display with system information.
        
        Args:
            frame: Video frame
            vehicle_count: Total number of vehicles detected
            
        Returns:
            Frame with HUD overlays
        """
        # Create overlay background
        overlay = frame.copy()
        
        # HUD position (top-left)
        hud_x, hud_y = 10, 10
        hud_width, hud_height = 380, 180
        
        # Draw semi-transparent background
        cv2.rectangle(
            overlay,
            (hud_x, hud_y),
            (hud_x + hud_width, hud_y + hud_height),
            (0, 0, 0),
            -1
        )
        cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)
        
        # Color based on density
        if density_level == "LOW":
            density_color = (0, 255, 0)  # Green
        elif density_level == "MEDIUM":
            density_color = (0, 165, 255)  # Orange
        else:  # HIGH
            density_color = (0, 0, 255)  # Red
        
        # Draw information
        text_y = hud_y + 30
        line_height = 30
        
        # 1. Vehicle Count
        cv2.putText(
            frame,
            f"Vehicle Count: {vehicle_count}",
            (hud_x + 10, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            2
        )
        text_y += line_height
        
        # 2. Traffic Density
        cv2.putText(
            frame,
            f"Density: {density_level}",
            (hud_x + 10, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            density_color,
            2
        )
        text_y += line_height
        
        # 3. Green Signal Time
        timer_color = (0, 255, 0) if green_time >= 30 else (0, 165, 255)
        cv2.putText(
            frame,
            f"Green Signal Time: {green_time}s",
            (hud_x + 10, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            timer_color,
            2
        )
        text_y += line_height
        
        # 4. Signal Duration Info
        cv2.putText(
            frame,
            f"Signal Duration: {green_time}s",
            (hud_x + 10, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (200, 200, 200),
            1
        )
        text_y += line_height
        
        # 5. Frame info
        cv2.putText(
            frame,
            f"Frame: {self.frame_count}",
            (hud_x + 10, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (150, 150, 150),
            1
        )
        
        # Draw large prominent signal time at bottom-center
        signal_text = f"Green Light Duration: {green_time}s ({density_level})"
        
        # Get text size for centering
        font_scale = 1.1
        thickness = 3
        (text_width, text_height), _ = cv2.getTextSize(
            signal_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness
        )
        
        # Position at bottom-center
        text_x = 50
        text_y_bottom = 120
        
        # Draw black outline for visibility
        outline_color = (0, 0, 0)
        outline_thickness = thickness + 2
        cv2.putText(
            frame,
            signal_text,
            (text_x, text_y_bottom),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            outline_color,
            outline_thickness
        )
        
        # Draw main text in green/white/red based on density
        main_color = (0, 255, 0) if density_level == "LOW" else \
                     (0, 255, 255) if density_level == "MEDIUM" else (0, 0, 255)
        cv2.putText(
            frame,
            signal_text,
            (text_x, text_y_bottom),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            main_color,
            thickness
        )
        
        # UrbanFlow branding
        cv2.putText(
            frame,
            "UrbanFlow AI",
            (self.width - 150, self.height - 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 255, 255),
            1
        )
        
        return frame
    
    def get_video_info(self) -> Dict:
        """
        Get information about the input video.
        
        Returns:
            Dictionary with video metadata
        """
        return {
            "path": self.video_path,
            "width": self.width,
            "height": self.height,
            "fps": self.fps,
            "total_frames": self.total_frames,
            "duration_seconds": self.total_frames / self.fps if self.fps > 0 else 0
        }


# =============================================================================
# Standalone Processing Function
# =============================================================================

def process_video_file(
    video_path: str,
    output_path: str = "urbanflow_output.mp4",
    model_path: str = None,
    frame_skip: int = 1
) -> Tuple[str, str]:
    """
    Convenience function to process a video file.
    
    Args:
        video_path: Path to input video
        output_path: Path for output video
        model_path: Path to YOLO model
        frame_skip: Frame skip for processing speed
        
    Returns:
        Tuple of (csv_log_path, output_video_path)
    """
    processor = VideoProcessor(
        video_path=video_path,
        output_path=output_path,
        model_path=model_path,
        frame_skip=frame_skip
    )
    
    return processor.process_video()