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 }