class mmpretrain.apis.ImageRetrievalInferencer(model, prototype, prototype_cache=None, prepare_batch_size=8, pretrained=True, device=None, **kwargs)[source]

The inferencer for image to image retrieval.

  • model (BaseModel | str | Config) – A model name or a path to the config file, or a BaseModel object. The model name can be found by ImageRetrievalInferencer.list_models() and you can also query it in Model Zoo Summary.

  • prototype (str | list | dict | DataLoader, BaseDataset) –

    The images to be retrieved. It can be the following types:

    • str: The directory of the the images.

    • list: A list of path of the images.

    • dict: A config dict of the a prototype dataset.

    • BaseDataset: A prototype dataset.

    • DataLoader: A data loader to load the prototype data.

  • prototype_cache (str, optional) – The path of the generated prototype features. If exists, directly load the cache instead of re-generate the prototype features. If not exists, save the generated features to the path. Defaults to None.

  • pretrained (str, optional) – Path to the checkpoint. If None, it will try to find a pre-defined weight from the model you specified (only work if the model is a model name). Defaults to None.

  • device (str, optional) – Device to run inference. If None, the available device will be automatically used. Defaults to None.

  • **kwargs – Other keyword arguments to initialize the model (only work if the model is a model name).


>>> from mmpretrain import ImageRetrievalInferencer
>>> inferencer = ImageRetrievalInferencer(
...     'resnet50-arcface_inshop',
...     prototype='./demo/',
...     prototype_cache='img_retri.pth')
>>> inferencer('demo/cat-dog.png', topk=2)[0][1]
{'match_score': tensor(0.4088, device='cuda:0'),
 'sample_idx': 3,
 'sample': {'img_path': './demo/dog.jpg'}}
__call__(inputs, return_datasamples=False, batch_size=1, **kwargs)[source]

Call the inferencer.

  • inputs (str | array | list) – The image path or array, or a list of images.

  • return_datasamples (bool) – Whether to return results as DataSample. Defaults to False.

  • batch_size (int) – Batch size. Defaults to 1.

  • resize (int, optional) – Resize the long edge of the image to the specified length before visualization. Defaults to None.

  • draw_score (bool) – Whether to draw the match scores. Defaults to True.

  • show (bool) – Whether to display the visualization result in a window. Defaults to False.

  • wait_time (float) – The display time (s). Defaults to 0, which means “forever”.

  • show_dir (str, optional) – If not None, save the visualization results in the specified directory. Defaults to None.


The inference results.

Return type:


static list_models(pattern=None)[source]

List all available model names.


pattern (str | None) – A wildcard pattern to match model names.


a list of model names.

Return type:


Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.