Spaces:
Runtime error
Runtime error
| import os | |
| from typing import Optional, Tuple | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from src.densenet import DenseNet121 | |
| def _preprocess_image(data_tensor: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Convert a (C, H, W) tensor into: | |
| - normalized float RGB image in [0, 1] | |
| - uint8 RGB image in [0, 255] | |
| """ | |
| image = data_tensor.cpu().numpy() | |
| image = np.transpose(image, (1, 2, 0)) | |
| image = (image - image.min()) / (image.max() - image.min()) | |
| image_uint8 = (image * 255).astype(np.uint8) | |
| return image, image_uint8 | |
| def _process_smoothgrad_map( | |
| smoothgrad_map: np.ndarray, | |
| img_shape: Tuple[int, int], | |
| percentile: float, | |
| colormap: str, | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Process raw SmoothGrad map into visualization-ready formats. | |
| """ | |
| smoothgrad_map = np.transpose(smoothgrad_map, (1, 2, 0)) | |
| smoothgrad_map = np.mean(smoothgrad_map, axis=2) | |
| smoothgrad_map = cv2.resize(smoothgrad_map, (img_shape[1], img_shape[0])) | |
| p_clip = np.percentile(smoothgrad_map, percentile) | |
| if p_clip > 0: | |
| smoothgrad_map = np.clip(smoothgrad_map, 0, p_clip) | |
| smoothgrad_map = (smoothgrad_map - smoothgrad_map.min()) / (p_clip - smoothgrad_map.min() + 1e-8) | |
| smoothgrad_heatmap = (smoothgrad_map * 255).astype(np.uint8) | |
| heatmap = cv2.applyColorMap( | |
| smoothgrad_heatmap, | |
| cv2.COLORMAP_HOT if colormap == "hot" else cv2.COLORMAP_VIRIDIS, | |
| ) | |
| return smoothgrad_map, heatmap | |
| def _get_target_layer(model: nn.Module, layer_name: Optional[str] = None): | |
| """ | |
| Return DenseNet feature layer used for GradCAM. | |
| Default layer is `denseblock3` to reduce center-biased CAMs compared with | |
| the original `norm5_last`. | |
| You can override with: | |
| - function arg: `layer_name` | |
| - env var: `XDL_TARGET_LAYER` | |
| Supported layer names: | |
| `denseblock3`, `transition2`, `transition1`, `denseblock4`, `transition3`, `norm5_last`. | |
| """ | |
| if not isinstance(model, DenseNet121): | |
| raise TypeError( | |
| f"Unsupported model type for this demo: {type(model).__name__}. Expected DenseNet121." | |
| ) | |
| requested = (layer_name or os.getenv("XDL_TARGET_LAYER") or "denseblock3").strip().lower() | |
| aliases = { | |
| "default": "denseblock3", | |
| "last": "norm5_last", | |
| "norm5": "norm5_last", | |
| } | |
| selected = aliases.get(requested, requested) | |
| features = model.densenet_model.features | |
| layer_map = { | |
| "denseblock3": features.denseblock3, | |
| "transition2": features.transition2, | |
| "transition1": features.transition1, | |
| "denseblock4": features.denseblock4, | |
| "transition3": features.transition3, | |
| "norm5_last": features[-1], | |
| } | |
| if selected not in layer_map: | |
| supported = ", ".join(layer_map.keys()) | |
| raise ValueError( | |
| f"Unsupported XDL target layer '{selected}'. Supported layers: {supported}" | |
| ) | |
| return layer_map[selected] | |
| def smoothgrad( | |
| model: nn.Module, | |
| input_tensor: torch.Tensor, | |
| target_class: int, | |
| n_samples: int = 100, | |
| noise_level: float = 0.05, | |
| batch_size: Optional[int] = None, | |
| use_amp: bool = False, | |
| ) -> np.ndarray: | |
| """ | |
| Compute SmoothGrad saliency map for one input tensor of shape (1, C, H, W). | |
| Notes: | |
| - This implementation batches noisy samples to reduce per-step overhead. | |
| - If `input_tensor` and `model` are on CUDA, computation runs on GPU. | |
| """ | |
| if n_samples <= 0: | |
| raise ValueError(f"n_samples must be > 0, got {n_samples}") | |
| if input_tensor.ndim != 4 or input_tensor.shape[0] != 1: | |
| raise ValueError( | |
| f"input_tensor must have shape (1, C, H, W), got {tuple(input_tensor.shape)}" | |
| ) | |
| if batch_size is None: | |
| batch_size = min(16, n_samples) | |
| if batch_size <= 0: | |
| raise ValueError(f"batch_size must be > 0, got {batch_size}") | |
| model_param = next(model.parameters(), None) | |
| if model_param is not None and model_param.device != input_tensor.device: | |
| raise ValueError( | |
| f"Model device ({model_param.device}) and input device ({input_tensor.device}) must match." | |
| ) | |
| model.eval() | |
| accumulated_gradients = torch.zeros_like(input_tensor) | |
| input_range = torch.max(input_tensor) - torch.min(input_tensor) | |
| scaled_noise = noise_level * input_range | |
| use_cuda_amp = use_amp and input_tensor.device.type == "cuda" | |
| for start in range(0, n_samples, batch_size): | |
| current_batch = min(batch_size, n_samples - start) | |
| expanded_input = input_tensor.expand(current_batch, -1, -1, -1) | |
| noise = torch.randn_like(expanded_input) * scaled_noise | |
| noisy_input = (expanded_input + noise).clone().detach().requires_grad_(True) | |
| if use_cuda_amp: | |
| with torch.autocast(device_type="cuda", dtype=torch.float16): | |
| output = model(noisy_input) | |
| else: | |
| output = model(noisy_input) | |
| if output.ndim == 1: | |
| output = output.unsqueeze(0) | |
| model.zero_grad(set_to_none=True) | |
| output[:, target_class].sum().backward() | |
| if noisy_input.grad is not None: | |
| accumulated_gradients += noisy_input.grad.data.sum(dim=0, keepdim=True) | |
| smooth = accumulated_gradients / n_samples | |
| smooth = torch.abs(smooth) | |
| smooth_min = torch.min(smooth) | |
| smooth_max = torch.max(smooth) | |
| smooth = (smooth - smooth_min) / (smooth_max - smooth_min + 1e-8) | |
| return smooth[0].cpu().numpy() | |