heatmap_generator.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from image_processor import ImageProcessor

class HeatmapGenerator:
    def __init__(self):
        """
        Initialize the heatmap generator for visualizing threat areas
        """
        self.image_processor = ImageProcessor()
        
        # Define colormap options
        self.colormap_options = {
            'hot': cv2.COLORMAP_HOT,       # Red-yellow-white, good for high intensity
            'jet': cv2.COLORMAP_JET,       # Blue-cyan-yellow-red, good for range
            'inferno': cv2.COLORMAP_INFERNO,  # Purple-red-yellow, good for threat
            'plasma': cv2.COLORMAP_PLASMA     # Purple-red-yellow, alternative
        }
        
        # Default colormap
        self.default_colormap = 'inferno'
    
    def generate_heatmap_from_diff(self, diff_image, threshold=0, blur_size=15):
        """
        Generate a heatmap directly from a difference image
        
        Args:
            diff_image: Difference image (0-255 range)
            threshold: Minimum difference value to consider (0-255)
            blur_size: Size of Gaussian blur kernel for smoothing
            
        Returns:
            Heatmap image
        """
        # Apply threshold to filter out low differences
        _, thresholded = cv2.threshold(diff_image, threshold, 255, cv2.THRESH_TOZERO)
        
        # Apply Gaussian blur to smooth the heatmap
        if blur_size > 0:
            blurred = cv2.GaussianBlur(thresholded, (blur_size, blur_size), 0)
        else:
            blurred = thresholded
        
        # Apply colormap
        heatmap = cv2.applyColorMap(blurred, self.colormap_options[self.default_colormap])
        
        # Convert to RGB for consistent display
        heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
        
        return heatmap_rgb
    
    def generate_heatmap_from_regions(self, image_shape, labeled_regions, sigma=40):
        """
        Generate a heatmap from labeled regions based on threat levels
        
        Args:
            image_shape: Shape of the original image (height, width)
            labeled_regions: List of regions with threat levels from ThreatLabeler
            sigma: Standard deviation for Gaussian kernel
            
        Returns:
            Heatmap image
        """
        # Create an empty heatmap
        height, width = image_shape[:2]
        heatmap = np.zeros((height, width), dtype=np.float32)
        
        # Define threat level weights with increased intensity
        threat_weights = {
            'low': 0.4,
            'medium': 0.7,
            'high': 1.0
        }
        
        # Add each region to the heatmap with appropriate weight
        for region in labeled_regions:
            bbox = region['bbox']
            threat_level = region['threat_level']
            diff_percentage = region['difference_percentage']
            
            # Calculate center of bounding box
            x, y, w, h = bbox
            center_x, center_y = x + w // 2, y + h // 2
            
            # Calculate intensity based on threat level and difference percentage with increased brightness
            intensity = threat_weights[threat_level] * (diff_percentage / 100) * 1.2
            
            # Create a Gaussian kernel centered at the region with increased sigma for more circular spread
            y_coords, x_coords = np.ogrid[:height, :width]
            dist_from_center = ((y_coords - center_y) ** 2 + (x_coords - center_x) ** 2) / (2 * sigma ** 2)
            kernel = np.exp(-dist_from_center) * intensity
            
            # Add to heatmap
            heatmap += kernel
        
        # Normalize heatmap to 0-255 range
        if np.max(heatmap) > 0:  # Avoid division by zero
            heatmap = (heatmap / np.max(heatmap) * 255).astype(np.uint8)
        else:
            heatmap = np.zeros((height, width), dtype=np.uint8)
        
        # Apply colormap
        colored_heatmap = cv2.applyColorMap(heatmap, self.colormap_options[self.default_colormap])
        colored_heatmap = cv2.cvtColor(colored_heatmap, cv2.COLOR_BGR2RGB)
        
        return colored_heatmap
    
    def overlay_heatmap(self, original_image, heatmap, alpha=0.6):
        """
        Overlay heatmap on original image
        
        Args:
            original_image: Original image
            heatmap: Heatmap image
            alpha: Transparency factor (0-1)
            
        Returns:
            Overlaid image
        """
        # Ensure images are the same size
        if original_image.shape[:2] != heatmap.shape[:2]:
            heatmap = cv2.resize(heatmap, (original_image.shape[1], original_image.shape[0]))
        
        # Overlay heatmap on original image
        return self.image_processor.overlay_images(original_image, heatmap, alpha)
    
    def generate_threat_heatmap(self, image, labeled_regions, overlay=True, alpha=0.6):
        """
        Generate a complete threat heatmap visualization
        
        Args:
            image: Original image
            labeled_regions: List of regions with threat levels
            overlay: Whether to overlay on original image
            alpha: Transparency for overlay
            
        Returns:
            Heatmap image or overlaid image
        """
        # Generate heatmap from labeled regions
        heatmap = self.generate_heatmap_from_regions(image.shape, labeled_regions)
        
        # Overlay on original image if requested
        if overlay:
            return self.overlay_heatmap(image, heatmap, alpha)
        else:
            return heatmap
    
    def save_heatmap_visualization(self, image, heatmap, output_path, dpi=300):
        """
        Save a side-by-side visualization of original image and heatmap
        
        Args:
            image: Original image
            heatmap: Heatmap image
            output_path: Path to save visualization
            dpi: Resolution for saved image
        """
        # Create figure with two subplots
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
        
        # Display original image
        ax1.imshow(image)
        ax1.set_title('Original Image')
        ax1.axis('off')
        
        # Display heatmap
        ax2.imshow(heatmap)
        ax2.set_title('Threat Heatmap')
        ax2.axis('off')
        
        # Save figure
        plt.tight_layout()
        plt.savefig(output_path, dpi=dpi, bbox_inches='tight')
        plt.close()
        
    def generate_multi_level_heatmap(self, image, labeled_regions):
        """
        Generate separate heatmaps for each threat level
        
        Args:
            image: Original image
            labeled_regions: List of regions with threat levels
            
        Returns:
            Dictionary with heatmaps for each threat level and combined
        """
        # Create separate lists for each threat level
        low_regions = [r for r in labeled_regions if r['threat_level'] == 'low']
        medium_regions = [r for r in labeled_regions if r['threat_level'] == 'medium']
        high_regions = [r for r in labeled_regions if r['threat_level'] == 'high']
        
        # Generate heatmaps for each level
        low_heatmap = self.generate_heatmap_from_regions(image.shape, low_regions)
        medium_heatmap = self.generate_heatmap_from_regions(image.shape, medium_regions)
        high_heatmap = self.generate_heatmap_from_regions(image.shape, high_regions)
        
        # Generate combined heatmap
        combined_heatmap = self.generate_heatmap_from_regions(image.shape, labeled_regions)
        
        # Overlay all on original image
        combined_overlay = self.overlay_heatmap(image, combined_heatmap)
        
        return {
            'low': low_heatmap,
            'medium': medium_heatmap,
            'high': high_heatmap,
            'combined': combined_heatmap,
            'overlay': combined_overlay
        }