add tests part

tests/__init__.py

+# Copyright (c) OpenMMLab. All rights reserved.

tests/data/dataset/a/1.JPG

+../../color.jpg
\ No newline at end of file

tests/data/dataset/ann.json

+{
+  "metainfo": {
+    "categories": [
+      {
+        "category_name": "first",
+        "id": 0
+      },
+      {
+        "category_name": "second",
+        "id": 1
+      }
+    ]
+  },
+  "data_list": [
+    {
+      "img_path": "a/1.JPG",
+      "gt_label": 0
+    },
+    {
+      "img_path": "b/2.jpeg",
+      "gt_label": 1
+    },
+    {
+      "img_path": "b/subb/2.jpeg",
+      "gt_label": 1
+    }
+  ]
+}

tests/data/dataset/ann.txt

+a/1.JPG 0
+b/2.jpeg 1
+b/subb/3.jpg 1

tests/data/dataset/ann_without_labels.txt

+a/1.JPG
+b/2.jpeg
+b/subb/3.jpg

tests/data/dataset/b/2.jpeg

+../../color.jpg
\ No newline at end of file

tests/data/dataset/b/subb/3.jpg

+../../../color.jpg
\ No newline at end of file

tests/data/dataset/classes.txt

+bus
+car

tests/data/dataset/multi-task.json

+{
+  "metainfo": {
+    "tasks": [
+      "gender",
+      "wear"
+    ]
+  },
+  "data_list": [
+    {
+      "img_path": "a/1.JPG",
+      "gt_label": {
+        "gender": 0
+      }
+    },
+    {
+      "img_path": "b/2.jpeg",
+      "gt_label": {
+        "gender": 0,
+        "wear": [
+          1,
+          0,
+          1,
+          0
+        ]
+      }
+    },
+    {
+      "img_path": "b/subb/3.jpg",
+      "gt_label": {
+        "gender": 1,
+        "wear": [
+          0,
+          1,
+          0,
+          1
+        ]
+      }
+    }
+  ]
+}

tests/data/dataset/multi_label_ann.json

+{
+  "metainfo": {
+    "categories": [
+      {
+        "category_name": "first",
+        "id": 0
+      },
+      {
+        "category_name": "second",
+        "id": 1
+      }
+    ]
+  },
+  "data_list": [
+    {
+      "img_path": "a/1.JPG",
+      "gt_label": [0]
+    },
+    {
+      "img_path": "b/2.jpeg",
+      "gt_label": [1]
+    },
+    {
+      "img_path": "b/subb/2.jpeg",
+      "gt_label": [0, 1]
+    }
+  ]
+}

tests/data/meta.yml

+Models:
+  - Name: test_model
+    Metadata:
+      FLOPs: 319000000
+      Parameters: 3500000
+    Results:
+      - Dataset: ImageNet-1k
+        Metrics:
+          Top 1 Accuracy: 71.86
+          Top 5 Accuracy: 90.42
+        Task: Image Classification
+    Weights: test_weight.pth
+    Config: test_config.py

tests/data/retinanet.py

+# Copyright (c) OpenMMLab. All rights reserved.
+# small RetinaNet
+num_classes = 3
+
+# model settings
+model = dict(
+    type='RetinaNet',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        start_level=1,
+        add_extra_convs='on_input',
+        num_outs=5),
+    bbox_head=dict(
+        type='RetinaHead',
+        num_classes=num_classes,
+        in_channels=256,
+        stacked_convs=1,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            octave_base_scale=4,
+            scales_per_octave=3,
+            ratios=[0.5, 1.0, 2.0],
+            strides=[8, 16, 32, 64, 128]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='FocalLoss',
+            use_sigmoid=True,
+            gamma=2.0,
+            alpha=0.25,
+            loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    # model training and testing settings
+    train_cfg=dict(
+        assigner=dict(
+            type='MaxIoUAssigner',
+            pos_iou_thr=0.5,
+            neg_iou_thr=0.4,
+            min_pos_iou=0,
+            ignore_iof_thr=-1),
+        allowed_border=-1,
+        pos_weight=-1,
+        debug=False),
+    test_cfg=dict(
+        nms_pre=1000,
+        min_bbox_size=0,
+        score_thr=0.05,
+        nms=dict(type='nms', iou_threshold=0.5),
+        max_per_img=100))
+
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(test=dict(pipeline=test_pipeline))

tests/data/vis_data.json

+{"lr": 0.1, "data_time": 0.0061125516891479496, "loss": 2.6531384229660033, "time": 0.14429793357849122, "epoch": 1, "step": 10}
+{"lr": 0.1, "data_time": 0.00030262470245361327, "loss": 2.9456406116485594, "time": 0.0219132661819458, "epoch": 1, "step": 20}
+{"lr": 0.1, "data_time": 0.00022499561309814454, "loss": 3.1025198698043823, "time": 0.021793675422668458, "epoch": 1, "step": 30}
+{"lr": 0.1, "data_time": 0.00023109912872314452, "loss": 2.5765398740768433, "time": 0.021819329261779784, "epoch": 1, "step": 40}
+{"lr": 0.1, "data_time": 0.00023169517517089843, "loss": 2.671005058288574, "time": 0.02181088924407959, "epoch": 1, "step": 50}
+{"lr": 0.1, "data_time": 0.00021798610687255858, "loss": 2.5273321866989136, "time": 0.021781444549560547, "epoch": 1, "step": 60}
+{"accuracy/top1": 18.80000114440918, "step": 1}
+{"lr": 0.1, "data_time": 0.0007575273513793946, "loss": 2.3254310727119445, "time": 0.02237672805786133, "epoch": 2, "step": 73}
+{"lr": 0.1, "data_time": 0.0002459049224853516, "loss": 2.194095492362976, "time": 0.021792054176330566, "epoch": 2, "step": 83}
+{"lr": 0.1, "data_time": 0.00027666091918945315, "loss": 2.207821953296661, "time": 0.021822547912597655, "epoch": 2, "step": 93}
+{"lr": 0.1, "data_time": 0.00025298595428466795, "loss": 2.090667963027954, "time": 0.02178535461425781, "epoch": 2, "step": 103}
+{"lr": 0.1, "data_time": 0.0002483367919921875, "loss": 2.18342148065567, "time": 0.021893739700317383, "epoch": 2, "step": 113}
+{"lr": 0.1, "data_time": 0.00030078887939453123, "loss": 2.2274346113204957, "time": 0.022345948219299316, "epoch": 2, "step": 123}
+{"accuracy/top1": 21.100000381469727, "step": 2}
+{"lr": 0.1, "data_time": 0.0008128643035888672, "loss": 2.017984461784363, "time": 0.02267434597015381, "epoch": 3, "step": 136}
+{"lr": 0.1, "data_time": 0.00023736953735351563, "loss": 2.0648953437805178, "time": 0.02174344062805176, "epoch": 3, "step": 146}
+{"lr": 0.1, "data_time": 0.00024063587188720702, "loss": 2.0859395623207093, "time": 0.022107195854187012, "epoch": 3, "step": 156}
+{"lr": 0.1, "data_time": 0.0002336740493774414, "loss": 2.1662048220634462, "time": 0.021825361251831054, "epoch": 3, "step": 166}
+{"lr": 0.1, "data_time": 0.0002296924591064453, "loss": 2.1007142066955566, "time": 0.021821355819702147, "epoch": 3, "step": 176}
+{"lr": 0.1, "data_time": 0.00023157596588134765, "loss": 2.0436240792274476, "time": 0.021722936630249025, "epoch": 3, "step": 186}
+{"accuracy/top1": 25.600000381469727, "step": 3}

tests/test_apis/test_inference.py

+# 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'})

tests/test_apis/test_model.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+from unittest import TestCase
+from unittest.mock import patch
+
+from mmengine import Config
+
+from mmpretrain.apis import ModelHub, get_model, init_model, list_models
+from mmpretrain.models import ImageClassifier, MobileNetV2
+
+
+class TestModelHub(TestCase):
+
+    def test_mmpretrain_models(self):
+        self.assertIn('resnet18_8xb32_in1k', ModelHub._models_dict)
+
+    def test_register_model_index(self):
+        model_index_path = osp.join(osp.dirname(__file__), '../data/meta.yml')
+
+        ModelHub.register_model_index(model_index_path)
+        self.assertIn('test_model', ModelHub._models_dict)
+        self.assertEqual(
+            ModelHub._models_dict['test_model'].config,
+            osp.abspath(
+                osp.join(osp.dirname(model_index_path), 'test_config.py')))
+
+        with self.assertRaisesRegex(ValueError, 'meta.yml'):
+            # test name conflict
+            ModelHub.register_model_index(model_index_path)
+
+        # test specify config prefix
+        del ModelHub._models_dict['test_model']
+        ModelHub.register_model_index(
+            model_index_path, config_prefix='configs')
+        self.assertEqual(ModelHub._models_dict['test_model'].config,
+                         osp.abspath(osp.join('configs', 'test_config.py')))
+
+    def test_get_model(self):
+        metainfo = ModelHub.get('resnet18_8xb32_in1k')
+        self.assertIsInstance(metainfo.weights, str)
+        self.assertIsInstance(metainfo.config, Config)
+
+
+class TestHubAPIs(TestCase):
+
+    def test_list_models(self):
+        models_names = list_models()
+        self.assertIsInstance(models_names, list)
+
+        models_names = list_models(pattern='swin*in1k')
+        for model_name in models_names:
+            self.assertTrue(
+                model_name.startswith('swin') and 'in1k' in model_name)
+
+    def test_get_model(self):
+        model = get_model('mobilenet-v2_8xb32_in1k')
+        self.assertIsInstance(model, ImageClassifier)
+        self.assertIsInstance(model.backbone, MobileNetV2)
+
+        with patch('mmengine.runner.load_checkpoint') as mock:
+            model = get_model('mobilenet-v2_8xb32_in1k', pretrained=True)
+            model = get_model('mobilenet-v2_8xb32_in1k', pretrained='test.pth')
+
+            weight = mock.call_args_list[0][0][1]
+            self.assertIn('https', weight)
+            weight = mock.call_args_list[1][0][1]
+            self.assertEqual('test.pth', weight)
+
+        with self.assertRaisesRegex(ValueError, 'Failed to find'):
+            get_model('unknown-model')
+
+    def test_init_model(self):
+        # test init from config object
+        cfg = ModelHub.get('mobilenet-v2_8xb32_in1k').config
+        model = init_model(cfg)
+        self.assertIsInstance(model, ImageClassifier)
+        self.assertIsInstance(model.backbone, MobileNetV2)
+
+        # test init from config file
+        cfg = ModelHub._models_dict['mobilenet-v2_8xb32_in1k'].config
+        self.assertIsInstance(cfg, str)
+        model = init_model(cfg)
+        self.assertIsInstance(model, ImageClassifier)
+        self.assertIsInstance(model.backbone, MobileNetV2)
+
+        # test modify configs of the model
+        model = init_model(cfg, head=dict(num_classes=10))
+        self.assertEqual(model.head.num_classes, 10)

tests/test_datasets/test_dataset_utils.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+import random
+import string
+from unittest.mock import patch
+
+import pytest
+
+from mmpretrain.datasets.utils import (check_integrity,
+                                       open_maybe_compressed_file, rm_suffix)
+
+
+def test_dataset_utils():
+    # test rm_suffix
+    assert rm_suffix('a.jpg') == 'a'
+    assert rm_suffix('a.bak.jpg') == 'a.bak'
+    assert rm_suffix('a.bak.jpg', suffix='.jpg') == 'a.bak'
+    assert rm_suffix('a.bak.jpg', suffix='.bak.jpg') == 'a'
+
+    # test check_integrity
+    rand_file = ''.join(random.sample(string.ascii_letters, 10))
+    assert not check_integrity(rand_file, md5=None)
+    assert not check_integrity(rand_file, md5=2333)
+    test_file = osp.join(osp.dirname(__file__), '../data/color.jpg')
+    assert check_integrity(test_file, md5='08252e5100cb321fe74e0e12a724ce14')
+    assert not check_integrity(test_file, md5=2333)
+
+
+@pytest.mark.parametrize('method,path', [('gzip.open', 'abc.gz'),
+                                         ('lzma.open', 'abc.xz'),
+                                         ('builtins.open', 'abc.txt'),
+                                         (None, 1)])
+def test_open_maybe_compressed_file(method, path):
+    if method:
+        with patch(method) as mock:
+            open_maybe_compressed_file(path)
+            mock.assert_called()
+    else:
+        assert open_maybe_compressed_file(path) == path