Upload heatmap_generator.py

heatmap_generator.py

@@ -0,0 +1,212 @@
+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
+        }