src/inference.py

"""
src/inference.py
Unified inference interface for FirenetCNN
Supports single-image, video, and webcam inference with Grad-CAM visualization
"""
import cv2
import numpy as np
import tensorflow as tf
from pathlib import Path
from typing import Optional, Tuple, Dict, Any
from PIL import Image
import time

from .model import FireNetModel
from .gradcam import GradCAM
class FireNetInference:
    """
    Unified inference engine for FirenetCNN
    
    Provides interface for:
    - Single image inference with optional Grad-CAM
    - Video inference with frame-by-frame processing
    - Real-time webcam inference
    - Model evaluation and reporting
    """
    
    def __init__(self, model_path: str = 'models/FirenetCNN.keras'):
        """
        Initialize the inference engine.
        
        Args:
            model_path: Path to the trained model file
        """
        self.model_path = model_path
        self.class_labels = FireNetModel.CLASS_LABELS
        self.reverse_class_map = FireNetModel.REVERSE_CLASS_MAP
        self.color_map = FireNetModel.TEXT_COLOR
        
        # Load model and Grad-CAM components
        self.model_wrapper = FireNetModel(model_path)
        self.model = self.model_wrapper.load_pretrained_model()
        self.gradcam = GradCAM(self.model, self.model_wrapper.last_conv_layer_name)
        
        # Create gradient model for Grad-CAM
        self.grad_model = tf.keras.Model(
            [self.model.inputs],
            [self.model.get_layer(self.model_wrapper.last_conv_layer_name).output, self.model.output]
        )
    
    def predict_image(self, image_path: str, 
                     apply_gradcam: bool = True, 
                     output_path: Optional[str] = None) -> Dict[str, Any]:
        """
        Perform inference on a single image.
        
        Args:
            image_path: Path to input image
            apply_gradcam: Whether to generate Grad-CAM heatmap
            output_path: Optional path to save annotated image
            
        Returns:
            Dictionary with prediction results and annotations
        """
        # Load image
        image = cv2.imread(str(image_path))
        if image is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        # Store original for potential output
        original_image = image.copy()
        height, width = image.shape[:2]
        
        # Preprocess for model
        input_frame = GradCAM.preprocess_frame(image)
        
        # Get predictions
        predictions = self.model.predict(input_frame, verbose=0)
        label, confidence, prob_array = GradCAM.get_confidence_and_prediction(
            predictions, self.class_labels
        )
        
        # Generate Grad-CAM if requested and prediction is fire or smoke
        heatmap = None
        superimposed_img = None
        
        if apply_gradcam and label in ['fire', 'smoke']:
            # Convert BGR to RGB for Grad-CAM processing
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # Prepare input for Grad-CAM
            img_array = tf.expand_dims(tf.convert_to_tensor(rgb_image / 255.0), 0)
            class_index = self.reverse_class_map[label]
            
            # Generate heatmap
            heatmap = self.gradcam.generate_heatmap(img_array, class_index)
            
            # Overlay heatmap
            superimposed_img = self.gradcam.overlay_heatmap(
                original_image, heatmap, alpha=0.5
            )
        
        # Create text overlay parameters
        text_color, bg_color = GradCAM.get_text_overlay_params(label, self.color_map)
        result_text = f"Class: {label} ({confidence*100:.2f}%)"
        
        # Add text overlay to image
        annotated_image = GradCAM.create_text_overlay(
            superimposed_img if superimposed_img is not None else original_image,
            result_text, text_color, bg_color
        )
        
        # Save annotated image if output path provided
        if output_path:
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            cv2.imwrite(str(output_path), annotated_image)
        
        # Return results as dictionary
        result = {
            'image_path': image_path,
            'label': label,
            'confidence': confidence,
            'probability_array': prob_array.tolist(),
            'class_labels': self.class_labels,
            'has_gradcam': heatmap is not None,
            'heatmap': heatmap.tolist() if heatmap is not None else None,
            'annotated_image': annotated_image if output_path is None else None,
            'processing_time': 0  # Will be measured externally if needed
        }
        
        return result
    
    def predict_video(self, video_path: str, 
                     output_path: Optional[str] = None,
                     skip_frames: int = 5,
                     apply_gradcam: bool = True) -> Dict[str, Any]:
        """
        Process a video file frame-by-frame.
        
        Args:
            video_path: Path to input video file
            output_path: Optional path to save processed video
            skip_frames: Process every Nth frame (for performance)
            apply_gradcam: Whether to generate Grad-CAM for fire/smoke frames
            
        Returns:
            Dictionary with video processing statistics
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Could not open video file: {video_path}")
        
        # Get video properties
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # Setup video writer if output path provided
        video_writer = None
        if output_path:
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        # Statistics
        stats = {
            'total_frames': frame_count,
            'processed_frames': 0,
            'predictions': [],
            'frame_by_frame': [],
            'final_frame': None,
            'processing_time_seconds': 0
        }
        
        # State tracking across frames
        last_label = "Initializing..."
        last_confidence = 0
        last_heatmap = None
        last_color = (255, 255, 255)  # White
        
        start_time = time.time()
        frame_number = 0
        
        while cap.isOpened() and frame_number < frame_count:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Process only every Nth frame for performance
            if frame_number % skip_frames == 0 or frame_number == 0:
                # Preprocess frame
                input_frame = GradCAM.preprocess_frame(frame)
                
                # Get predictions
                predictions = self.model.predict(input_frame, verbose=0)
                label, confidence, prob_array = GradCAM.get_confidence_and_prediction(
                    predictions, self.class_labels
                )
                
                # Update state
                last_label = label
                last_confidence = confidence
                
                # Generate Grad-CAM if requested and prediction is fire or smoke
                if apply_gradcam and label in ['fire', 'smoke']:
                    rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    img_array = tf.expand_dims(tf.convert_to_tensor(rgb_image / 255.0), 0)
                    class_index = self.reverse_class_map[label]
                    
                    last_heatmap = self.gradcam.generate_heatmap(img_array, class_index)
                    last_color = self.color_map.get(label, (255, 255, 255))
                else:
                    last_heatmap = None
                    last_color = self.color_map.get(label, (255, 255, 255))
                
                # Update statistics
                stats['processed_frames'] += 1
            
            # Apply Grad-CAM if available
            display_frame = frame.copy()
            if last_heatmap is not None:
                display_frame = self.gradcam.overlay_heatmap(frame, last_heatmap, alpha=0.5)
            
            # Create text overlay
            text_color, bg_color = GradCAM.get_text_overlay_params(last_label, self.color_map)
            result_text = f"Class: {last_label} ({last_confidence*100:.2f}%)"
            display_frame = GradCAM.create_text_overlay(display_frame, result_text, text_color, bg_color)
            
            # Save frame if video writer provided
            if video_writer:
                video_writer.write(display_frame)
            
            # Store frame info
            frame_info = {
                'frame_number': frame_number,
                'label': last_label,
                'confidence': last_confidence,
                'processed': (frame_number % skip_frames == 0 or frame_number == 0),
                'has_gradcam': last_heatmap is not None,
                'color_bgr': last_color
            }
            stats['frame_by_frame'].append(frame_info)
            
            # Store last frame for final output
            stats['final_frame'] = display_frame.copy()
            
            frame_number += 1
        
        # Cleanup
        cap.release()
        if video_writer:
            video_writer.release()
        
        # Calculate processing time
        processing_time = time.time() - start_time
        stats['processing_time_seconds'] = processing_time
        
        return stats
    
    def predict_webcam(self, window_name: str = "Live Webcam Inference (Slow)",
                      apply_gradcam: bool = False,
                      max_frames: Optional[int] = None) -> Dict[str, Any]:
        """
        Perform real-time inference using webcam.
        
        Args:
            window_name: Name of the display window
            apply_gradcam: Whether to generate Grad-CAM (normally False for webcam)
            max_frames: Maximum number of frames to process (None for infinite)
            
        Returns:
            Dictionary with webcam inference results
        """
        cap = cv2.VideoCapture(0)
        if not cap.isOpened():
            raise RuntimeError("Could not open webcam.")
        
        # Statistics
        stats = {
            'frames_processed': 0,
            'predictions': [],
            'detected_classes': set(),
            'processing_times': [],
            'live_feed_active': True
        }
        
        frame_count = 0
        
        print(f"Webcam started. Press 'q' to quit.")
        
        try:
            while cap.isOpened() and (max_frames is None or frame_count < max_frames):
                ret, frame = cap.read()
                if not ret:
                    print("Error: Failed to capture frame.")
                    break
                
                # Start timing
                start_time = time.time()
                
                # Preprocess frame
                input_frame = GradCAM.preprocess_frame(frame)
                
                # Get predictions
                predictions = self.model.predict(input_frame, verbose=0)
                label, confidence, prob_array = GradCAM.get_confidence_and_prediction(
                    predictions, self.class_labels
                )
                
                # End timing
                processing_time = time.time() - start_time
                stats['processing_times'].append(processing_time)
                
                # Update stats
                stats['frames_processed'] += 1
                stats['detected_classes'].add(label)
                frame_info = {
                    'frame_number': frame_count,
                    'label': label,
                    'confidence': confidence,
                    'processing_time': processing_time
                }
                stats['predictions'].append(frame_info)
                
                # Generate Grad-CAM if requested and prediction is fire or smoke
                heatmap = None
                if apply_gradcam and label in ['fire', 'smoke']:
                    rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    img_array = tf.expand_dims(tf.convert_to_tensor(rgb_image / 255.0), 0)
                    class_index = self.reverse_class_map[label]
                    
                    heatmap = self.gradcam.generate_heatmap(img_array, class_index)
                
                # Create text overlay
                text_color, bg_color = GradCAM.get_text_overlay_params(label, self.color_map)
                result_text = f"Class: {label} ({confidence*100:.2f}%)"
                display_frame = GradCAM.create_text_overlay(frame, result_text, text_color, bg_color)
                
                # Apply Grad-CAM if available
                if heatmap is not None:
                    display_frame = self.gradcam.overlay_heatmap(display_frame, heatmap, alpha=0.5)
                
                # Display the frame
                cv2.imshow(window_name, display_frame)
                
                # Check for quit
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
                
                frame_count += 1
        
        finally:
            # Cleanup
            cap.release()
            cv2.destroyAllWindows()
            stats['live_feed_active'] = False
        
        return stats
    
    def create_gradcam_demo_image(self, image_path: str, class_to_highlight: Optional[str] = None) -> Dict[str, Any]:
        """
        Create a demonstration image with Grad-CAM for all classes.
        
        Args:
            image_path: Path to input image
            class_to_highlight: Optional specific class to highlight (default: auto-select best)
            
        Returns:
            Dictionary with demo images for each class
        """
        # Load and preprocess image
        image = cv2.imread(str(image_path))
        if image is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        img_array = tf.expand_dims(tf.convert_to_tensor(rgb_image / 255.0), 0)
        
        # Get predictions
        predictions = self.model.predict(img_array, verbose=0)
        prob_array = predictions[0]
        
        # Determine which class to highlight (default: class with highest probability)
        if class_to_highlight is None:
            max_idx = np.argmax(prob_array)
            class_to_highlight = self.class_labels[max_idx]
        
        # Create demo images for each class
        demo_images = {}
        
        for class_label in self.class_labels:
            # Get class index
            class_index = self.reverse_class_map[class_label]
            
            # Generate heatmap
            heatmap = self.gradcam.generate_heatmap(img_array, class_index)
            
            # Overlay heatmap
            if class_label == class_to_highlight:
                # For highlighted class, show the original heatmap
                overlaid = self.gradcam.overlay_heatmap(
                    rgb_image, heatmap, alpha=0.6
                )
                overlaid = cv2.cvtColor(overlaid, cv2.COLOR_RGB2BGR)
            else:
                # For other classes, show grayscale heatmap
                heatmap_bgr = cv2.applyColorMap(
                    np.uint8(255 * heatmap), cv2.COLORMAP_JET
                )
                overlaid = heatmap_bgr
            
            # Create text overlay
            text_color, bg_color = GradCAM.get_text_overlay_params(class_label, self.color_map)
            result_text = f"{class_label} ({prob_array[class_index]*100:.1f}%)"
            demo_images[class_label] = GradCAM.create_text_overlay(
                overlaid, result_text, text_color, bg_color
            )
        
        return {
            'original_image': cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR),
            'prediction': {
                'predicted_class': self.class_labels[np.argmax(prob_array)],
                'confidence': float(np.max(prob_array)),
                'probabilities': prob_array.tolist()
            },
            'demo_images': demo_images,
            'class_to_highlight': class_to_highlight
        }
    
    @classmethod
    def convert_model_format(cls, input_path: str, output_path: str) -> None:
        """
        Convert model from legacy format to modern Keras format.
        
        Args:
            input_path: Path to input model (HDF5 format)
            output_path: Path to save converted model
        """
        from tensorflow.keras.models import load_model
        
        model = load_model(input_path, compile=False)
        model.save(output_path, save_format='keras')
        print(f"Model converted and saved to: {output_path}")
    
    @classmethod
    def evaluate_model_on_dataset(cls, model_path: str, 
                                 test_dir: str,
                                 output_report: Optional[str] = None) -> Dict[str, Any]:
        """
        Evaluate model performance on a dataset.
        
        Args:
            model_path: Path to model file
            test_dir: Path to test dataset directory (with class subdirectories)
            output_report: Optional path to save evaluation report
            
        Returns:
            Dictionary with evaluation results
        """
        from sklearn.metrics import classification_report, confusion_matrix
        import matplotlib.pyplot as plt
        import seaborn as sns
        
        # Load model
        model_wrapper = FireNetModel(model_path)
        model = model_wrapper.load_pretrained_model()
        
        # Create data generators
        test_datagen = ImageDataGenerator(rescale=1./255.)
        test_generator = test_datagen.flow_from_directory(
            test_dir,
            target_size=FireNetModel.IMAGE_SIZE,
            batch_size=FireNetModel.BATCH_SIZE,
            class_mode='categorical',
            shuffle=False
        )
        
        # Make predictions
        y_pred = model.predict(test_generator, verbose=0)
        y_pred_classes = np.argmax(y_pred, axis=1)
        
        # Get true labels
        y_true = test_generator.classes
        
        # Generate classification report
        report = classification_report(
            y_true, 
            y_pred_classes, 
            target_names=FireNetModel.CLASS_LABELS,
            output_dict=True
        )
        
        # Generate confusion matrix
        cm = confusion_matrix(y_true, y_pred_classes)
        
        # Create and save visualization
        plt.figure(figsize=(10, 8))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                   xticklabels=FireNetModel.CLASS_LABELS,
                   yticklabels=FireNetModel.CLASS_LABELS)
        plt.title('Confusion Matrix')
        plt.ylabel('True Label')
        plt.xlabel('Predicted Label')
        
        if output_report:
            plt.savefig(output_report, dpi=150, bbox_inches='tight')
            plt.close()
        
        # Save report as text file if requested
        if output_report and output_report.endswith('.txt'):
            with open(output_report, 'w') as f:
                f.write("Forest Fire Detection Model Evaluation Report\n")
                f.write("=" * 60 + "\n\n")
                f.write(f"Total samples: {len(y_true)}\n")
                f.write(f"Classes: {', '.join(FireNetModel.CLASS_LABELS)}\n\n")
                f.write("Classification Report:\n")
                f.write(classification_report(y_true, y_pred_classes, 
                                             target_names=FireNetModel.CLASS_LABELS) + "\n\n")
                f.write("Confusion Matrix:\n")
                f.write(str(cm) + "\n")
                
                # Add detailed metrics
                f.write("\nDetailed Metrics:\n")
                for i, class_name in enumerate(FireNetModel.CLASS_LABELS):
                    precision = report[class_name]['precision']
                    recall = report[class_name]['recall']
                    f1 = report[class_name]['f1-score']
                    support = report[class_name]['support']
                    f.write(f"{class_name:10} - Precision: {precision:.3f}, "
                           f"Recall: {recall:.3f}, F1-score: {f1:.3f}, "
                           f"Support: {support}\n")
        
        return {
            'classification_report': report,
            'confusion_matrix': cm.tolist(),
            'total_samples': len(y_true),
            'accuracy': report['accuracy'],
            'macro_avg_f1': report['macro avg']['f1-score'],
            'weighted_avg_f1': report['weighted avg']['f1-score']
        }