def infer(self,img):
'''
Args:
img: image path or a image array
Returns:
'''
test_pipeline = [LoadImage()] + self.config.test_pipeline
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [torch.device(self.device)])[0]
# forward the model
with torch.no_grad():
result = self.deviceLocator(return_loss=False, rescale=True, **data)
# filter bb according to threshold
filterResult = []
for cls , bbs in enumerate(result[0]): # result:tuple (bb,segm)
thr = self.clsThr[cls]
filterResult.append(bbs[np.nonzero(bbs[:, -1] >= thr)])
return filterResult
def inference_detector(model, img, cfg, device):
if isinstance(cfg, str):
cfg = mmcv.Config.fromfile(cfg)
device = torch.device(device)
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
test_pipeline = cfg.data.test.pipeline
test_pipeline = Compose(test_pipeline)
# prepare data
data = test_pipeline(data)
tensor = data['img'][0].unsqueeze(0).to(device)
img_metas = data['img_metas']
scale_factor = img_metas[0].data['scale_factor']
scale_factor = torch.tensor(scale_factor, dtype=torch.float32, device=device)
with torch.no_grad():
result = model(tensor)
result = list(result)
result[1] = result[1]/scale_factor
return result
def inference_detector(cfg, model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
# build the data pipeline
# We don't want to crop bottom
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[2:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
# forward the model
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
return result
def __init__(self,
ann_file,
pipeline,
load_and_dump_config_name: str = 'load_and_dump_config',
composer_config_name: str = 'composer_config',
generated_objects_fields: Tuple[str,
str] = ('bboxes', 'labels'),
test_mode=False):
self._load_config_filename = ann_file
self._test_mode = test_mode
self._load_and_dump_config_name = load_and_dump_config_name
self._composer_config_name = composer_config_name
self._pipeline = Compose(pipeline)
self._categories_dict = {}
self._generated_objects_fields = generated_objects_fields
self._generated_objects_default_field = self._generated_objects_fields == (
'bboxes', 'labels')
self._trassir_composer: TrassirComposer = self.load_trassir_composer(
self._load_config_filename)
if not self._test_mode:
self._set_group_flag()
# Need for coco wrapper (CocoMapEval)
self._coco = None
self._img_ids = None
self._cat_ids = None
def load_transformed_gt_info(self):
CLASSES = self.CLASSES
img_infos = self.img_infos
transfroms = Compose([
dict(type='LoadImageFromFile'),
dict(type='LoadAnnotations', with_bbox=True),
dict(type='Resize', img_scale=[self._resize], keep_ratio=False)
])
gt_boxes_all = dict()
for idx in range(len(img_infos)):
img_info = img_infos[idx]
ann_info = self.get_ann_info(idx)
results = dict(img_info=img_info, ann_info=ann_info)
if self.proposals is not None:
results['proposals'] = self.proposals[idx]
self.pre_pipeline(results)
results = transfroms(results)
gt_boxes = results['gt_bboxes']
labels = ann_info['labels']
assert len(gt_boxes) == len(labels)
for i, label in enumerate(labels):
w = gt_boxes[i][2] - gt_boxes[i][0]
h = gt_boxes[i][3] - gt_boxes[i][1]
classname = CLASSES[label - 1]
if classname not in gt_boxes_all:
gt_boxes_all[classname] = []
gt_boxes_all[classname].append([w, h])
return gt_boxes_all
def get_data(img, cfg, device):
# import ipdb; ipdb.set_trace()
test_pipeline = Compose([LoadImage()] + cfg.test_pipeline[1:])
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
return data
async def async_inference_detector(model, img):
"""Async inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
Awaitable detection results.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# We don't restore `torch.is_grad_enabled()` value during concurrent
# inference since execution can overlap
torch.set_grad_enabled(False)
result = await model.aforward_test(rescale=True, **data)
return result
def __init__(self,
transforms,
img_scale,
pts_scale_ratio,
flip=False,
flip_direction='horizontal',
pcd_horizontal_flip=False,
pcd_vertical_flip=False):
self.transforms = Compose(transforms)
self.img_scale = img_scale if isinstance(img_scale,
list) else [img_scale]
self.pts_scale_ratio = pts_scale_ratio \
if isinstance(pts_scale_ratio, list) else[float(pts_scale_ratio)]
assert mmcv.is_list_of(self.img_scale, tuple)
assert mmcv.is_list_of(self.pts_scale_ratio, float)
self.flip = flip
self.pcd_horizontal_flip = pcd_horizontal_flip
self.pcd_vertical_flip = pcd_vertical_flip
self.flip_direction = flip_direction if isinstance(
flip_direction, list) else [flip_direction]
assert mmcv.is_list_of(self.flip_direction, str)
if not self.flip and self.flip_direction != ['horizontal']:
warnings.warn(
'flip_direction has no effect when flip is set to False')
if (self.flip and not any([(t['type'] == 'RandomFlip3D'
or t['type'] == 'RandomFlip')
for t in transforms])):
warnings.warn(
'flip has no effect when RandomFlip is not in transforms')
def get_img_meta(model, idx, dataset_train):
cfg = model.cfg
# get gt_ann
#dataset_train = build_dataset(cfg.data.train)
ann_info = dataset_train.get_ann_info(idx)
img_path = '{}{}'.format(cfg.data.train.img_prefix,
ann_info['seg_map'].replace('png', 'jpg'))
gt_bboxes = ann_info['bboxes']
gt_labels = ann_info['labels']
gt_segms = ann_info['masks']
# get img_meta
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
data = dict(img=img_path)
data = test_pipeline(data)
# dict_keys(['img_meta', 'img'])
img_metas = data['img_meta'][0].data
# dict_keys(['filename', 'ori_shape', 'img_shape', 'pad_shape', 'scale_factor', 'flip', 'img_norm_cfg'])
img_tensor = data['img'][0].data
# get FPN feat_size
c, h, w = img_tensor.shape
input_img_feat = img_tensor.view(1, c, h, w)
device = input_img_feat.device
fpn_feats = model.extract_feat(input_img_feat)
return img_path, fpn_feats, gt_bboxes, gt_labels, img_metas, gt_segms
def __init__(self, datadir, pipeline, stage='train', repeat=1):
self.images = get_spec_files(datadir, ext=IMG_EXT, oswalk=True)
if stage == 'train':
self.images = self.images * repeat
self.pipeline = Compose(pipeline)
self.flag = np.zeros(len(self), 'u1')
self.stage = stage
def export_onnx(model, img, onnx_save_path, verbose=False, output_names=None):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
if hasattr(model, 'forward_export'):
model.forward = model.forward_export
else:
raise NotImplementedError(
'Export onnx is currently not currently supported with {}'.format(
model.__class__.__name__))
torch.onnx.export(
model, # model being run
data['img'][0], # model input (or a tuple for multiple inputs)
onnx_save_path, # where to save the model (can be a file or file-like object)
verbose=verbose,
output_names=output_names)
async def async_inference_detector(model, img):
"""Async inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
img (str | ndarray): Either image files or loaded images.
Returns:
Awaitable detection results.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# We don't restore `torch.is_grad_enabled()` value during concurrent
# inference since execution can overlap
torch.set_grad_enabled(False)
result = await model.aforward_test(rescale=True, **data)
return result
class CutConfig(object):
# process module
train_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='CutROI', training=False),
dict(type='CutImage', training=False, window=(1000, 1000), step=(500, 500), order_index=False,
is_keep_none=True)
]
compose = Compose(train_pipeline)
# data module
img_dir = "/home/lifeng/undone-work/DefectNet/tools/data/tile/raw/tile_round1_testA_20201231/testA_imgs"
test_file = "/home/lifeng/undone-work/dataset/detection/tile/annotations/instance_testA.json"
save_file = "/home/lifeng/undone-work/DetCompetition/mmdet-v2/work_dirs/tile/baseline_cut_1000x1000/do_submit_testA.json"
original_coco = COCO(test_file)
label2name = {x['id']: x['name'] for x in original_coco.dataset['categories']}
main_thread_lock = threading.Lock()
save_results = []
num_workers = 7
process_cnt = 0
# inference module
device = 'cuda:0'
config_file = '/home/lifeng/undone-work/DefectNet/configs/tile/baseline_model_2000x2000.py'
checkpoint_file = '/data/liphone/detcomp/mmdet-v2/tile/baseline_cut_1000x1000/epoch_12.pth'
model = init_detector(config_file, checkpoint_file, device=device)
def inference_detector(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
data['img'][0] = data['img'][0].cuda()
data['img_metas'] = data['img_metas'][0].data
else:
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
# forward the model
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)[0]
return result
def text_model_inference(model, input_sentence):
"""Inference text(s) with the entity recognizer.
Args:
model (nn.Module): The loaded recognizer.
input_sentence (str): A text entered by the user.
Returns:
result (dict): Predicted results.
"""
assert isinstance(input_sentence, str)
cfg = model.cfg
test_pipeline = Compose(cfg.data.test.pipeline)
data = {'text': input_sentence, 'label': {}}
# build the data pipeline
data = test_pipeline(data)
if isinstance(data['img_metas'], dict):
img_metas = data['img_metas']
else:
img_metas = data['img_metas'].data
assert isinstance(img_metas, dict)
img_metas = {
'input_ids': img_metas['input_ids'].unsqueeze(0),
'attention_masks': img_metas['attention_masks'].unsqueeze(0),
'token_type_ids': img_metas['token_type_ids'].unsqueeze(0),
'labels': img_metas['labels'].unsqueeze(0)
}
# forward the model
with torch.no_grad():
result = model(None, img_metas, return_loss=False)
return result
def inference_detector2(model, img_path):
cfg = model.cfg
device = next(model.parameters()).device # model device
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
data = dict(img=img_path)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# Use torchvision ops for CPU mode instead
for m in model.modules():
if isinstance(m, (RoIPool, RoIAlign)):
if not m.aligned:
# aligned=False is not implemented on CPU
# set use_torchvision on-the-fly
m.use_torchvision = True
warnings.warn('We set use_torchvision=True in CPU mode.')
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
imgs = data['img'][0]
img_metas = data['img_metas'][0]
return imgs, img_metas
def inference_detector(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# forward the model
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
return result
def inference_detector(model, img, cfg, device):
if isinstance(cfg, str):
cfg = mmcv.Config.fromfile(cfg)
device = torch.device(device)
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
tensor = data['img'][0].unsqueeze(0).to(device)
img_metas = data['img_metas']
scale_factor = img_metas[0].data['scale_factor']
scale_factor = torch.tensor(scale_factor,
dtype=torch.float32,
device=device)
with torch.no_grad():
result = model(tensor)
result = list(result)
result[1] = result[1] / scale_factor
return result
def get_fake_input(cfg, orig_img_shape=(128, 128, 3), device='cuda'):
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
data = dict(img=np.zeros(orig_img_shape, dtype=np.uint8))
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
return data
def inference_recognizer(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
mask_h, mask_w = cfg.data.test.pipeline[1]['img_scale']
text_max_len = cfg.text_max_len + 2
data = dict(img=img,
target_variable=torch.zeros(text_max_len, 1, dtype=torch.long),
mask=np.zeros((1, mask_h, mask_w)))
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device.index])[0]
# forward the model
with torch.no_grad():
pre, score = model(return_loss=False, rescale=True, **data)
return pre, score
def backbone(self, images, **kwargs):
r"""Returns list of backbone features and transformed images as well as meta info.
"""
from mmdet.apis.inference import inference_detector, LoadImage
from mmdet.datasets.pipelines import Compose
from mmcv.parallel import collate, scatter
model = self.module
cfg = model.cfg
device = next(model.parameters()).device
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
results = []
for img in images:
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
img = data['img'][0]
img_meta = data['img_meta'][0]
data['img'] = img
data['img_meta'] = img_meta
data['feats'] = model.extract_feat(img)
results.append(data)
#print(img.shape, img_meta)
#return model.backbone(images.tensors), images, original_image_sizes
return results
def prefetch_img_metas(cfg, ori_wh):
w, h = ori_wh
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
test_pipeline = Compose(cfg.data.test.pipeline)
data = {'img': np.zeros((h, w, 3), dtype=np.uint8)}
data = test_pipeline(data)
img_metas = data['img_metas'][0].data
return img_metas
def ImagePreprocess(image, cfg):
data = dict(img=image)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data['img'] = torch.stack(data['img'], dim=0)
return data
def inference_model(model, imgs, frame_id):
if isinstance(imgs, (list, tuple)):
is_batch = True
else:
imgs = [imgs]
is_batch = False
cfg = model.cfg
device = next(model.parameters()).device # model device
if isinstance(imgs[0], np.ndarray):
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
test_pipeline = Compose(cfg.data.test.pipeline)
datas = []
for img in imgs:
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img, frame_id=frame_id)
else:
# add information into dict
data = dict(img_info=dict(filename=img, frame_id=frame_id),
img_prefix=None)
# build the data pipeline
data = test_pipeline(data)
datas.append(data)
data = collate(datas, samples_per_gpu=len(imgs))
# just get the actual data from DataContainer
data['img_metas'] = [img_metas.data[0] for img_metas in data['img_metas']]
data['img'] = [img.data[0] for img in data['img']]
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
for m in model.modules():
assert not isinstance(
m, RoIPool
), 'CPU inference with RoIPool is not supported currently.'
# forward the model
with torch.no_grad():
results = model(return_loss=False,
rescale=True,
detection_only=True,
**data)
if not is_batch:
return results[0]
else:
return results
def __init__(self,
ann_file,
pipeline,
classes=None,
data_root=None,
img_prefix='',
seg_prefix=None,
proposal_file=None,
test_mode=False,
filter_empty_gt=True,
mosaic_ratio=0.5
):
self.ann_file = ann_file
self.data_root = data_root
self.img_prefix = img_prefix
self.seg_prefix = seg_prefix
self.proposal_file = proposal_file
self.test_mode = test_mode
self.filter_empty_gt = filter_empty_gt
self.CLASSES = self.get_classes(classes)
self.mosaic_ratio = mosaic_ratio
# join paths if data_root is specified
if self.data_root is not None:
if not osp.isabs(self.ann_file):
self.ann_file = osp.join(self.data_root, self.ann_file)
if not (self.img_prefix is None or osp.isabs(self.img_prefix)):
self.img_prefix = osp.join(self.data_root, self.img_prefix)
if not (self.seg_prefix is None or osp.isabs(self.seg_prefix)):
self.seg_prefix = osp.join(self.data_root, self.seg_prefix)
if not (self.proposal_file is None
or osp.isabs(self.proposal_file)):
self.proposal_file = osp.join(self.data_root,
self.proposal_file)
# load annotations (and proposals)
self.data_infos = self.load_annotations(self.ann_file)
# filter data infos if classes are customized
if self.custom_classes:
self.data_infos = self.get_subset_by_classes()
if self.proposal_file is not None:
self.proposals = self.load_proposals(self.proposal_file)
else:
self.proposals = None
# filter images too small
if not test_mode:
valid_inds = self._filter_imgs()
self.data_infos = [self.data_infos[i] for i in valid_inds]
if self.proposals is not None:
self.proposals = [self.proposals[i] for i in valid_inds]
# set group flag for the sampler
if not self.test_mode:
self._set_group_flag()
# processing pipeline
self.pipeline = Compose(pipeline)
def data_init(self, img_path):
# its main purpose is to obtain the img_metas.
# besides, we compare the image we transform with self.data['img'].
cfg = self.model.cfg
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
self.data = dict(img=img_path)
self.data = test_pipeline(self.data)
self.data = collate([self.data], samples_per_gpu=1)
self.data = scatter(self.data, [device])[0]
def _prepare_input_img(imgs, test_pipeline: Iterable[dict]):
"""Inference image(s) with the detector.
Args:
imgs (str/ndarray or list[str/ndarray] or tuple[str/ndarray]):
Either image files or loaded images.
test_pipeline (Iterable[dict]): Test pipline of configuration.
Returns:
result (dict): Predicted results.
"""
if isinstance(imgs, (list, tuple)):
if not isinstance(imgs[0], (np.ndarray, str)):
raise AssertionError('imgs must be strings or numpy arrays')
elif isinstance(imgs, (np.ndarray, str)):
imgs = [imgs]
else:
raise AssertionError('imgs must be strings or numpy arrays')
test_pipeline = replace_ImageToTensor(test_pipeline)
test_pipeline = Compose(test_pipeline)
datas = []
for img in imgs:
# prepare data
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
data = test_pipeline(data)
# get tensor from list to stack for batch mode (text detection)
datas.append(data)
if isinstance(datas[0]['img'], list) and len(datas) > 1:
raise Exception('aug test does not support '
f'inference with batch size '
f'{len(datas)}')
data = collate(datas, samples_per_gpu=len(imgs))
# process img_metas
if isinstance(data['img_metas'], list):
data['img_metas'] = [
img_metas.data[0] for img_metas in data['img_metas']
]
else:
data['img_metas'] = data['img_metas'].data
if isinstance(data['img'], list):
data['img'] = [img.data for img in data['img']]
if isinstance(data['img'][0], list):
data['img'] = [img[0] for img in data['img']]
else:
data['img'] = data['img'].data
return data
def inference_detector(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# Use torchvision ops for CPU mode instead
for m in model.modules():
if isinstance(m, (RoIPool, RoIAlign)):
if not m.aligned:
# aligned=False is not implemented on CPU
# set use_torchvision on-the-fly
m.use_torchvision = True
warnings.warn('We set use_torchvision=True in CPU mode.')
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
# forward the model
torch.cuda.synchronize()
start_time = time.perf_counter()
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)[0]
torch.cuda.synchronize()
inf_time = time.perf_counter() - start_time
return result, inf_time
async def async_inference_detector(model, imgs):
"""Async inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
img (str | ndarray): Either image files or loaded images.
Returns:
Awaitable detection results.
"""
if not isinstance(imgs, (list, tuple)):
imgs = [imgs]
cfg = model.cfg
device = next(model.parameters()).device # model device
if isinstance(imgs[0], np.ndarray):
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
test_pipeline = Compose(cfg.data.test.pipeline)
datas = []
for img in imgs:
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
data = test_pipeline(data)
datas.append(data)
data = collate(datas, samples_per_gpu=len(imgs))
# just get the actual data from DataContainer
data['img_metas'] = [img_metas.data[0] for img_metas in data['img_metas']]
data['img'] = [img.data[0] for img in data['img']]
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
for m in model.modules():
assert not isinstance(
m, RoIPool
), 'CPU inference with RoIPool is not supported currently.'
# We don't restore `torch.is_grad_enabled()` value during concurrent
# inference since execution can overlap
torch.set_grad_enabled(False)
results = await model.aforward_test(rescale=True, **data)
return results
def inference_detector(model, img):
cfg = model.cfg
device = next(model.parameters()).device
test_pipeline = [LoadImage()] + cfg.test_pipeline[1:]
test_pipeline = Compose(test_pipeline)
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
return result
