combined_visualization.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
from image_processor import ImageProcessor
from heatmap_generator import HeatmapGenerator

class CombinedVisualizer:
    def __init__(self):
        """
        Initialize the combined visualizer for creating overlaid threat visualizations
        """
        self.image_processor = ImageProcessor()
        self.heatmap_generator = HeatmapGenerator()
    
    def create_combined_visualization(self, image_pair_results, output_path, alpha_diff=0.4, alpha_low=0.3, alpha_medium=0.4, dpi=300):
        """
        Create a combined visualization that overlays difference image, low and medium threat heatmaps,
        and bounding boxes on top of each other.
        
        Args:
            image_pair_results: Dictionary containing all processing results
            output_path: Path to save the visualization
            alpha_diff: Transparency for difference image overlay
            alpha_low: Transparency for low threat heatmap overlay
            alpha_medium: Transparency for medium threat heatmap overlay
            dpi: Resolution for saved image
            
        Returns:
            Path to the generated visualization
        """
        # Extract required components from results
        original_image = image_pair_results['original_image']
        difference_image = image_pair_results['difference_image']
        bounding_boxes = image_pair_results['bounding_boxes']
        multi_heatmaps = image_pair_results.get('multi_heatmaps', {})
        labeled_regions = [r for r in image_pair_results.get('labeled_regions', []) 
                          if 'bbox' in r and 'threat_level' in r]
        
        # If labeled_regions not provided, extract from bounding boxes
        if not labeled_regions and 'threat_summary' in image_pair_results:
            # Create simplified labeled regions from bounding boxes
            for bbox in bounding_boxes:
                # Default to medium threat if specific threat level not available
                labeled_regions.append({
                    'bbox': bbox,
                    'threat_level': 'medium',
                    'difference_percentage': 50  # Default value
                })
        
        # Start with a copy of the original image
        combined_image = original_image.copy()
        
        # 1. Overlay the difference image with transparency
        # Convert difference image to RGB if it's grayscale
        if len(difference_image.shape) == 2 or difference_image.shape[2] == 1:
            diff_colored = cv2.applyColorMap(difference_image, cv2.COLORMAP_HOT)
            diff_colored = cv2.cvtColor(diff_colored, cv2.COLOR_BGR2RGB)
        else:
            diff_colored = difference_image
        
        # Overlay difference image
        combined_image = self.image_processor.overlay_images(combined_image, diff_colored, alpha_diff)
        
        # 2. Overlay low threat heatmap if available
        if 'low' in multi_heatmaps:
            low_heatmap = multi_heatmaps['low']
            combined_image = self.image_processor.overlay_images(combined_image, low_heatmap, alpha_low)
        
        # 3. Overlay medium threat heatmap if available
        if 'medium' in multi_heatmaps:
            medium_heatmap = multi_heatmaps['medium']
            combined_image = self.image_processor.overlay_images(combined_image, medium_heatmap, alpha_medium)
        
        # 4. Draw bounding boxes with threat level colors
        threat_colors = {
            'low': (0, 255, 0),      # Green
            'medium': (0, 165, 255), # Orange
            'high': (0, 0, 255)      # Red
        }
        
        # Draw bounding boxes based on threat levels
        for region in labeled_regions:
            bbox = region['bbox']
            threat_level = region['threat_level']
            x, y, w, h = bbox
            
            # Get color for this threat level (default to red if not found)
            color = threat_colors.get(threat_level, (0, 0, 255))
            
            # Convert BGR to RGB for matplotlib
            color_rgb = (color[2]/255, color[1]/255, color[0]/255)
            
            # Draw rectangle with threat level color
            cv2.rectangle(combined_image, (x, y), (x + w, y + h), color, 2)
            
            # Add label text with threat level
            if 'difference_percentage' in region:
                label_text = f"{threat_level.upper()}: {region['difference_percentage']:.1f}%"
            else:
                label_text = f"{threat_level.upper()}"
                
            cv2.putText(combined_image, label_text, (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        # Save the combined visualization
        plt.figure(figsize=(12, 8))
        plt.imshow(combined_image)
        plt.title('Combined Threat Visualization')
        plt.axis('off')
        plt.tight_layout()
        plt.savefig(output_path, dpi=dpi, bbox_inches='tight')
        plt.close()
        
        # Also save the raw image for potential further processing
        raw_output_path = output_path.replace('.png', '_raw.png')
        self.image_processor.save_image(combined_image, raw_output_path)
        
        return output_path
    
    def create_combined_visualization_from_files(self, original_path, difference_path, 
                                               low_heatmap_path, medium_heatmap_path,
                                               bounding_boxes, output_path, 
                                               alpha_diff=0.4, alpha_low=0.3, alpha_medium=0.4):
        """
        Create a combined visualization from individual image files
        
        Args:
            original_path: Path to original image
            difference_path: Path to difference image
            low_heatmap_path: Path to low threat heatmap
            medium_heatmap_path: Path to medium threat heatmap
            bounding_boxes: List of bounding boxes (x, y, w, h)
            output_path: Path to save the visualization
            alpha_diff: Transparency for difference image overlay
            alpha_low: Transparency for low threat heatmap overlay
            alpha_medium: Transparency for medium threat heatmap overlay
            
        Returns:
            Path to the generated visualization
        """
        # Load images
        original_image = self.image_processor.load_image(original_path)
        difference_image = self.image_processor.load_image(difference_path)
        
        # Load heatmaps if paths are provided
        low_heatmap = None
        medium_heatmap = None
        
        if low_heatmap_path:
            low_heatmap = self.image_processor.load_image(low_heatmap_path)
            
        if medium_heatmap_path:
            medium_heatmap = self.image_processor.load_image(medium_heatmap_path)
        
        # Create a mock image_pair_results dictionary
        image_pair_results = {
            'original_image': original_image,
            'difference_image': difference_image,
            'bounding_boxes': bounding_boxes,
            'multi_heatmaps': {}
        }
        
        if low_heatmap is not None:
            image_pair_results['multi_heatmaps']['low'] = low_heatmap
            
        if medium_heatmap is not None:
            image_pair_results['multi_heatmaps']['medium'] = medium_heatmap
        
        # Call the main visualization method
        return self.create_combined_visualization(
            image_pair_results, output_path, alpha_diff, alpha_low, alpha_medium
        )

# Example usage
if __name__ == "__main__":
    import os
    from deepfake_detector import DeepfakeDetector
    from labeling import ThreatLabeler
    
    # Initialize components
    detector = DeepfakeDetector()
    labeler = ThreatLabeler()
    heatmap_gen = HeatmapGenerator()
    img_processor = ImageProcessor()
    visualizer = CombinedVisualizer()
    
    # Example paths
    image1_path = "path/to/original.jpg"
    image2_path = "path/to/modified.jpg"
    output_dir = "path/to/output"
    
    # Ensure output directory exists
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Process images
    results = detector.process_image_pair(image1_path, image2_path)
    
    # Label regions
    original_image = img_processor.load_image(image1_path)
    labeled_image, labeled_regions = labeler.label_regions(
        original_image, results['difference_image'], results['bounding_boxes']
    )
    
    # Generate multi-level heatmaps
    multi_heatmaps = heatmap_gen.generate_multi_level_heatmap(original_image, labeled_regions)
    
    # Prepare results for combined visualization
    image_pair_results = {
        'original_image': original_image,
        'difference_image': results['difference_image'],
        'bounding_boxes': results['bounding_boxes'],
        'multi_heatmaps': multi_heatmaps,
        'labeled_regions': labeled_regions
    }
    
    # Create combined visualization
    output_path = os.path.join(output_dir, "combined_visualization.png")
    visualizer.create_combined_visualization(image_pair_results, output_path)
    
    print(f"Combined visualization saved to: {output_path}")