src/utils/cnn_feature_visualization.py

import argparse
import datetime
import os
import random
from typing import Dict, List, Optional, Tuple

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
from torchvision import datasets, transforms

from CNN_model import TrafficSignCNN


class CNNFeatureMapVisualizer:
    """Visualize CNN activations, filters, and class confidence."""

    def __init__(
        self,
        model: torch.nn.Module,
        device: Optional[torch.device] = None,
        class_names: Optional[List[str]] = None,
        mean: Tuple[float, float, float] = (0.5, 0.5, 0.5),
        std: Tuple[float, float, float] = (0.5, 0.5, 0.5),
    ) -> None:
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = model.to(self.device)
        self.model.eval()
        self.class_names = class_names
        self.mean = mean
        self.std = std

    @staticmethod
    def load_model(
        checkpoint_path: Optional[str],
        num_classes: int = 43,
        device: Optional[torch.device] = None,
    ) -> Tuple[torch.nn.Module, bool, str]:
        """Load model from checkpoint; fallback to initialized model if unavailable."""
        device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = TrafficSignCNN(num_classes=num_classes).to(device)

        if not checkpoint_path:
            return model, False, 'No checkpoint provided. Using initialized weights.'

        if not os.path.exists(checkpoint_path):
            msg = f'Checkpoint not found at {checkpoint_path}. Using initialized weights.'
            return model, False, msg

        payload = torch.load(checkpoint_path, map_location=device)

        state_dict = payload
        if isinstance(payload, dict):
            if 'state_dict' in payload:
                state_dict = payload['state_dict']
            elif 'model_state_dict' in payload:
                state_dict = payload['model_state_dict']

        model.load_state_dict(state_dict, strict=True)
        return model, True, f'Loaded checkpoint from {checkpoint_path}'

    def _to_display_image(self, image_tensor: torch.Tensor) -> np.ndarray:
        if image_tensor.ndim == 4:
            image_tensor = image_tensor[0]

        img = image_tensor.detach().cpu().float().clone()
        if img.shape[0] == 1:
            img = img.repeat(3, 1, 1)

        mean = torch.tensor(self.mean, dtype=img.dtype).view(-1, 1, 1)
        std = torch.tensor(self.std, dtype=img.dtype).view(-1, 1, 1)
        if img.shape[0] == mean.shape[0]:
            img = img * std + mean

        img = img.clamp(0.0, 1.0)
        img = img.permute(1, 2, 0).numpy()
        return img

    def _extract_feature_maps(self, image_batch: torch.Tensor) -> Dict[str, torch.Tensor]:
        image_batch = image_batch.to(self.device)
        with torch.no_grad():
            if hasattr(self.model, 'extract_feature_maps'):
                return self.model.extract_feature_maps(image_batch)

            activations: Dict[str, torch.Tensor] = {}
            hooks = []

            for layer_name in ('conv1', 'conv2', 'conv3'):
                if hasattr(self.model, layer_name):
                    module = getattr(self.model, layer_name)
                    hooks.append(
                        module.register_forward_hook(
                            lambda _m, _inp, out, name=layer_name: activations.__setitem__(name, out.detach())
                        )
                    )

            _ = self.model(image_batch)
            for hook in hooks:
                hook.remove()
            return activations

    def _predict_topk(self, image_batch: torch.Tensor, top_k: int = 5) -> Tuple[np.ndarray, np.ndarray]:
        image_batch = image_batch.to(self.device)
        with torch.no_grad():
            logits = self.model(image_batch)
            probs = torch.softmax(logits, dim=1)
            k = min(top_k, probs.shape[1])
            top_probs, top_idx = torch.topk(probs, k=k, dim=1)

        return top_probs[0].detach().cpu().numpy(), top_idx[0].detach().cpu().numpy()

    def _plot_activation_grid(
        self,
        feature_map: torch.Tensor,
        layer_name: str,
        save_path: str,
        max_channels: int = 16,
    ) -> None:
        fmap = feature_map[0].detach().cpu()
        channels = min(max_channels, fmap.shape[0])
        cols = 4
        rows = int(np.ceil(channels / cols))

        fig, axes = plt.subplots(rows, cols, figsize=(3.5 * cols, 3.0 * rows))
        axes = np.array(axes).reshape(-1)

        for idx, ax in enumerate(axes):
            if idx < channels:
                sns.heatmap(fmap[idx].numpy(), cmap='viridis', cbar=False, ax=ax)
                ax.set_title(f'{layer_name} ch{idx}')
            ax.axis('off')

        fig.suptitle(f'Feature Maps - {layer_name}', fontsize=14)
        fig.tight_layout()
        fig.savefig(save_path, dpi=200, bbox_inches='tight')
        plt.close(fig)

    def _plot_aggregated_maps(
        self,
        image_tensor: torch.Tensor,
        feature_maps: Dict[str, torch.Tensor],
        save_path: str,
    ) -> None:
        img = self._to_display_image(image_tensor)

        fig, axes = plt.subplots(1, 4, figsize=(18, 4.5))
        axes[0].imshow(img)
        axes[0].set_title('Input Image')
        axes[0].axis('off')

        for i, layer in enumerate(('conv1', 'conv2', 'conv3'), start=1):
            if layer in feature_maps:
                fmap = feature_maps[layer][0].detach().cpu()
                agg = torch.mean(fmap, dim=0).numpy()
                sns.heatmap(agg, cmap='magma', cbar=False, ax=axes[i])
                axes[i].set_title(f'{layer} mean activation')
                axes[i].axis('off')
            else:
                axes[i].axis('off')
                axes[i].set_title(f'{layer} unavailable')

        fig.suptitle('How the CNN Understands the Image (Aggregated Activations)', fontsize=14)
        fig.tight_layout()
        fig.savefig(save_path, dpi=200, bbox_inches='tight')
        plt.close(fig)

    def _plot_filters(self, save_dir: str, max_kernels: int = 16) -> None:
        for layer_name in ('conv1', 'conv2', 'conv3'):
            if not hasattr(self.model, layer_name):
                continue

            layer = getattr(self.model, layer_name)
            if not hasattr(layer, 'weight'):
                continue

            weights = layer.weight.detach().cpu()
            kernels = min(max_kernels, weights.shape[0])
            cols = 4
            rows = int(np.ceil(kernels / cols))

            fig, axes = plt.subplots(rows, cols, figsize=(3.5 * cols, 3.0 * rows))
            axes = np.array(axes).reshape(-1)

            for idx, ax in enumerate(axes):
                if idx < kernels:
                    kernel = weights[idx]
                    if kernel.shape[0] == 3:
                        rgb = kernel.permute(1, 2, 0).numpy()
                        rgb_min = rgb.min()
                        rgb_max = rgb.max()
                        if rgb_max > rgb_min:
                            rgb = (rgb - rgb_min) / (rgb_max - rgb_min)
                        ax.imshow(rgb)
                    else:
                        # For deeper conv layers, visualize filter strength map.
                        norm_map = torch.norm(kernel, dim=0).numpy()
                        sns.heatmap(norm_map, cmap='coolwarm', cbar=False, ax=ax)
                    ax.set_title(f'{layer_name} k{idx}')
                ax.axis('off')

            fig.suptitle(f'Learned Filters - {layer_name}', fontsize=14)
            fig.tight_layout()
            out = os.path.join(save_dir, f'filters_{layer_name}.png')
            fig.savefig(out, dpi=200, bbox_inches='tight')
            plt.close(fig)

    def _plot_confidence(
        self,
        top_probs: np.ndarray,
        top_idx: np.ndarray,
        save_path: str,
        true_label: Optional[int] = None,
    ) -> None:
        labels = []
        for cls_id in top_idx:
            if self.class_names and int(cls_id) < len(self.class_names):
                labels.append(self.class_names[int(cls_id)])
            else:
                labels.append(f'class_{int(cls_id)}')

        fig, ax = plt.subplots(figsize=(8, 5))
        sns.barplot(x=top_probs, y=labels, color='#4A86E8', orient='h', ax=ax)
        pred_label = labels[0]
        pred_conf = float(top_probs[0])

        title = f'Prediction: {pred_label} ({pred_conf:.3f})'
        if true_label is not None:
            title += f' | True: {true_label}'
        ax.set_title(title)
        ax.set_xlabel('Probability')
        ax.set_ylabel('Class')
        fig.tight_layout()
        fig.savefig(save_path, dpi=200, bbox_inches='tight')
        plt.close(fig)

    def visualize_sample(
        self,
        image_tensor: torch.Tensor,
        sample_id: str,
        output_dir: str,
        true_label: Optional[int] = None,
        max_channels: int = 16,
        top_k: int = 5,
    ) -> None:
        if image_tensor.ndim == 3:
            image_batch = image_tensor.unsqueeze(0)
        else:
            image_batch = image_tensor

        feature_maps = self._extract_feature_maps(image_batch)

        sample_dir = os.path.join(output_dir, f'sample_{sample_id}')
        os.makedirs(sample_dir, exist_ok=True)

        self._plot_aggregated_maps(
            image_tensor=image_batch,
            feature_maps=feature_maps,
            save_path=os.path.join(sample_dir, 'activations_aggregated.png'),
        )

        for layer in ('conv1', 'conv2', 'conv3'):
            if layer in feature_maps:
                self._plot_activation_grid(
                    feature_map=feature_maps[layer],
                    layer_name=layer,
                    save_path=os.path.join(sample_dir, f'feature_maps_{layer}.png'),
                    max_channels=max_channels,
                )

        top_probs, top_idx = self._predict_topk(image_batch, top_k=top_k)
        self._plot_confidence(
            top_probs=top_probs,
            top_idx=top_idx,
            save_path=os.path.join(sample_dir, 'prediction_confidence.png'),
            true_label=true_label,
        )

    def run_from_test_loader(
        self,
        test_loader,
        num_samples: int,
        output_dir: str,
        max_channels: int = 16,
        top_k: int = 5,
    ) -> None:
        dataset = test_loader.dataset
        if len(dataset) == 0:
            raise ValueError('Test dataset is empty.')

        self._plot_filters(output_dir)

        for idx in range(num_samples):
            sample_index = random.randrange(0, len(dataset))
            image, label = dataset[sample_index]
            self.visualize_sample(
                image_tensor=image,
                sample_id=f'{idx+1}_idx{sample_index}',
                output_dir=output_dir,
                true_label=int(label),
                max_channels=max_channels,
                top_k=top_k,
            )


def _build_test_loader(
    test_dir: str,
    image_size: int,
    batch_size: int,
    num_workers: int,
    mean: Tuple[float, float, float],
    std: Tuple[float, float, float],
):
    transform = transforms.Compose(
        [
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std),
        ]
    )
    test_dataset = datasets.ImageFolder(root=test_dir, transform=transform)
    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
    )
    return test_loader, test_dataset.classes


def _timestamped_dir(base_output_dir: str) -> str:
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    out = os.path.join(base_output_dir, f'cnn_feature_maps_{timestamp}')
    os.makedirs(out, exist_ok=True)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description='CNN feature maps, filters, and confidence visualizer')
    parser.add_argument('--checkpoint', type=str, default='', help='Path to model checkpoint (.pth)')
    parser.add_argument('--test-dir', type=str, default='', help='ImageFolder-compatible test directory')
    parser.add_argument('--output-dir', type=str, default='visualization_outputs', help='Base output directory')
    parser.add_argument('--num-classes', type=int, default=43, help='Number of classes in classifier head')
    parser.add_argument('--samples', type=int, default=3, help='Random test samples to visualize')
    parser.add_argument('--top-k', type=int, default=5, help='Top-k classes for confidence bar plot')
    parser.add_argument('--image-size', type=int, default=32, help='Input image size')
    parser.add_argument('--batch-size', type=int, default=32, help='Test loader batch size')
    parser.add_argument('--num-workers', type=int, default=0, help='Dataloader workers')
    parser.add_argument('--max-channels', type=int, default=16, help='Max channels per layer in activation grid')
    parser.add_argument('--seed', type=int, default=42, help='Random seed for sample selection')
    args = parser.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    mean = (0.5, 0.5, 0.5)
    std = (0.5, 0.5, 0.5)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model, loaded, msg = CNNFeatureMapVisualizer.load_model(
        checkpoint_path=args.checkpoint,
        num_classes=args.num_classes,
        device=device,
    )
    print(msg)
    if not loaded:
        print('Warning: continuing with fallback initialized weights.')

    output_dir = _timestamped_dir(args.output_dir)
    visualizer = CNNFeatureMapVisualizer(model=model, device=device, class_names=None, mean=mean, std=std)

    if args.test_dir and os.path.isdir(args.test_dir):
        test_loader, class_names = _build_test_loader(
            test_dir=args.test_dir,
            image_size=args.image_size,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
            mean=mean,
            std=std,
        )
        visualizer.class_names = class_names
        visualizer.run_from_test_loader(
            test_loader=test_loader,
            num_samples=args.samples,
            output_dir=output_dir,
            max_channels=args.max_channels,
            top_k=args.top_k,
        )
        print(f'Done. Saved visualization outputs to: {output_dir}')
        return

    print('No valid --test-dir supplied. Running synthetic fallback sample.')
    visualizer._plot_filters(output_dir)
    synthetic = torch.rand(1, 3, args.image_size, args.image_size)
    visualizer.visualize_sample(
        image_tensor=synthetic,
        sample_id='synthetic_1',
        output_dir=output_dir,
        true_label=None,
        max_channels=args.max_channels,
        top_k=args.top_k,
    )
    print(f'Done. Saved visualization outputs to: {output_dir}')


if __name__ == '__main__':
    main()