114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
|
# Copyright (c) OpenMMLab. All rights reserved.
|
||
|
|
import os.path as osp
|
||
|
|
from tempfile import TemporaryDirectory
|
||
|
|
from unittest import TestCase
|
||
|
|
from unittest.mock import ANY, MagicMock, patch
|
||
|
|
|
||
|
|
from mmcv.image import imread
|
||
|
|
|
||
|
|
from mmcls.apis import (ImageClassificationInferencer, ModelHub, get_model,
|
||
|
|
inference_model)
|
||
|
|
from mmcls.models import MobileNetV3
|
||
|
|
from mmcls.structures import ClsDataSample
|
||
|
|
from mmcls.visualization import ClsVisualizer
|
||
|
|
|
||
|
|
MODEL = 'mobilenet-v3-small-050_3rdparty_in1k'
|
||
|
|
WEIGHT = 'https://download.openmmlab.com/mmclassification/v0/mobilenet_v3/mobilenet-v3-small-050_3rdparty_in1k_20221114-e0b86be1.pth' # noqa: E501
|
||
|
|
CONFIG = ModelHub.get(MODEL).config
|
||
|
|
|
||
|
|
|
||
|
|
class TestImageClassificationInferencer(TestCase):
|
||
|
|
|
||
|
|
def test_init(self):
|
||
|
|
# test input BaseModel
|
||
|
|
model = get_model(MODEL)
|
||
|
|
inferencer = ImageClassificationInferencer(model)
|
||
|
|
self.assertEqual(model.cfg, inferencer.cfg)
|
||
|
|
self.assertIsInstance(inferencer.model.backbone, MobileNetV3)
|
||
|
|
|
||
|
|
# test input model name
|
||
|
|
with patch('mmengine.runner.load_checkpoint') as mock:
|
||
|
|
inferencer = ImageClassificationInferencer(MODEL)
|
||
|
|
self.assertIsInstance(inferencer.model.backbone, MobileNetV3)
|
||
|
|
mock.assert_called_once_with(ANY, WEIGHT, map_location='cpu')
|
||
|
|
|
||
|
|
# test input config path
|
||
|
|
inferencer = ImageClassificationInferencer(CONFIG.filename)
|
||
|
|
self.assertIsInstance(inferencer.model.backbone, MobileNetV3)
|
||
|
|
|
||
|
|
# test input config object
|
||
|
|
inferencer = ImageClassificationInferencer(CONFIG)
|
||
|
|
self.assertIsInstance(inferencer.model.backbone, MobileNetV3)
|
||
|
|
|
||
|
|
# test specify weights
|
||
|
|
with patch('mmengine.runner.load_checkpoint') as mock:
|
||
|
|
ImageClassificationInferencer(MODEL, weights='custom.pth')
|
||
|
|
mock.assert_called_once_with(ANY, 'custom.pth', map_location='cpu')
|
||
|
|
|
||
|
|
def test_call(self):
|
||
|
|
img_path = osp.join(osp.dirname(__file__), '../data/color.jpg')
|
||
|
|
img = imread(img_path)
|
||
|
|
|
||
|
|
# test inference classification model
|
||
|
|
inferencer = ImageClassificationInferencer(MODEL)
|
||
|
|
results = inferencer(img_path)[0]
|
||
|
|
self.assertEqual(
|
||
|
|
results.keys(),
|
||
|
|
{'pred_score', 'pred_scores', 'pred_label', 'pred_class'})
|
||
|
|
|
||
|
|
# test return_datasample=True
|
||
|
|
results = inferencer(img, return_datasamples=True)[0]
|
||
|
|
self.assertIsInstance(results, ClsDataSample)
|
||
|
|
|
||
|
|
def test_visualize(self):
|
||
|
|
img_path = osp.join(osp.dirname(__file__), '../data/color.jpg')
|
||
|
|
img = imread(img_path)
|
||
|
|
|
||
|
|
inferencer = ImageClassificationInferencer(MODEL)
|
||
|
|
self.assertIsNone(inferencer.visualizer)
|
||
|
|
|
||
|
|
with TemporaryDirectory() as tmpdir:
|
||
|
|
inferencer(img, show_dir=tmpdir)
|
||
|
|
self.assertIsInstance(inferencer.visualizer, ClsVisualizer)
|
||
|
|
self.assertTrue(osp.exists(osp.join(tmpdir, '0.png')))
|
||
|
|
|
||
|
|
inferencer.visualizer = MagicMock(wraps=inferencer.visualizer)
|
||
|
|
inferencer(
|
||
|
|
img_path, rescale_factor=2., draw_score=False, show_dir=tmpdir)
|
||
|
|
self.assertTrue(osp.exists(osp.join(tmpdir, 'color.png')))
|
||
|
|
inferencer.visualizer.add_datasample.assert_called_once_with(
|
||
|
|
'color',
|
||
|
|
ANY,
|
||
|
|
ANY,
|
||
|
|
show=False,
|
||
|
|
rescale_factor=2.,
|
||
|
|
draw_gt=False,
|
||
|
|
draw_pred=True,
|
||
|
|
draw_score=False,
|
||
|
|
out_file=osp.join(tmpdir, 'color.png'))
|
||
|
|
|
||
|
|
|
||
|
|
class TestInferenceAPIs(TestCase):
|
||
|
|
|
||
|
|
def test_inference_model(self):
|
||
|
|
# test backward compatibility
|
||
|
|
img_path = osp.join(osp.dirname(__file__), '../data/color.jpg')
|
||
|
|
img = imread(img_path)
|
||
|
|
|
||
|
|
model = get_model(MODEL, pretrained=True)
|
||
|
|
results = inference_model(model, img_path)
|
||
|
|
self.assertEqual(
|
||
|
|
results.keys(),
|
||
|
|
{'pred_score', 'pred_scores', 'pred_label', 'pred_class'})
|
||
|
|
|
||
|
|
results = inference_model(model, img)
|
||
|
|
self.assertEqual(
|
||
|
|
results.keys(),
|
||
|
|
{'pred_score', 'pred_scores', 'pred_label', 'pred_class'})
|
||
|
|
|
||
|
|
# test input model name
|
||
|
|
results = inference_model(MODEL, img)
|
||
|
|
self.assertEqual(
|
||
|
|
results.keys(),
|
||
|
|
{'pred_score', 'pred_scores', 'pred_label', 'pred_class'})
|