ParallelLLC/Segmentation

1"""2Visualization Utilities3 4This module provides comprehensive visualization tools for segmentation results,5attention maps, and experiment comparisons in few-shot and zero-shot learning.6"""7 8import torch9import numpy as np10import matplotlib.pyplot as plt11import matplotlib.patches as patches12from matplotlib.colors import ListedColormap13import seaborn as sns14from typing import Dict, List, Tuple, Optional, Union15import cv216from PIL import Image17import os18 19 20class SegmentationVisualizer:21    """Visualization tools for segmentation results."""22    23    def __init__(self, figsize: Tuple[int, int] = (15, 10)):24        self.figsize = figsize25        26        # Color maps for different classes27        self.class_colors = {28            'building': [1.0, 0.0, 0.0],      # Red29            'road': [0.0, 1.0, 0.0],          # Green30            'vegetation': [0.0, 0.0, 1.0],    # Blue31            'water': [1.0, 1.0, 0.0],         # Yellow32            'shirt': [1.0, 0.5, 0.0],         # Orange33            'pants': [0.5, 0.0, 1.0],         # Purple34            'dress': [0.0, 1.0, 1.0],         # Cyan35            'shoes': [1.0, 0.0, 1.0],         # Magenta36            'robot': [0.5, 0.5, 0.5],         # Gray37            'tool': [0.8, 0.4, 0.2],          # Brown38            'safety': [0.2, 0.8, 0.2]         # Light Green39        }40    41    def visualize_segmentation(42        self, 43        image: torch.Tensor, 44        predictions: Dict[str, torch.Tensor], 45        ground_truth: Optional[Dict[str, torch.Tensor]] = None,46        title: str = "Segmentation Results"47    ) -> plt.Figure:48        """Visualize segmentation results with optional ground truth comparison."""49        num_classes = len(predictions)50        has_gt = ground_truth is not None51        52        # Calculate subplot layout53        if has_gt:54            cols = 355            rows = max(2, num_classes)56        else:57            cols = 258            rows = max(1, num_classes)59        60        fig, axes = plt.subplots(rows, cols, figsize=(cols * 5, rows * 4))61        if rows == 1:62            axes = axes.reshape(1, -1)63        64        # Original image65        image_np = image.permute(1, 2, 0).cpu().numpy()66        # Denormalize if needed67        if image_np.min() < 0 or image_np.max() > 1:68            image_np = (image_np - image_np.min()) / (image_np.max() - image_np.min())69        70        axes[0, 0].imshow(image_np)71        axes[0, 0].set_title("Original Image")72        axes[0, 0].axis('off')73        74        # Combined prediction overlay75        if cols > 1:76            combined_pred = self.create_combined_mask(predictions)77            axes[0, 1].imshow(image_np)78            axes[0, 1].imshow(combined_pred, alpha=0.6, cmap='tab10')79            axes[0, 1].set_title("Combined Predictions")80            axes[0, 1].axis('off')81        82        # Ground truth overlay83        if has_gt and cols > 2:84            combined_gt = self.create_combined_mask(ground_truth)85            axes[0, 2].imshow(image_np)86            axes[0, 2].imshow(combined_gt, alpha=0.6, cmap='tab10')87            axes[0, 2].set_title("Ground Truth")88            axes[0, 2].axis('off')89        90        # Individual class predictions91        for i, (class_name, pred_mask) in enumerate(predictions.items()):92            row = i + 1 if has_gt else i93            col_offset = 094            95            # Prediction mask96            pred_np = pred_mask.cpu().numpy()97            axes[row, col_offset].imshow(pred_np, cmap='gray')98            axes[row, col_offset].set_title(f"Prediction: {class_name}")99            axes[row, col_offset].axis('off')100            101            # Overlay on original image102            col_offset += 1103            axes[row, col_offset].imshow(image_np)104            axes[row, col_offset].imshow(pred_np, alpha=0.6, cmap='Reds')105            axes[row, col_offset].set_title(f"Overlay: {class_name}")106            axes[row, col_offset].axis('off')107            108            # Ground truth comparison109            if has_gt and class_name in ground_truth:110                col_offset += 1111                gt_mask = ground_truth[class_name]112                gt_np = gt_mask.cpu().numpy()113                114                # Create comparison visualization115                comparison = np.zeros((*gt_np.shape, 3))116                comparison[gt_np > 0.5] = [0, 1, 0]  # Green for ground truth117                comparison[pred_np > 0.5] = [1, 0, 0]  # Red for prediction118                comparison[(gt_np > 0.5) & (pred_np > 0.5)] = [1, 1, 0]  # Yellow for overlap119                120                axes[row, col_offset].imshow(image_np)121                axes[row, col_offset].imshow(comparison, alpha=0.6)122                axes[row, col_offset].set_title(f"Comparison: {class_name}")123                axes[row, col_offset].axis('off')124        125        plt.tight_layout()126        return fig127    128    def create_combined_mask(self, masks: Dict[str, torch.Tensor]) -> np.ndarray:129        """Create a combined mask visualization for multiple classes."""130        if not masks:131            return np.zeros((512, 512))132        133        # Get the shape from the first mask134        first_mask = list(masks.values())[0]135        combined = np.zeros((*first_mask.shape, 3))136        137        for i, (class_name, mask) in enumerate(masks.items()):138            mask_np = mask.cpu().numpy()139            color = self.class_colors.get(class_name, [1, 1, 1])140            141            # Apply color to mask142            for c in range(3):143                combined[:, :, c] += mask_np * color[c]144        145        # Normalize146        combined = np.clip(combined, 0, 1)147        return combined148    149    def visualize_attention_maps(150        self, 151        image: torch.Tensor, 152        attention_maps: torch.Tensor, 153        class_names: List[str],154        title: str = "Attention Maps"155    ) -> plt.Figure:156        """Visualize attention maps for different classes."""157        num_classes = len(class_names)158        fig, axes = plt.subplots(2, num_classes, figsize=(num_classes * 4, 8))159        160        # Original image161        image_np = image.permute(1, 2, 0).cpu().numpy()162        if image_np.min() < 0 or image_np.max() > 1:163            image_np = (image_np - image_np.min()) / (image_np.max() - image_np.min())164        165        for i in range(num_classes):166            axes[0, i].imshow(image_np)167            axes[0, i].set_title(f"Original - {class_names[i]}")168            axes[0, i].axis('off')169        170        # Attention maps171        attention_np = attention_maps.cpu().numpy()172        for i in range(min(num_classes, attention_np.shape[0])):173            attention_map = attention_np[i]174            175            # Resize attention map to image size176            attention_map = cv2.resize(attention_map, (image_np.shape[1], image_np.shape[0]))177            178            axes[1, i].imshow(attention_map, cmap='hot')179            axes[1, i].set_title(f"Attention - {class_names[i]}")180            axes[1, i].axis('off')181        182        plt.tight_layout()183        return fig184    185    def visualize_prompt_points(186        self, 187        image: torch.Tensor, 188        prompts: List[Dict],189        title: str = "Prompt Points"190    ) -> plt.Figure:191        """Visualize prompt points and boxes on the image."""192        fig, ax = plt.subplots(1, 1, figsize=(10, 10))193        194        # Original image195        image_np = image.permute(1, 2, 0).cpu().numpy()196        if image_np.min() < 0 or image_np.max() > 1:197            image_np = (image_np - image_np.min()) / (image_np.max() - image_np.min())198        199        ax.imshow(image_np)200        201        # Plot prompts202        colors = plt.cm.Set3(np.linspace(0, 1, len(prompts)))203        204        for i, prompt in enumerate(prompts):205            color = colors[i]206            207            if prompt['type'] == 'point':208                x, y = prompt['data']209                ax.scatter(x, y, c=[color], s=100, marker='o', 210                          label=f"{prompt['class']} (point)")211                212            elif prompt['type'] == 'box':213                x1, y1, x2, y2 = prompt['data']214                rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, 215                                       linewidth=2, edgecolor=color, 216                                       facecolor='none', 217                                       label=f"{prompt['class']} (box)")218                ax.add_patch(rect)219        220        ax.set_title(title)221        ax.legend()222        ax.axis('off')223        224        return fig225 226 227class ExperimentVisualizer:228    """Visualization tools for experiment results and comparisons."""229    230    def __init__(self):231        self.segmentation_visualizer = SegmentationVisualizer()232    233    def plot_metrics_comparison(234        self, 235        results: Dict[str, List[float]], 236        metric_name: str = "IoU",237        title: str = "Metrics Comparison"238    ) -> plt.Figure:239        """Plot comparison of metrics across different methods/strategies."""240        fig, ax = plt.subplots(1, 1, figsize=(10, 6))241        242        # Prepare data243        methods = list(results.keys())244        values = [np.mean(results[method]) for method in methods]245        errors = [np.std(results[method]) for method in methods]246        247        # Create bar plot248        bars = ax.bar(methods, values, yerr=errors, capsize=5, alpha=0.7)249        250        # Add value labels on bars251        for bar, value in zip(bars, values):252            height = bar.get_height()253            ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,254                   f'{value:.3f}', ha='center', va='bottom')255        256        ax.set_title(title)257        ax.set_ylabel(metric_name)258        ax.set_xlabel("Methods")259        ax.grid(True, alpha=0.3)260        261        plt.xticks(rotation=45)262        plt.tight_layout()263        264        return fig265    266    def plot_learning_curves(267        self, 268        episode_metrics: List[Dict[str, float]], 269        metric_name: str = "iou"270    ) -> plt.Figure:271        """Plot learning curves over episodes."""272        fig, ax = plt.subplots(1, 1, figsize=(12, 6))273        274        # Extract metric values275        episodes = range(1, len(episode_metrics) + 1)276        values = [ep.get(metric_name, 0) for ep in episode_metrics]277        278        # Plot learning curve279        ax.plot(episodes, values, 'b-', linewidth=2, label=f'{metric_name.upper()}')280        281        # Add moving average282        window_size = min(10, len(values) // 4)283        if window_size > 1:284            moving_avg = np.convolve(values, np.ones(window_size)/window_size, mode='valid')285            ax.plot(episodes[window_size-1:], moving_avg, 'r--', linewidth=2, 286                   label=f'Moving Average (window={window_size})')287        288        ax.set_title(f"Learning Curve - {metric_name.upper()}")289        ax.set_xlabel("Episode")290        ax.set_ylabel(metric_name.upper())291        ax.grid(True, alpha=0.3)292        ax.legend()293        294        plt.tight_layout()295        return fig296    297    def plot_shot_analysis(298        self, 299        shot_results: Dict[int, List[float]], 300        metric_name: str = "iou"301    ) -> plt.Figure:302        """Plot performance analysis across different numbers of shots."""303        fig, ax = plt.subplots(1, 1, figsize=(10, 6))304        305        # Prepare data306        shots = sorted(shot_results.keys())307        means = [np.mean(shot_results[shot]) for shot in shots]308        stds = [np.std(shot_results[shot]) for shot in shots]309        310        # Create line plot with error bars311        ax.errorbar(shots, means, yerr=stds, marker='o', linewidth=2, 312                   capsize=5, capthick=2)313        314        ax.set_title(f"Performance vs Number of Shots - {metric_name.upper()}")315        ax.set_xlabel("Number of Shots")316        ax.set_ylabel(f"Mean {metric_name.upper()}")317        ax.grid(True, alpha=0.3)318        319        plt.tight_layout()320        return fig321    322    def plot_prompt_strategy_comparison(323        self, 324        strategy_results: Dict[str, Dict[str, float]], 325        metric_name: str = "mean_iou"326    ) -> plt.Figure:327        """Plot comparison of different prompt strategies."""328        fig, ax = plt.subplots(1, 1, figsize=(12, 6))329        330        # Prepare data331        strategies = list(strategy_results.keys())332        values = [strategy_results[s].get(metric_name, 0) for s in strategies]333        errors = [strategy_results[s].get(f'std_{metric_name.split("_")[-1]}', 0) 334                 for s in strategies]335        336        # Create bar plot337        bars = ax.bar(strategies, values, yerr=errors, capsize=5, alpha=0.7)338        339        # Add value labels340        for bar, value in zip(bars, values):341            height = bar.get_height()342            ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,343                   f'{value:.3f}', ha='center', va='bottom')344        345        ax.set_title(f"Prompt Strategy Comparison - {metric_name}")346        ax.set_ylabel(metric_name.replace('_', ' ').title())347        ax.set_xlabel("Strategy")348        ax.grid(True, alpha=0.3)349        350        plt.xticks(rotation=45)351        plt.tight_layout()352        353        return fig354    355    def create_comprehensive_report(356        self, 357        experiment_results: Dict,358        output_dir: str,359        experiment_name: str = "experiment"360    ):361        """Create a comprehensive visualization report."""362        os.makedirs(output_dir, exist_ok=True)363        364        # Create summary plots365        if 'episode_metrics' in experiment_results:366            # Learning curves367            for metric in ['iou', 'dice', 'precision', 'recall']:368                fig = self.plot_learning_curves(369                    experiment_results['episode_metrics'], 370                    metric371                )372                fig.savefig(os.path.join(output_dir, f'{experiment_name}_learning_curve_{metric}.png'))373                plt.close(fig)374        375        if 'class_metrics' in experiment_results:376            # Class-wise performance377            class_results = experiment_results['class_metrics']378            for class_name, metrics in class_results.items():379                if isinstance(metrics, list):380                    fig = self.plot_learning_curves(metrics, 'iou')381                    fig.savefig(os.path.join(output_dir, f'{experiment_name}_class_{class_name}.png'))382                    plt.close(fig)383        384        if 'shot_analysis' in experiment_results:385            # Shot analysis386            for metric in ['iou', 'dice']:387                fig = self.plot_shot_analysis(388                    experiment_results['shot_analysis'], 389                    metric390                )391                fig.savefig(os.path.join(output_dir, f'{experiment_name}_shot_analysis_{metric}.png'))392                plt.close(fig)393        394        if 'strategy_comparison' in experiment_results:395            # Strategy comparison396            for metric in ['mean_iou', 'mean_dice']:397                fig = self.plot_prompt_strategy_comparison(398                    experiment_results['strategy_comparison'], 399                    metric400                )401                fig.savefig(os.path.join(output_dir, f'{experiment_name}_strategy_comparison_{metric}.png'))402                plt.close(fig)403        404        print(f"Comprehensive report saved to {output_dir}")405 406 407class AttentionVisualizer:408    """Specialized visualizer for attention mechanisms."""409    410    def __init__(self):411        self.segmentation_visualizer = SegmentationVisualizer()412    413    def visualize_cross_attention(414        self, 415        image: torch.Tensor, 416        text_tokens: List[str], 417        attention_weights: torch.Tensor,418        title: str = "Cross-Attention Visualization"419    ) -> plt.Figure:420        """Visualize cross-attention between image and text tokens."""421        fig, axes = plt.subplots(2, 2, figsize=(15, 12))422        423        # Original image424        image_np = image.permute(1, 2, 0).cpu().numpy()425        if image_np.min() < 0 or image_np.max() > 1:426            image_np = (image_np - image_np.min()) / (image_np.max() - image_np.min())427        428        axes[0, 0].imshow(image_np)429        axes[0, 0].set_title("Original Image")430        axes[0, 0].axis('off')431        432        # Text tokens433        axes[0, 1].text(0.1, 0.5, '\n'.join(text_tokens), fontsize=12, 434                       verticalalignment='center')435        axes[0, 1].set_title("Text Tokens")436        axes[0, 1].axis('off')437        438        # Attention heatmap439        attention_np = attention_weights.cpu().numpy()440        sns.heatmap(attention_np, ax=axes[1, 0], cmap='viridis')441        axes[1, 0].set_title("Attention Heatmap")442        axes[1, 0].set_xlabel("Text Tokens")443        axes[1, 0].set_ylabel("Image Patches")444        445        # Attention overlay on image446        # Resize attention to image size447        attention_map = np.mean(attention_np, axis=1)448        attention_map = attention_map.reshape(int(np.sqrt(len(attention_map))), -1)449        attention_map = cv2.resize(attention_map, (image_np.shape[1], image_np.shape[0]))450        451        axes[1, 1].imshow(image_np)452        axes[1, 1].imshow(attention_map, alpha=0.6, cmap='hot')453        axes[1, 1].set_title("Attention Overlay")454        axes[1, 1].axis('off')455        456        plt.tight_layout()457        return fig