Winning HackIreland 2024: Building Invision-Transformer CCTV AI 🏆🎥🤖

import torch
import torch.nn as nn
from transformers import (
    ViTModel, 
    GPT2LMHeadModel, 
    GPT2Tokenizer,
    AutoImageProcessor
)
import cv2
import numpy as np
from typing import List, Dict, Tuple, Optional

class InvisionTransformer(nn.Module):
    """
    Hybrid architecture combining Vision Transformer with GPT-2
    for CCTV footage analysis and natural language alert generation
    """
    
    def __init__(self, 
                 vit_model_name: str = "google/vit-base-patch16-224",
                 gpt2_model_name: str = "gpt2-medium"):
        super(InvisionTransformer, self).__init__()
        
        # Vision component - Understanding what's happening
        self.vision_transformer = ViTModel.from_pretrained(vit_model_name)
        self.image_processor = AutoImageProcessor.from_pretrained(vit_model_name)
        
        # Language component - Describing what's seen
        self.language_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
        self.tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # Cross-modal fusion layer
        self.vision_projection = nn.Linear(768, 768)  # ViT hidden size
        self.temporal_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=768, nhead=8),
            num_layers=2
        )
        
        # Security event classifier
        self.event_classifier = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, len(self.security_event_types))
        )
        
        self.security_event_types = [
            "normal_activity",
            "suspicious_loitering", 
            "unattended_object",
            "unauthorized_access",
            "person_fall",
            "crowd_gathering",
            "vehicle_intrusion",
            "fire_smoke_detection"
        ]
    
    def process_video_sequence(self, video_frames: torch.Tensor) -> Dict:
        """Process a sequence of video frames to detect security events"""
        batch_size, seq_len, channels, height, width = video_frames.shape
        
        # Process each frame through Vision Transformer
        frame_features = []
        for i in range(seq_len):
            frame = video_frames[:, i]  # Shape: (batch_size, C, H, W)
            
            # Get ViT features
            vit_outputs = self.vision_transformer(pixel_values=frame)
            frame_feature = vit_outputs.last_hidden_state[:, 0]  # CLS token
            
            frame_features.append(frame_feature)
        
        # Stack frame features: (batch_size, seq_len, hidden_size)
        sequence_features = torch.stack(frame_features, dim=1)
        
        # Apply temporal encoding to understand sequence patterns
        # Reshape for transformer: (seq_len, batch_size, hidden_size)
        sequence_features = sequence_features.transpose(0, 1)
        temporal_features = self.temporal_encoder(sequence_features)
        
        # Get sequence representation (average across time)
        sequence_repr = temporal_features.mean(dim=0)  # (batch_size, hidden_size)
        
        # Classify security event type
        event_logits = self.event_classifier(sequence_repr)
        event_probs = torch.softmax(event_logits, dim=-1)
        
        return {
            'sequence_features': sequence_repr,
            'event_logits': event_logits,
            'event_probabilities': event_probs,
            'temporal_features': temporal_features.transpose(0, 1)
        }

class VideoAnalysisProcessor:
    def __init__(self, model: InvisionTransformer):
        self.model = model
        self.frame_buffer = []
        self.sequence_length = 16  # Analyze 16-frame sequences
        self.suspicious_threshold = 0.7
        
    def process_live_feed(self, video_stream_url: str) -> None:
        """Process live CCTV feed in real-time"""
        cap = cv2.VideoCapture(video_stream_url)
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Preprocess frame
            processed_frame = self._preprocess_frame(frame)
            self.frame_buffer.append(processed_frame)
            
            # Analyze when we have enough frames
            if len(self.frame_buffer) >= self.sequence_length:
                analysis_result = self._analyze_sequence(
                    self.frame_buffer[-self.sequence_length:]
                )
                
                # Generate alert if suspicious activity detected
                if analysis_result['max_suspicion'] > self.suspicious_threshold:
                    alert = self._generate_natural_language_alert(analysis_result)
                    self._send_security_alert(alert)
                
                # Sliding window approach
                self.frame_buffer = self.frame_buffer[-8:]  # Keep half for overlap
    
    def _preprocess_frame(self, frame: np.ndarray) -> torch.Tensor:
        """Preprocess video frame for Vision Transformer"""
        # Resize to expected input size
        frame_resized = cv2.resize(frame, (224, 224))
        
        # Convert BGR to RGB
        frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
        
        # Normalize pixel values
        frame_normalized = frame_rgb / 255.0
        
        # Convert to tensor and add batch dimension
        frame_tensor = torch.from_numpy(frame_normalized).float()
        frame_tensor = frame_tensor.permute(2, 0, 1)  # HWC -> CHW
        
        return frame_tensor
    
    def _analyze_sequence(self, frame_sequence: List[torch.Tensor]) -> Dict:
        """Analyze a sequence of frames for security events"""
        # Stack frames into sequence tensor
        sequence_tensor = torch.stack(frame_sequence).unsqueeze(0)  # Add batch dim
        
        with torch.no_grad():
            analysis = self.model.process_video_sequence(sequence_tensor)
        
        # Interpret results
        event_probs = analysis['event_probabilities'][0]  # Remove batch dim
        max_prob_idx = torch.argmax(event_probs)
        max_prob = event_probs[max_prob_idx].item()
        
        detected_event = self.model.security_event_types[max_prob_idx]
        
        return {
            'detected_event': detected_event,
            'confidence': max_prob,
            'max_suspicion': max_prob if detected_event != 'normal_activity' else 0,
            'all_probabilities': event_probs.tolist(),
            'sequence_features': analysis['sequence_features']
        }

class SecurityAlertGenerator:
    def __init__(self, language_model, tokenizer):
        self.language_model = language_model
        self.tokenizer = tokenizer
        
        # Template prompts for different security events
        self.alert_templates = {
            "suspicious_loitering": [
                "A person has been loitering in the {location} area for {duration} minutes. Possible suspicious activity.",
                "Individual detected remaining stationary near {location} for an extended period ({duration} minutes).",
                "Suspicious behavior: Person lingering at {location} for {duration} minutes without clear purpose."
            ],
            "unattended_object": [
                "Unattended bag or object detected in {location} for {duration} minutes. Potential security risk.",
                "Object left unattended at {location}. Duration: {duration} minutes. Requires investigation.",
                "Security alert: Suspicious package at {location}, unattended for {duration} minutes."
            ],
            "person_fall": [
                "Person down detected in {location}. Possible medical emergency requiring immediate attention.",
                "Fall incident detected at {location}. Individual appears to need medical assistance.",
                "Medical alert: Person has fallen in {location} area. Emergency response may be required."
            ],
            "unauthorized_access": [
                "Unauthorized access attempt detected at {location}. Security breach possible.",
                "Access control violation at {location}. Individual attempting entry without proper authorization.",
                "Security warning: Unauthorized person attempting to access restricted area at {location}."
            ]
        }
    
    def generate_contextual_alert(self, 
                                event_type: str, 
                                confidence: float,
                                location: str = "monitored area",
                                duration: int = None,
                                additional_context: Dict = None) -> str:
        """Generate natural language security alert with context"""
        
        if event_type == "normal_activity":
            return None  # No alert needed
        
        # Select appropriate template
        templates = self.alert_templates.get(event_type, [
            f"Security event detected: {event_type} in {location}"
        ])
        
        # Choose template based on confidence (higher confidence -> more urgent language)
        if confidence > 0.9:
            template_idx = 0  # Most urgent
        elif confidence > 0.8:
            template_idx = min(1, len(templates) - 1)
        else:
            template_idx = min(2, len(templates) - 1)
        
        base_alert = templates[template_idx]
        
        # Fill in context variables
        alert = base_alert.format(
            location=location,
            duration=duration or "unknown",
            confidence=f"{confidence:.1%}"
        )
        
        # Add confidence level
        if confidence < 0.8:
            alert += f" (Confidence: {confidence:.1%})"
        
        # Use language model to enhance the alert if needed
        if additional_context:
            enhanced_alert = self._enhance_with_language_model(alert, additional_context)
            return enhanced_alert
        
        return alert
    
    def _enhance_with_language_model(self, base_alert: str, context: Dict) -> str:
        """Use GPT-2 to enhance alert with additional context"""
        # Create prompt for language model
        prompt = f"Security Alert Context: {base_alert}\nAdditional Details: {context}\nEnhanced Alert:"
        
        # Tokenize prompt
        inputs = self.tokenizer.encode(prompt, return_tensors='pt')
        
        # Generate enhanced alert
        with torch.no_grad():
            outputs = self.language_model.generate(
                inputs,
                max_length=inputs.shape[1] + 50,
                num_return_sequences=1,
                temperature=0.7,
                pad_token_id=self.tokenizer.eos_token_id,
                do_sample=True
            )
        
        # Decode and extract enhanced alert
        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        enhanced_alert = generated_text.split("Enhanced Alert:")[-1].strip()
        
        return enhanced_alert if enhanced_alert else base_alert

class EdgeOptimizedInvision:
    def __init__(self, model_path: str):
        # Load quantized model for faster inference
        self.model = torch.jit.load(model_path)
        self.model.eval()
        
        # Initialize processing queues
        self.frame_queue = queue.Queue(maxsize=100)
        self.alert_queue = queue.Queue(maxsize=50)
        
        # Start processing threads
        self.processing_thread = threading.Thread(target=self._process_frames)
        self.alert_thread = threading.Thread(target=self._handle_alerts)
        
        self.processing_thread.start()
        self.alert_thread.start()
    
    def _process_frames(self):
        """Background thread for processing video frames"""
        frame_buffer = []
        
        while True:
            try:
                frame = self.frame_queue.get(timeout=1)
                frame_buffer.append(frame)
                
                if len(frame_buffer) >= 16:  # Process 16-frame sequences
                    result = self._analyze_sequence_optimized(frame_buffer)
                    
                    if result['requires_alert']:
                        self.alert_queue.put(result)
                    
                    # Sliding window
                    frame_buffer = frame_buffer[8:]
                    
            except queue.Empty:
                continue
    
    def _analyze_sequence_optimized(self, frames: List[np.ndarray]) -> Dict:
        """Optimized sequence analysis for edge deployment"""
        # Convert frames to tensor batch
        frame_tensors = []
        for frame in frames:
            # Aggressive preprocessing for speed
            frame_small = cv2.resize(frame, (112, 112))  # Smaller input
            frame_tensor = torch.from_numpy(frame_small).float() / 255.0
            frame_tensors.append(frame_tensor.permute(2, 0, 1))
        
        sequence_tensor = torch.stack(frame_tensors).unsqueeze(0)
        
        # Fast inference with reduced precision
        with torch.no_grad():
            with torch.cuda.amp.autocast():  # Mixed precision
                outputs = self.model(sequence_tensor)
        
        # Quick thresholding for alerts
        max_prob = torch.max(outputs['event_probabilities'])
        requires_alert = max_prob > 0.75  # Lower threshold for edge
        
        return {
            'requires_alert': requires_alert,
            'confidence': max_prob.item(),
            'event_type': self._get_event_type(outputs),
            'timestamp': time.time()
        }

class SecurityDatasetGenerator:
    def __init__(self):
        # Use publicly available datasets and synthetic data
        self.base_datasets = [
            "UCF-Crime",  # Crime video dataset
            "Avenue Dataset",  # Abnormal event detection
            "ShanghaiTech Campus",  # Anomaly detection
            "CUHK Avenue"  # Surveillance videos
        ]
        
        self.synthetic_generator = SyntheticSecurityVideoGenerator()
    
    def create_training_dataset(self) -> List[Tuple[torch.Tensor, str, str]]:
        """Create comprehensive training dataset for security events"""
        dataset = []
        
        # Process each base dataset
        for base_dataset in self.base_datasets:
            videos, labels = self._load_base_dataset(base_dataset)
            
            for video, label in zip(videos, labels):
                # Extract 16-frame sequences
                sequences = self._extract_sequences(video, sequence_length=16)
                
                for sequence in sequences:
                    # Generate natural language description
                    description = self._generate_description(sequence, label)
                    dataset.append((sequence, label, description))
        
        # Add synthetic data to augment rare events
        synthetic_data = self.synthetic_generator.generate_rare_events(
            n_samples=5000,
            event_types=["unauthorized_access", "person_fall", "fire_smoke_detection"]
        )
        
        dataset.extend(synthetic_data)
        
        return dataset
    
    def _generate_description(self, video_sequence: torch.Tensor, event_label: str) -> str:
        """Generate natural language descriptions for training"""
        # Use rule-based generation for training data
        description_templates = {
            "suspicious_loitering": [
                "A person is standing still in the area for an extended period",
                "Individual appears to be waiting or loitering without clear purpose",
                "Person remains stationary in location longer than typical"
            ],
            "unattended_object": [
                "A bag or object has been left alone in the monitored area",
                "Unattended item detected without owner present",
                "Object placed and person moved away, item remains"
            ],
            # ... more templates
        }
        
        templates = description_templates.get(event_label, ["Security event detected"])
        return random.choice(templates)

class SyntheticSecurityVideoGenerator:
    """Generate synthetic security footage for rare events"""
    
    def __init__(self):
        self.human_pose_estimator = HumanPoseEstimator()
        self.scene_generator = SceneGenerator()
    
    def generate_rare_events(self, n_samples: int, event_types: List[str]) -> List:
        """Generate synthetic video data for rare security events"""
        synthetic_data = []
        
        for event_type in event_types:
            for _ in range(n_samples // len(event_types)):
                # Generate base scene
                scene = self.scene_generator.create_surveillance_scene()
                
                # Add human figures with appropriate poses/actions
                if event_type == "person_fall":
                    video = self._generate_fall_sequence(scene)
                elif event_type == "unauthorized_access":
                    video = self._generate_access_attempt(scene)
                elif event_type == "fire_smoke_detection":
                    video = self._generate_fire_scene(scene)
                
                description = f"Synthetic {event_type} event for training"
                synthetic_data.append((video, event_type, description))
        
        return synthetic_data

class HackerDemoScenarios:
    def __init__(self, invision_system):
        self.system = invision_system
        
    def demo_scenario_1_suspicious_loitering(self):
        """Demo: Person loitering near building entrance"""
        print("🎥 Playing Demo Video: Suspicious Loitering...")
        
        # Simulate real-time processing
        video_path = "demo_videos/loitering_scenario.mp4"
        
        # Process video and show real-time alerts
        for frame_batch in self._process_video_in_batches(video_path):
            result = self.system.analyze_frames(frame_batch)
            
            if result['alert_generated']:
                print(f"🚨 ALERT: {result['alert_text']}")
                print(f"   Confidence: {result['confidence']:.1%}")
                print(f"   Timestamp: {result['timestamp']}")
                print()
    
    def demo_scenario_2_unattended_bag(self):
        """Demo: Unattended bag detection"""
        print("🎥 Playing Demo Video: Unattended Object...")
        
        # Show how system detects object placement and abandonment
        alerts = [
            "Person placed bag near entrance and walked away",
            "Unattended bag detected in lobby area for 2 minutes",
            "Security alert: Suspicious package at main entrance, unattended for 5 minutes"
        ]
        
        for i, alert in enumerate(alerts):
            time.sleep(3)  # Simulate real-time progression
            print(f"🚨 ALERT {i+1}: {alert}")
    
    def demo_live_camera_feed(self):
        """Demo with actual camera feed from laptop"""
        print("🎥 Starting Live Camera Demo...")
        print("Walk around, sit down, stand up - watch the AI describe your actions!")
        
        cap = cv2.VideoCapture(0)  # Laptop camera
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Process frame
            analysis = self.system.analyze_single_frame(frame)
            
            # Display frame with annotations
            annotated_frame = self._annotate_frame(frame, analysis)
            cv2.imshow('Invision-Transformer Live Demo', annotated_frame)
            
            # Print analysis
            if analysis['description']:
                print(f"AI Description: {analysis['description']}")
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        cap.release()
        cv2.destroyAllWindows()