diff --git a/mmdet/apis/__init__.py b/mmdet/apis/__init__.py index 39544f2bf41a46708943e6f1672fef0b1df03e11..030b7de41026755359535cc309e39c7b4e0efb66 100644 --- a/mmdet/apis/__init__.py +++ b/mmdet/apis/__init__.py @@ -1,8 +1,8 @@ from .env import init_dist, get_root_logger, set_random_seed from .train import train_detector -from .inference import inference_detector +from .inference import inference_detector, show_result __all__ = [ 'init_dist', 'get_root_logger', 'set_random_seed', 'train_detector', - 'inference_detector' + 'inference_detector', 'show_result' ] diff --git a/mmdet/apis/inference.py b/mmdet/apis/inference.py index 0addd598e5043220e9bbc1ad594c757fafa75481..a87323cee1aba6e97c75f2a563c5337bc5fe32ff 100644 --- a/mmdet/apis/inference.py +++ b/mmdet/apis/inference.py @@ -23,19 +23,29 @@ def _prepare_data(img, img_transform, cfg, device): return dict(img=[img], img_meta=[img_meta]) -def inference_detector(model, imgs, cfg, device='cuda:0'): +def _inference_single(model, img, img_transform, cfg, device): + img = mmcv.imread(img) + data = _prepare_data(img, img_transform, cfg, device) + with torch.no_grad(): + result = model(return_loss=False, rescale=True, **data) + return result + + +def _inference_generator(model, imgs, img_transform, cfg, device): + for img in imgs: + yield _inference_single(model, img, img_transform, cfg, device) - imgs = imgs if isinstance(imgs, list) else [imgs] + +def inference_detector(model, imgs, cfg, device='cuda:0'): img_transform = ImageTransform( size_divisor=cfg.data.test.size_divisor, **cfg.img_norm_cfg) model = model.to(device) model.eval() - for img in imgs: - img = mmcv.imread(img) - data = _prepare_data(img, img_transform, cfg, device) - with torch.no_grad(): - result = model(return_loss=False, rescale=True, **data) - yield result + + if not isinstance(imgs, list): + return _inference_single(model, imgs, img_transform, cfg, device) + else: + return _inference_generator(model, imgs, img_transform, cfg, device) def show_result(img, result, dataset='coco', score_thr=0.3): @@ -46,6 +56,7 @@ def show_result(img, result, dataset='coco', score_thr=0.3): ] labels = np.concatenate(labels) bboxes = np.vstack(result) + img = mmcv.imread(img) mmcv.imshow_det_bboxes( img.copy(), bboxes,