Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_inference.py 4.4 KB
Newer Older
limm's avatar
add tests part  
limm committed
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 
# 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 mmpretrain.apis import (ImageClassificationInferencer, ModelHub,
                             get_model, inference_model)
from mmpretrain.models import MobileNetV3
from mmpretrain.structures import DataSample
from mmpretrain.visualization import UniversalVisualizer

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._config, inferencer.config)
        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, pretrained='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, DataSample)

    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, UniversalVisualizer)
            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.visualize_cls.assert_called_once_with(
                ANY,
                ANY,
                classes=inferencer.classes,
                resize=None,
                show=False,
                wait_time=0,
                rescale_factor=2.,
                draw_gt=False,
                draw_pred=True,
                draw_score=False,
                name='color',
                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'})