Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_transforms_loading.py 4.18 KB
Newer Older
liukuikun's avatar
[Feature] Add Part1 of data transform (#1736)  
liukuikun 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
117
118
119
120
121
122 
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os.path as osp

import numpy as np

from mmcv.transforms import LoadAnnotation, LoadImageFromFile


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
        assert results['height'] == 300
        assert results['width'] == 400
        assert results['ori_height'] == 300
        assert results['ori_width'] == 400
        assert repr(transform) == transform.__class__.__name__ + \
            "(to_float32=False, color_type='color', " + \
            "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


class TestLoadAnnotation:

    def setup_class(cls):
        data_prefix = osp.join(osp.dirname(__file__), '../data')
        seg_map = osp.join(data_prefix, 'grayscale.jpg')
        cls.results = {
            'seg_map':
            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):
        transform = LoadAnnotation(
            with_bbox=True,
            with_label=False,
            with_seg=False,
            with_kps=False,
        )
        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()

    def test_load_labels(self):
        transform = LoadAnnotation(
            with_bbox=False,
            with_label=True,
            with_seg=False,
            with_kps=False,
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_bboxes_labels' in results
        assert (results['gt_bboxes_labels'] == np.array([1, 2])).all()

    def test_load_kps(self):
        transform = LoadAnnotation(
            with_bbox=False,
            with_label=False,
            with_seg=False,
            with_kps=True,
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_keypoints' in results
        assert (results['gt_keypoints'] == np.array([[[1, 2, 3]],
                                                     [[4, 5, 6]]])).all()

    def test_load_seg_map(self):
        transform = LoadAnnotation(
            with_bbox=False,
            with_label=False,
            with_seg=True,
            with_kps=False,
        )
        results = transform(copy.deepcopy(self.results))
        assert 'gt_semantic_seg' in results
        assert results['gt_semantic_seg'].shape[:2] == (300, 400)

    def test_repr(self):
        transform = LoadAnnotation(
            with_bbox=True,
            with_label=False,
            with_seg=False,
            with_kps=False,
        )
        assert repr(transform) == ('LoadAnnotation(with_bbox=True, '
                                   'with_label=False, with_seg=False, '
                                   "with_kps=False, imdecode_backend='cv2', "
                                   "file_client_args={'backend': 'disk'})")