AttentionMapsHelper Module

Provides utilities for generating and visualizing attention maps in deep learning models.

class HMB.AttentionMapsHelper.LogitsModelWrapper(model)

Bases: Module

Wrapper for a model to return only logits from the forward pass.

Parameters:

model (torch.nn.Module) – The model to wrap.

Returns:

Logits output from the model.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

HMB.AttentionMapsHelper.HuggingFaceModel(modelName, numClasses, modelCheckpointPath, size, device, meanValues=[0.485, 0.456, 0.406], stdValues=[0.229, 0.224, 0.225])

Load a Hugging Face Vision Transformer model and its checkpoint, and prepare preprocessing transforms.

Parameters:

• modelName (str) – Name of the Hugging Face model to use.

• numClasses (int) – Number of output classes.

• modelCheckpointPath (str) – Path to the trained model checkpoint.

• size (int) – Image size for resizing and visualization.

• device (str or torch.device) – Device to run the model on (“cuda” or “cpu”).

• meanValues (list) – Mean values for normalization.

• stdValues (list) – Standard deviation values for normalization.

Returns:

A tuple containing:

• model (torch.nn.Module): Loaded model.

• vitTargetLayer (torch.nn.Module): Target layer for CAM extraction.

• transform (torchvision.transforms.Compose): Image preprocessing pipeline.

Return type:

tuple

Notes

• Loads the model weights from the checkpoint.

• Sets the model to evaluation mode.

• Prepares the image transform for inference.

Examples

import HMB.AttentionMapsHelper as amh

model, vitTargetLayer, transform = amh.HuggingFaceModel(
  modelName="google/vit-base-patch16-224",
  numClasses=3,
  modelCheckpointPath="/path/to/checkpoint.pth",
  size=224,
  device="cuda",
  meanValues=[0.485, 0.456, 0.406],
  stdValues=[0.229, 0.224, 0.225]
)

HMB.AttentionMapsHelper.GetDefaultVitTargetLayer(model)

Dynamically select a suitable target layer for CAM from a timm model. Tries common patterns for ViT, Swin, MaxViT, ConvNeXtV2, etc. Skips Identity layers as they are not suitable for CAM.

Parameters:

model (torch.nn.Module) – The model to inspect.

Returns:

The selected target layer for CAM.

Return type:

torch.nn.Module

HMB.AttentionMapsHelper.TimmModel(modelName, numClasses, modelCheckpointPath, device, targetLayer=None)

Load a timm Vision Transformer model and its checkpoint, and prepare preprocessing transforms.

Parameters:

• modelName (str) – Name of the timm model to use.

• numClasses (int) – Number of output classes.

• modelCheckpointPath (str) – Path to the trained model checkpoint.

• device (str or torch.device) – Device to run the model on (“cuda” or “cpu”).

• targetLayer (str or None) – Optional string specifying the target layer for CAM.

Returns:

A tuple containing:

• model (torch.nn.Module): Loaded model.

• vitTargetLayer (torch.nn.Module): Target layer for CAM extraction.

• transform (torchvision.transforms.Compose): Image preprocessing pipeline.

Return type:

tuple

Notes

• Loads the model weights from the checkpoint.

• Sets the model to evaluation mode.

• Prepares the image transform for inference.

• Tested on:

  – timm/eva02_large_patch14_448.mim_m38m_ft_in22k_in1k from https://huggingface.co/timm/eva02_large_patch14_448.mim_m38m_ft_in22k_in1k

Examples

import HMB.AttentionMapsHelper as amh

model, vitTargetLayer, transform = amh.TimmModel(
  modelName="eva02_large_patch14_448.mim_m38m_ft_in22k_in1k",
  numClasses=3,
  modelCheckpointPath="/path/to/checkpoint.pth",
  device="cuda"
)

class HMB.AttentionMapsHelper.AttentionMapsVisualizer(baseFolder, dataFolder, modelName, modelCheckpointPath, modelType='Timm', size=448, doReshape=False, device=None)

Bases: object

Visualize attention maps for images using various CAM methods on a Vision Transformer (ViT) model.

Parameters:

• baseFolder (str) – Base directory containing model and dataset.

• dataFolder (str) – Directory containing image data organized by class.

• modelName (str) – Name of the timm model to use.

• modelCheckpointPath (str) – Path to the trained model checkpoint.

• modelType (str) – Type of model (“Timm” and “HuggingFace”). Default is “Timm”.

• size (int) – Image size for resizing and visualization. Default is 448.

• doReshape (bool) – Whether to reshape transformer outputs for CAM. Default is False.

• device (str or torch.device) – Device to run the model on (“cuda” or “cpu”).

Variables:

• model (torch.nn.Module) – Loaded vision transformer model.

• vitTargetLayer (torch.nn.Module) – Target layer for CAM extraction.

• transform (torchvision.transforms.Compose) – Image preprocessing pipeline.

• classes (list) – List of class names from dataset.

• numClasses (int) – Number of classes.

Notes

• Supports GradCAM, ScoreCAM, EigenCAM, and other CAM methods.

• Can select specific images or random images per class for visualization.

• Figure size, CAM methods, and output options are customizable.

• Saves the resulting attention map grid as a PNG file.

Examples

import HMB.AttentionMapsHelper as amh

visualizer = amh.AttentionMapsVisualizer(
  baseFolder="/path/to/base/folder",
  dataFolder="/path/to/data/folder",
  modelName="eva02_large_patch14_448.mim_m38m_ft_in22k_in1k",
  modelCheckpointPath="/path/to/checkpoint.pth",
  modelType="Timm",
  size=448,
  device="cuda"
)
visualizer.VisualizeAttentionMaps(
  cams=["GradCAM", "ScoreCAM"],
  figSize=(14, 8),
  imagesPerClass=2,
  save=True,
  display=True,
  outPrefix="AttentionMaps",
  dpi=300,
  fontSize=10,
  alpha=0.4,
  selectImages=None,
  allowedExtensions=(".jpg", ".jpeg", ".png", ".bmp")
)

static ReshapeTransform(outputs, height, width)

Reshape the transformer outputs to a 4D tensor suitable for CAM extraction.

Parameters:

• outputs (torch.Tensor) – Outputs from the transformer model.

• height (int) – Height of the feature map.

• width (int) – Width of the feature map.

Returns:

Reshaped tensor of shape (batch_size, channels, height, width).

Return type:

torch.Tensor

VisualizeAttentionMaps(cams=None, alpha=0.35, save=True, display=False, dpi=300, outPrefix='AttentionMapResults', figSize=(12, 10), fontSize=12, imagesPerClass=1, selectImages=None, allowedExtensions=('.jpg', '.jpeg', '.png', '.bmp'), doAverage=True)

Visualize and save attention maps for images using specified CAM methods.

Parameters:

• cams (list) – List of CAM classes to use (default: [GradCAM, ScoreCAM, EigenCAM]).

• alpha (float) – Transparency for overlaying heatmap.

• save (bool) – Whether to save the resulting figure.

• display (bool) – Whether to display the figure.

• dpi (int) – DPI for saving the figure.

• outPrefix (str) – Prefix for output file name.

• figSize (tuple) – Figure size in inches.

• imagesPerClass (int) – Number of images per class to visualize.

• selectImages (dict or None) – Optional dict mapping class names to list of image filenames.

• allowedExtensions (tuple) – Allowed image file extensions.

• doAverage (bool) – Whether to compute and display the average overlay per class.

Notes

• If selectImages is provided, it should be a dictionary where keys are class names and values are lists of image filenames to visualize for that class. If not provided, random images will be selected.

• The resulting figure will have rows corresponding to classes and columns corresponding to CAM methods and images.

• The output file will be saved in the base folder with a timestamp.

Raises:

• AssertionError – If no image files are found in the data folder.

• ValueError – If an unsupported CAM method is specified.