zensvi.cv.Embeddings

class zensvi.cv.Embeddings(model_name: str = 'resnet-18', cuda: bool = False, tensor: bool = True, verbosity: int = 1)

A class for extracting image embeddings using pre-trained models.

load_image(image_path)

Load and preprocess an image from a file path.

Parameters:

image_path – Path to the image file.

Returns:

Loaded and resized image.

Return type:

PIL.Image

get_model_and_layer()

Get the pre-trained model and extraction layer.

Returns:

Contains:

  • torch.nn.Module: The pre-trained model

  • torch.nn.Module: The extraction layer

Return type:

tuple

get_image_embedding(image_path: List[str] | str, tensor: bool = None, cuda: bool = None)

Extract embedding for a single image.

Parameters:

  • image_path – Path to the image file.

  • tensor – Whether to return the embedding as a PyTorch tensor. If None, uses the instance default.

  • cuda – Whether to use CUDA for computation. If None, uses the instance default.

Returns:

The image embedding.

Return type:

Union[torch.Tensor, numpy.ndarray]

generate_embedding(images_path: List[str] | str, dir_embeddings_output: str, batch_size: int = 100, maxWorkers: int = 8, verbosity: int = None)

Generate and save embeddings for multiple images.

Parameters:

  • images_path – Either a directory path containing images or a list of image paths.

  • dir_embeddings_output – Directory where embeddings will be saved as parquet files.

  • batch_size – Number of images to process in each batch. Default is 100.

  • maxWorkers – Maximum number of worker threads for image processing. Default is 8.

  • verbosity – Level of verbosity for progress bars. 0 = no progress bars, 1 = outer loops only, 2 = all loops. If None, uses the instance’s verbosity level.

Returns:

True if embeddings were successfully generated and saved.

Return type:

bool

search_similar_images(image_key: str, embeddings_dir: str, number_of_items: int = 10)

Search for similar images using embeddings.

Parameters:

  • image_key – Key of the query image.

  • embeddings_dir – Directory containing the embeddings parquet file.

  • number_of_items – Number of similar images to return. Default is 10.

Returns:

List of tuples containing (similarity_score, image_key, image_filename).

Return type:

list

get_all_models()

Get all the available models.

Returns:

Dictionary of available models with their configurations.

Return type:

dict