test_transforms_loading.py 4.75 KB
Newer Older
1
2
3
4
5
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os.path as osp

import numpy as np
liukuikun's avatar
liukuikun committed
6
import pytest
7

8
from mmcv.transforms import LoadAnnotations, LoadImageFromFile
9
10
11
12
13
14
15
16
17
18
19
20
21


class TestLoadImageFromFile:

    def test_load_img(self):
        data_prefix = osp.join(osp.dirname(__file__), '../data')

        results = dict(img_path=osp.join(data_prefix, 'color.jpg'))
        transform = LoadImageFromFile()
        results = transform(copy.deepcopy(results))
        assert results['img_path'] == osp.join(data_prefix, 'color.jpg')
        assert results['img'].shape == (300, 400, 3)
        assert results['img'].dtype == np.uint8
22
23
        assert results['img_shape'] == (300, 400)
        assert results['ori_shape'] == (300, 400)
24
        assert repr(transform) == transform.__class__.__name__ + \
liukuikun's avatar
liukuikun committed
25
            "(ignore_empty=False, to_float32=False, color_type='color', " + \
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
            "imdecode_backend='cv2', file_client_args={'backend': 'disk'})"

        # to_float32
        transform = LoadImageFromFile(to_float32=True)
        results = transform(copy.deepcopy(results))
        assert results['img'].dtype == np.float32

        # gray image
        results = dict(img_path=osp.join(data_prefix, 'grayscale.jpg'))
        transform = LoadImageFromFile()
        results = transform(copy.deepcopy(results))
        assert results['img'].shape == (300, 400, 3)
        assert results['img'].dtype == np.uint8

        transform = LoadImageFromFile(color_type='unchanged')
        results = transform(copy.deepcopy(results))
        assert results['img'].shape == (300, 400)
        assert results['img'].dtype == np.uint8

liukuikun's avatar
liukuikun committed
45
46
47
48
49
50
51
52
53
        # test load empty
        fake_img_path = osp.join(data_prefix, 'fake.jpg')
        results['img_path'] = fake_img_path
        transform = LoadImageFromFile(ignore_empty=False)
        with pytest.raises(FileNotFoundError):
            transform(copy.deepcopy(results))
        transform = LoadImageFromFile(ignore_empty=True)
        assert transform(copy.deepcopy(results)) is None

54

55
class TestLoadAnnotations:
56
57
58
59
60

    def setup_class(cls):
        data_prefix = osp.join(osp.dirname(__file__), '../data')
        seg_map = osp.join(data_prefix, 'grayscale.jpg')
        cls.results = {
61
            'seg_map_path':
62
63
64
65
66
67
68
69
70
71
72
73
74
            seg_map,
            'instances': [{
                'bbox': [0, 0, 10, 20],
                'bbox_label': 1,
                'keypoints': [1, 2, 3]
            }, {
                'bbox': [10, 10, 110, 120],
                'bbox_label': 2,
                'keypoints': [4, 5, 6]
            }]
        }

    def test_load_bboxes(self):
75
        transform = LoadAnnotations(
76
77
78
            with_bbox=True,
            with_label=False,
            with_seg=False,
79
            with_keypoints=False,
80
81
82
83
84
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_bboxes' in results
        assert (results['gt_bboxes'] == np.array([[0, 0, 10, 20],
                                                  [10, 10, 110, 120]])).all()
liukuikun's avatar
liukuikun committed
85
        assert results['gt_bboxes'].dtype == np.float32
86
87

    def test_load_labels(self):
88
        transform = LoadAnnotations(
89
90
91
            with_bbox=False,
            with_label=True,
            with_seg=False,
92
            with_keypoints=False,
93
94
95
96
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_bboxes_labels' in results
        assert (results['gt_bboxes_labels'] == np.array([1, 2])).all()
97
        assert results['gt_bboxes_labels'].dtype == np.int64
98
99

    def test_load_kps(self):
100
        transform = LoadAnnotations(
101
102
103
            with_bbox=False,
            with_label=False,
            with_seg=False,
104
            with_keypoints=True,
105
106
107
108
109
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_keypoints' in results
        assert (results['gt_keypoints'] == np.array([[[1, 2, 3]],
                                                     [[4, 5, 6]]])).all()
liukuikun's avatar
liukuikun committed
110
        assert results['gt_keypoints'].dtype == np.float32
111
112

    def test_load_seg_map(self):
113
        transform = LoadAnnotations(
114
115
116
            with_bbox=False,
            with_label=False,
            with_seg=True,
117
            with_keypoints=False,
118
119
        )
        results = transform(copy.deepcopy(self.results))
120
121
        assert 'gt_seg_map' in results
        assert results['gt_seg_map'].shape[:2] == (300, 400)
liukuikun's avatar
liukuikun committed
122
        assert results['gt_seg_map'].dtype == np.uint8
123
124

    def test_repr(self):
125
        transform = LoadAnnotations(
126
127
128
            with_bbox=True,
            with_label=False,
            with_seg=False,
129
            with_keypoints=False,
130
        )
131
        assert repr(transform) == (
132
            'LoadAnnotations(with_bbox=True, '
133
134
135
            'with_label=False, with_seg=False, '
            "with_keypoints=False, imdecode_backend='cv2', "
            "file_client_args={'backend': 'disk'})")