object_detection.py

"""
Object Detection Module for DetectifAI

This module handles:
- Fire detection using fire_YOLO11.pt
- Knife and gun detection using weapon_YOLO11.pt
- Multi-model forking approach for parallel inference
- Integration with video processing pipeline
- Object-based event generation
"""

import cv2
import torch
import numpy as np
import os
import logging
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
from ultralytics import YOLO
import time

logger = logging.getLogger(__name__)

@dataclass
class DetectedObject:
    """Represents a detected object"""
    class_name: str
    confidence: float
    bbox: Tuple[int, int, int, int]  # (x1, y1, x2, y2)
    center_point: Tuple[int, int]
    area: float
    frame_timestamp: float
    detection_model: str

@dataclass
class ObjectDetectionResult:
    """Result of object detection on a frame"""
    frame_path: str
    timestamp: float
    detected_objects: List[DetectedObject]
    total_detections: int
    detection_confidence_avg: float
    processing_time: float

class ObjectDetector:
    """Main object detection class using YOLOv11 models"""
    
    def __init__(self, config):
        """
        Initialize object detector with trained models
        
        Args:
            config: VideoProcessingConfig object with object detection settings
        """
        self.config = config
        self.models = {}
        self.class_names = {}
        self.confidence_threshold = config.object_detection_confidence
        self.device = 'cuda' if torch.cuda.is_available() and config.use_gpu_acceleration else 'cpu'
        
        logger.info(f"Initializing ObjectDetector on device: {self.device}")
        
        # Load models
        self._load_models()
        
        # Statistics
        self.detection_stats = {
            'total_frames_processed': 0,
            'total_objects_detected': 0,
            'detection_times': [],
            'objects_by_class': {},
            'confidence_scores': []
        }
    
    def _load_models(self):
        """Load YOLOv11 models separately: fire_YOLO11.pt and weapon_YOLO11.pt (multi-model forking)"""
        try:
            # Fire detection model
            fire_model_path = os.path.join(self.config.models_dir, "fire_YOLO11.pt")
            if os.path.exists(fire_model_path):
                logger.info(f"Loading fire detection model: {fire_model_path}")
                self.models['fire'] = YOLO(fire_model_path)
                self.models['fire'].to(self.device)
                # Class names mapping for fire model: 0='Fire' (only detecting Fire class, ignoring class 1)
                self.class_names['fire'] = ['Fire']
                logger.info("✅ Fire detection model loaded successfully (detecting only 'Fire' class)")
            else:
                logger.warning(f"Fire model not found at: {fire_model_path}")
            
            # Weapon detection model (gun + knife)
            weapon_model_path = os.path.join(self.config.models_dir, "weapon_YOLO11.pt")
            if os.path.exists(weapon_model_path):
                logger.info(f"Loading weapon detection model: {weapon_model_path}")
                self.models['weapon'] = YOLO(weapon_model_path)
                self.models['weapon'].to(self.device)
                # Class names mapping for weapon model: 0='gun', 1='knife' (CORRECTED ORDER)
                self.class_names['weapon'] = ['gun', 'knife']
                logger.info("✅ Weapon detection model loaded successfully (gun, knife)")
            else:
                logger.warning(f"Weapon model not found at: {weapon_model_path}")
            
            if not self.models:
                logger.error("❌ No object detection models loaded!")
                raise FileNotFoundError("No object detection models found")
            
            logger.info(f"📊 Loaded {len(self.models)} object detection models: {list(self.models.keys())}")
            
        except Exception as e:
            logger.error(f"❌ Failed to load object detection models: {e}")
            raise
    
    def detect_objects_in_frame(self, frame_path: str, timestamp: float) -> ObjectDetectionResult:
        """
        Detect objects in a single frame
        
        Args:
            frame_path: Path to the frame image
            timestamp: Timestamp of the frame in video
            
        Returns:
            ObjectDetectionResult with all detected objects
        """
        start_time = time.time()
        
        # Load frame
        frame = cv2.imread(frame_path)
        if frame is None:
            logger.error(f"Could not load frame: {frame_path}")
            return ObjectDetectionResult(
                frame_path=frame_path,
                timestamp=timestamp,
                detected_objects=[],
                total_detections=0,
                detection_confidence_avg=0.0,
                processing_time=0.0
            )
        
        detected_objects = []
        
        # Run detection with each model
        for model_name, model in self.models.items():
            try:
                # Run inference
                results = model(frame, conf=self.confidence_threshold, verbose=False)
                
                # Process results
                for result in results:
                    if result.boxes is not None:
                        boxes = result.boxes.xyxy.cpu().numpy()  # x1, y1, x2, y2
                        confidences = result.boxes.conf.cpu().numpy()
                        classes = result.boxes.cls.cpu().numpy().astype(int)
                        
                        for i, (box, conf, cls) in enumerate(zip(boxes, confidences, classes)):
                            # For fire model, only process class 0 (Fire), skip class 1
                            if model_name == 'fire' and cls != 0:
                                continue
                            
                            # Get class name
                            if model_name in self.class_names and cls < len(self.class_names[model_name]):
                                class_name = self.class_names[model_name][cls]
                            else:
                                class_name = f"unknown_{cls}"
                            
                            # Apply specific confidence thresholds based on object type
                            confidence_threshold = self.confidence_threshold  # default
                            if class_name.lower() == 'fire':
                                confidence_threshold = getattr(self.config, 'fire_detection_confidence', 0.4)
                            elif class_name in ['knife', 'gun']:
                                confidence_threshold = getattr(self.config, 'weapon_detection_confidence', 0.7)
                            
                            # Skip detection if confidence is below specific threshold
                            if float(conf) < confidence_threshold:
                                continue
                            
                            # Calculate center point and area
                            x1, y1, x2, y2 = box.astype(int)
                            center_x = int((x1 + x2) / 2)
                            center_y = int((y1 + y2) / 2)
                            area = (x2 - x1) * (y2 - y1)
                            
                            detected_object = DetectedObject(
                                class_name=class_name,
                                confidence=float(conf),
                                bbox=(x1, y1, x2, y2),
                                center_point=(center_x, center_y),
                                area=area,
                                frame_timestamp=timestamp,
                                detection_model=model_name
                            )
                            
                            detected_objects.append(detected_object)
                            
                            # Update statistics
                            if class_name not in self.detection_stats['objects_by_class']:
                                self.detection_stats['objects_by_class'][class_name] = 0
                            self.detection_stats['objects_by_class'][class_name] += 1
                            self.detection_stats['confidence_scores'].append(float(conf))
            
            except Exception as e:
                logger.error(f"Error running {model_name} detection: {e}")
                continue
        
        # Calculate processing time and statistics
        processing_time = time.time() - start_time
        self.detection_stats['detection_times'].append(processing_time)
        self.detection_stats['total_frames_processed'] += 1
        self.detection_stats['total_objects_detected'] += len(detected_objects)
        
        # Calculate average confidence
        avg_confidence = np.mean([obj.confidence for obj in detected_objects]) if detected_objects else 0.0
        
        result = ObjectDetectionResult(
            frame_path=frame_path,
            timestamp=timestamp,
            detected_objects=detected_objects,
            total_detections=len(detected_objects),
            detection_confidence_avg=float(avg_confidence),
            processing_time=processing_time
        )
        
        if detected_objects:
            object_summary = ", ".join([f"{obj.class_name}({obj.confidence:.2f})" for obj in detected_objects])
            logger.info(f"🎯 Detected {len(detected_objects)} objects at {timestamp:.2f}s: {object_summary}")
        
        return result
    
    def detect_objects_in_keyframes(self, keyframes: List) -> List[ObjectDetectionResult]:
        """
        Run object detection on all keyframes
        
        Args:
            keyframes: List of KeyframeResult objects from video processing
            
        Returns:
            List of ObjectDetectionResult objects
        """
        logger.info(f"🔍 Running object detection on {len(keyframes)} keyframes")
        
        detection_results = []
        
        for i, keyframe in enumerate(keyframes):
            try:
                frame_path = keyframe.frame_data.frame_path
                timestamp = keyframe.frame_data.timestamp
                
                # Run detection
                result = self.detect_objects_in_frame(frame_path, timestamp)
                detection_results.append(result)
                
                # Progress logging
                if (i + 1) % 10 == 0 or i == len(keyframes) - 1:
                    logger.info(f"📊 Object detection progress: {i + 1}/{len(keyframes)} frames processed")
                
            except Exception as e:
                logger.error(f"Error detecting objects in keyframe {i}: {e}")
                continue
        
        # Log final statistics
        total_objects = sum(r.total_detections for r in detection_results)
        frames_with_objects = sum(1 for r in detection_results if r.total_detections > 0)
        avg_processing_time = np.mean([r.processing_time for r in detection_results]) if detection_results else 0
        
        logger.info(f"🎯 Object Detection Summary:")
        logger.info(f"   📊 Total objects detected: {total_objects}")
        logger.info(f"   📊 Frames with objects: {frames_with_objects}/{len(keyframes)}")
        logger.info(f"   📊 Average processing time: {avg_processing_time:.3f}s per frame")
        logger.info(f"   📊 Objects by class: {self.detection_stats['objects_by_class']}")
        
        return detection_results
    
    def create_object_based_events(self, detection_results: List[ObjectDetectionResult], 
                                 temporal_window: float = 5.0) -> List[Dict[str, Any]]:
        """
        Create events based on object detections
        
        Args:
            detection_results: List of ObjectDetectionResult objects
            temporal_window: Time window for grouping detections (seconds)
            
        Returns:
            List of object-based events
        """
        logger.info(f"🎯 Creating object-based events from {len(detection_results)} detection results")
        
        # Filter results with detections
        results_with_objects = [r for r in detection_results if r.total_detections > 0]
        
        if not results_with_objects:
            logger.info("No objects detected, no object-based events created")
            return []
        
        # Group detections by object class
        events_by_class = {}
        
        for result in results_with_objects:
            for obj in result.detected_objects:
                class_name = obj.class_name
                
                if class_name not in events_by_class:
                    events_by_class[class_name] = []
                
                events_by_class[class_name].append({
                    'timestamp': result.timestamp,
                    'confidence': obj.confidence,
                    'bbox': obj.bbox,
                    'frame_path': result.frame_path,
                    'object': obj
                })
        
        # Create temporal events for each class
        object_events = []
        event_id_counter = 1000  # Start from 1000 to differentiate from motion events
        
        for class_name, detections in events_by_class.items():
            # Sort by timestamp
            detections.sort(key=lambda x: x['timestamp'])
            
            # Group into temporal windows
            current_event_detections = []
            current_event_start = None
            
            for detection in detections:
                timestamp = detection['timestamp']
                
                if current_event_start is None:
                    # Start new event
                    current_event_start = timestamp
                    current_event_detections = [detection]
                elif timestamp - current_event_start <= temporal_window:
                    # Add to current event
                    current_event_detections.append(detection)
                else:
                    # Finish current event and start new one
                    if current_event_detections:
                        event = self._create_event_from_detections(
                            class_name, current_event_detections, event_id_counter
                        )
                        object_events.append(event)
                        event_id_counter += 1
                    
                    # Start new event
                    current_event_start = timestamp
                    current_event_detections = [detection]
            
            # Don't forget the last event
            if current_event_detections:
                event = self._create_event_from_detections(
                    class_name, current_event_detections, event_id_counter
                )
                object_events.append(event)
                event_id_counter += 1
        
        logger.info(f"✅ Created {len(object_events)} object-based events")
        for event in object_events:
            logger.info(f"   🎯 {event['event_type']}: {event['start_timestamp']:.2f}s - {event['end_timestamp']:.2f}s "
                       f"(confidence: {event['confidence']:.2f})")
        
        return object_events
    
    def _create_event_from_detections(self, class_name: str, detections: List[Dict], 
                                    event_id: int) -> Dict[str, Any]:
        """Create an event from a group of detections"""
        start_time = min(d['timestamp'] for d in detections)
        end_time = max(d['timestamp'] for d in detections)
        confidences = [d['confidence'] for d in detections]
        avg_confidence = np.mean(confidences)
        max_confidence = max(confidences)
        
        # Determine event type and importance
        event_type = f"{class_name}_detection"
        importance_score = max_confidence * len(detections) * 2.0  # Higher importance for object events
        
        # Get keyframes with detections
        keyframes = [d['frame_path'] for d in detections]
        
        # Create description
        description = f"{class_name.title()} detected with {avg_confidence:.2f} average confidence over {len(detections)} frames"
        
        return {
            'event_id': f"obj_event_{event_id:04d}",
            'start_timestamp': start_time,
            'end_timestamp': end_time,
            'event_type': event_type,
            'confidence': avg_confidence,
            'max_confidence': max_confidence,
            'keyframes': keyframes,
            'importance_score': importance_score,
            'motion_intensity': 0.0,  # Object events don't have motion intensity
            'description': description,
            'object_class': class_name,
            'detection_count': len(detections),
            'duration': end_time - start_time,
            'detection_details': detections
        }
    
    def get_detection_statistics(self) -> Dict[str, Any]:
        """Get comprehensive detection statistics"""
        stats = self.detection_stats.copy()
        
        if stats['detection_times']:
            stats['avg_detection_time'] = np.mean(stats['detection_times'])
            stats['max_detection_time'] = max(stats['detection_times'])
            stats['min_detection_time'] = min(stats['detection_times'])
        
        if stats['confidence_scores']:
            stats['avg_confidence'] = np.mean(stats['confidence_scores'])
            stats['max_confidence'] = max(stats['confidence_scores'])
            stats['min_confidence'] = min(stats['confidence_scores'])
        
        return stats
    
    def annotate_frame_with_detections(self, frame_path: str, 
                                     detection_result: ObjectDetectionResult,
                                     output_path: str = None) -> str:
        """
        Annotate frame with bounding boxes and labels
        
        Args:
            frame_path: Path to input frame
            detection_result: ObjectDetectionResult for the frame
            output_path: Optional output path, auto-generated if None
            
        Returns:
            Path to annotated frame
        """
        frame = cv2.imread(frame_path)
        if frame is None:
            logger.error(f"Could not load frame for annotation: {frame_path}")
            return frame_path
        
        # Draw bounding boxes and labels
        for obj in detection_result.detected_objects:
            x1, y1, x2, y2 = obj.bbox
            
            # Choose color based on object class (BGR format)
            color_map = {
                'fire': (255, 255, 0),    # Neon Cyan/Blue
                'knife': (0, 255, 255),   # Neon Yellow
                'gun': (0, 255, 0)        # Neon Green
            }
            color = color_map.get(obj.class_name, (255, 255, 255))  # Default white
            
            # Draw bounding box
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            
            # Draw label with confidence
            label = f"{obj.class_name}: {obj.confidence:.2f}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
            
            # Draw label background
            cv2.rectangle(frame, (x1, y1 - label_size[1] - 10), 
                         (x1 + label_size[0], y1), color, -1)
            
            # Draw label text
            cv2.putText(frame, label, (x1, y1 - 5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        # Generate output path if not provided
        if output_path is None:
            base_name = os.path.splitext(os.path.basename(frame_path))[0]
            output_dir = os.path.dirname(frame_path)
            output_path = os.path.join(output_dir, f"{base_name}_annotated.jpg")
        
        # Save annotated frame
        cv2.imwrite(output_path, frame)
        return output_path


class ObjectDetectionIntegrator:
    """Integration layer between object detection and video processing pipeline"""
    
    def __init__(self, config):
        self.config = config
        self.detector = ObjectDetector(config) if config.enable_object_detection else None
    
    def process_keyframes_with_object_detection(self, keyframes: List) -> Tuple[List, List[Dict[str, Any]]]:
        """
        Process keyframes with object detection and create object-based events
        
        Args:
            keyframes: List of KeyframeResult objects
            
        Returns:
            Tuple of (detection_results, object_events)
        """
        if not self.config.enable_object_detection or not self.detector:
            logger.info("Object detection disabled, skipping...")
            return [], []
        
        logger.info("🎯 Starting object detection integration")
        
        # Run object detection on keyframes
        detection_results = self.detector.detect_objects_in_keyframes(keyframes)
        
        # Create annotated frames for keyframes WITH detections
        annotated_frames = []
        frames_with_detections = []
        
        for result in detection_results:
            if result.total_detections > 0:
                # Create annotated version of the frame
                annotated_path = self.detector.annotate_frame_with_detections(
                    result.frame_path, result
                )
                
                # Store metadata about frames with detections
                frames_with_detections.append({
                    'original_path': result.frame_path,
                    'annotated_path': annotated_path,
                    'timestamp': result.timestamp,
                    'detection_count': result.total_detections,
                    'objects': [obj.class_name for obj in result.detected_objects],
                    'confidence_avg': result.detection_confidence_avg
                })
                
                annotated_frames.append(annotated_path)
                
                logger.info(f"🎯 Annotated frame at {result.timestamp:.2f}s with {result.total_detections} detections")
        
        # Create object-based events
        object_events = self.detector.create_object_based_events(
            detection_results, 
            temporal_window=self.config.object_event_temporal_window
        )
        
        # Store detection metadata in config for later retrieval
        if hasattr(self.config, 'output_base_dir'):
            detection_metadata = {
                'total_keyframes': len(keyframes),
                'frames_with_detections': len(frames_with_detections),
                'annotated_frames': annotated_frames,
                'detection_summary': frames_with_detections,
                'objects_detected': self.detector.detection_stats['objects_by_class'].copy()
            }
            
            # Save metadata to output directory
            metadata_path = os.path.join(self.config.output_base_dir, 'detection_metadata.json')
            os.makedirs(os.path.dirname(metadata_path), exist_ok=True)
            
            import json
            with open(metadata_path, 'w') as f:
                json.dump(detection_metadata, f, indent=2)
                
            logger.info(f"📊 Detection metadata saved: {metadata_path}")
        
        logger.info(f"✅ Object detection integration complete: {len(object_events)} events created")
        logger.info(f"📊 Annotated {len(annotated_frames)} frames with detections out of {len(keyframes)} total keyframes")
        
        return detection_results, object_events
    
    def create_annotated_video(self, video_path: str, detection_results: List, output_path: str = None) -> str:
        """
        Create an annotated video with bounding boxes drawn on frames with detections
        
        Args:
            video_path: Path to the original video
            detection_results: List of ObjectDetectionResult from keyframe detection
            output_path: Optional output path for annotated video
            
        Returns:
            Path to the created annotated video
        """
        if not self.detector or not detection_results:
            logger.warning("No detector or detection results available for video annotation")
            return None
        
        logger.info(f"🎨 Creating annotated video with bounding boxes...")
        
        # Open input video
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            logger.error(f"Cannot open video: {video_path}")
            return None
        
        # Get video properties
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # Build detection lookup by timestamp
        detection_lookup = {}
        for result in detection_results:
            if result.total_detections > 0:
                detection_lookup[result.timestamp] = result
        
        # Create output path if not provided
        if output_path is None:
            video_dir = os.path.dirname(video_path)
            video_name = os.path.splitext(os.path.basename(video_path))[0]
            output_path = os.path.join(video_dir, f"{video_name}_annotated.mp4")
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # Create video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        if not out.isOpened():
            logger.error(f"Cannot create output video: {output_path}")
            cap.release()
            return None
        
        frame_count = 0
        frames_annotated = 0
        
        logger.info(f"Processing {total_frames} frames at {fps} FPS...")
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Calculate timestamp
            timestamp = round(frame_count / fps, 2)
            
            # Check if this timestamp has detections
            if timestamp in detection_lookup:
                result = detection_lookup[timestamp]
                
                # Draw bounding boxes and labels
                for obj in result.detected_objects:
                    x1, y1, x2, y2 = obj.bbox
                    
                    # Choose color based on object class (BGR format)
                    color_map = {
                        'fire': (255, 255, 0),    # Neon Cyan/Blue
                        'knife': (0, 255, 255),   # Neon Yellow
                        'gun': (0, 255, 0)        # Neon Green
                    }
                    color = color_map.get(obj.class_name, (255, 255, 255))
                    
                    # Draw bounding box
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    
                    # Draw label with confidence
                    label = f"{obj.class_name}: {obj.confidence:.2f}"
                    label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                    
                    # Draw label background
                    cv2.rectangle(frame, (x1, y1 - label_size[1] - 10), 
                                 (x1 + label_size[0], y1), color, -1)
                    
                    # Draw label text
                    cv2.putText(frame, label, (x1, y1 - 5), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                
                frames_annotated += 1
            
            # Write frame to output video
            out.write(frame)
            frame_count += 1
            
            # Progress logging
            if frame_count % 100 == 0:
                progress = (frame_count / total_frames) * 100
                logger.info(f"Progress: {progress:.1f}% ({frame_count}/{total_frames} frames)")
        
        # Release resources
        cap.release()
        out.release()
        
        logger.info(f"✅ Annotated video created: {output_path}")
        logger.info(f"📊 Annotated {frames_annotated} frames out of {total_frames} total frames")
        
        return output_path
    
    def get_object_detection_summary(self) -> Dict[str, Any]:
        """Get summary of object detection results"""
        if not self.detector:
            return {'enabled': False}
        
        stats = self.detector.get_detection_statistics()
        stats['enabled'] = True
        return stats