Spaces:
Runtime error
Runtime error
File size: 4,649 Bytes
4d0eb62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Callable, List, Optional, Union
import torch
from mmcv.image import imread
from mmengine.config import Config
from mmengine.dataset import Compose, default_collate
from mmpretrain.registry import TRANSFORMS
from .base import BaseInferencer, InputType
from .model import list_models
class FeatureExtractor(BaseInferencer):
"""The inferencer for extract features.
Args:
model (BaseModel | str | Config): A model name or a path to the config
file, or a :obj:`BaseModel` object. The model name can be found
by ``FeatureExtractor.list_models()`` and you can also query it in
:doc:`/modelzoo_statistics`.
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).
Example:
>>> from mmpretrain import FeatureExtractor
>>> inferencer = FeatureExtractor('resnet50_8xb32_in1k', backbone=dict(out_indices=(0, 1, 2, 3)))
>>> feats = inferencer('demo/demo.JPEG', stage='backbone')[0]
>>> for feat in feats:
>>> print(feat.shape)
torch.Size([256, 56, 56])
torch.Size([512, 28, 28])
torch.Size([1024, 14, 14])
torch.Size([2048, 7, 7])
""" # noqa: E501
def __call__(self,
inputs: InputType,
batch_size: int = 1,
**kwargs) -> dict:
"""Call the inferencer.
Args:
inputs (str | array | list): The image path or array, or a list of
images.
batch_size (int): Batch size. Defaults to 1.
**kwargs: Other keyword arguments accepted by the `extract_feat`
method of the model.
Returns:
tensor | Tuple[tensor]: The extracted features.
"""
ori_inputs = self._inputs_to_list(inputs)
inputs = self.preprocess(ori_inputs, batch_size=batch_size)
preds = []
for data in inputs:
preds.extend(self.forward(data, **kwargs))
return preds
@torch.no_grad()
def forward(self, inputs: Union[dict, tuple], **kwargs):
inputs = self.model.data_preprocessor(inputs, False)['inputs']
outputs = self.model.extract_feat(inputs, **kwargs)
def scatter(feats, index):
if isinstance(feats, torch.Tensor):
return feats[index]
else:
# Sequence of tensor
return type(feats)([scatter(item, index) for item in feats])
results = []
for i in range(inputs.shape[0]):
results.append(scatter(outputs, i))
return results
def _init_pipeline(self, cfg: Config) -> Callable:
test_pipeline_cfg = cfg.test_dataloader.dataset.pipeline
if test_pipeline_cfg[0]['type'] == 'LoadImageFromFile':
# Image loading is finished in `self.preprocess`.
test_pipeline_cfg = test_pipeline_cfg[1:]
test_pipeline = Compose(
[TRANSFORMS.build(t) for t in test_pipeline_cfg])
return test_pipeline
def preprocess(self, inputs: List[InputType], batch_size: int = 1):
def load_image(input_):
img = imread(input_)
if img is None:
raise ValueError(f'Failed to read image {input_}.')
return dict(
img=img,
img_shape=img.shape[:2],
ori_shape=img.shape[:2],
)
pipeline = Compose([load_image, self.pipeline])
chunked_data = self._get_chunk_data(map(pipeline, inputs), batch_size)
yield from map(default_collate, chunked_data)
def visualize(self):
raise NotImplementedError(
"The FeatureExtractor doesn't support visualization.")
def postprocess(self):
raise NotImplementedError(
"The FeatureExtractor doesn't need postprocessing.")
@staticmethod
def list_models(pattern: Optional[str] = None):
"""List all available model names.
Args:
pattern (str | None): A wildcard pattern to match model names.
Returns:
List[str]: a list of model names.
"""
return list_models(pattern=pattern)
|