ImagesToEmbeddings Module¶

Implements different methods to convert images to embeddings.

class HMB.ImagesToEmbeddings.TransformersEmbeddingModel(modelName, device)[source]¶

Bases: object

A class to extract embeddings from images using pre-trained models from the Hugging Face Transformers library.

\[\mathrm{embedding} = \mathrm{model}(I)_{\mathrm{CLS}}\]

where the CLS token (or first token) is used as the image-level embedding.

Initialize the TransformersEmbeddingModel with a specified model name and device.

Parameters:

modelName (str) – Name of the pre-trained model to load from Hugging Face.
device (str or torch.device) – Device to run the model on (e.g., “cuda”, “cpu”).

__init__(modelName, device)[source]¶

Initialize the TransformersEmbeddingModel with a specified model name and device.

Parameters:

modelName (str) – Name of the pre-trained model to load from Hugging Face.
device (str or torch.device) – Device to run the model on (e.g., “cuda”, “cpu”).

LoadModel()[source]¶

Load the pre-trained model and processor from the specified model name.

Returns:: The loaded pre-trained model. processor (transformers.AutoImageProcessor): The loaded image processor.
Return type:: model (torch.nn.Module)

GetEmbedding(imagePath)[source]¶

Extract embedding from an image using the loaded model and processor.

Parameters:: imagePath (str) – Path to the input image.
Returns:: The extracted embedding as a numpy array.
Return type:: embedding (numpy.ndarray)

HMB.ImagesToEmbeddings.ExtractEmbeddingsTimm(datasetFolder, outputPicklePath, modelName='hf-hub:paige-ai/Virchow2', mlpLayer=None, actLayer=<class 'torch.nn.modules.activation.SiLU'>, device=None)[source]¶

Extract embeddings from images in a dataset folder using a specified model from the timm library.

Parameters:

datasetFolder (str) – Path to the root folder containing subfolders for each class, each with images.
outputPicklePath (str) – Path to save the output pickle file containing the embeddings lookup table.
modelName (str) – Name of the timm model to use. Default is “hf-hub:paige-ai/Virchow2”.
mlpLayer (nn.Module) – MLP layer class to use in the model. Default is None.
actLayer (nn.Module) – Activation layer class to use in the model. Default is torch.nn.SiLU.
device (str or torch.device, optional) – Device to run the model on (e.g., “cuda”, “cpu”). If None, uses CUDA if available.

Examples

from HMB.ImagesToEmbeddings import ExtractEmbeddingsTimm
datasetFolder = "path/to/dataset"
outputPickle = "embeddings.pkl"
ExtractEmbeddingsTimm(datasetFolder, outputPickle)

Notes

The function composes a per-image embedding by concatenating the class token and the mean of patch tokens:

e = [class_token ; mean(patch_tokens)]