1from __future__ import print_function
2from __future__ import division
3
4import os.path as osp
5import mxnet as mx
6import numpy as np
7from mxnet import autograd, gluon
8import gluoncv as gcv
9from gluoncv.data import transforms
10from gluoncv.data import batchify
11from gluoncv.data.batchify import Tuple, Stack, Pad
12from gluoncv.data.transforms.presets import ssd
13from gluoncv.data.transforms.presets import rcnn
14from gluoncv.data.transforms.presets import yolo
15from gluoncv.data.transforms.presets import center_net
16from .tiny_datasets import COCODetectionTiny, COCOInstanceTiny, VOCDetectionTiny, VOCSegmentationTiny
17
18def test_bbox_crop():
19    bbox = np.array([[10, 20, 200, 500], [150, 200, 400, 300]])
20    np.testing.assert_allclose(transforms.bbox.crop(bbox, None), bbox)
21    np.testing.assert_allclose(
22        transforms.bbox.crop(bbox, (20, 30, 200, 200), allow_outside_center=True),
23        np.array([[  0,   0, 180, 200], [130, 170, 200, 200]]))
24    np.testing.assert_allclose(
25        transforms.bbox.crop(bbox, (20, 30, 200, 300), allow_outside_center=False),
26        np.array([[  0,   0, 180, 300]]))
27
28def test_bbox_flip():
29    bbox = np.array([[10, 20, 200, 500], [150, 200, 400, 300]])
30    size = (500, 1000)
31    np.testing.assert_allclose(
32        transforms.bbox.flip(bbox, size, False, False),
33        bbox)
34    np.testing.assert_allclose(
35        transforms.bbox.flip(bbox, size, False, True),
36        np.array([[ 10, 500, 200, 980], [150, 700, 400, 800]]))
37    np.testing.assert_allclose(
38        transforms.bbox.flip(bbox, size, True, False),
39        np.array([[300,  20, 490, 500], [100, 200, 350, 300]]))
40    np.testing.assert_allclose(
41        transforms.bbox.flip(bbox, size, True, True),
42        np.array([[300, 500, 490, 980], [100, 700, 350, 800]]))
43
44def test_bbox_resize():
45    bbox = np.array([[10, 20, 200, 500], [150, 200, 400, 300]], dtype=np.float32)
46    in_size = (600, 1000)
47    out_size = (200, 300)
48    np.testing.assert_allclose(
49        transforms.bbox.resize(bbox, in_size, out_size),
50        np.array([[  3.333333,   6.,  66.66667 , 150.], [ 50.,  60.000004, 133.33334 ,  90.]]),
51        rtol=1e-3)
52
53def test_bbox_translate():
54    bbox = np.array([[10, 20, 200, 500], [150, 200, 400, 300]], dtype=np.float32)
55    xoff = np.random.randint(-100, 100)
56    yoff = np.random.randint(-100, 100)
57    expected = bbox.copy()
58    expected[:, (0, 2)] += xoff
59    expected[:, (1, 3)] += yoff
60    np.testing.assert_allclose(transforms.bbox.translate(bbox, xoff, yoff), expected)
61
62def test_image_imresize():
63    image = mx.random.normal(shape=(240, 480, 3)).astype(np.uint8)
64    out = transforms.image.imresize(image, 300, 300)
65    np.testing.assert_allclose(out.shape, (300, 300, 3))
66
67def test_image_resize_long():
68    image = mx.random.normal(shape=(240, 480, 3)).astype(np.uint8)
69    out = transforms.image.resize_long(image, 300)
70    np.testing.assert_allclose(out.shape, (150, 300, 3))
71
72def test_image_random_pca():
73    image = mx.random.normal(shape=(240, 120, 3)).astype(np.float32)
74    out = transforms.image.random_pca_lighting(image, 0.1)
75    np.testing.assert_allclose(out.shape, image.shape)  # no value check
76
77def test_image_random_expand():
78    image = mx.random.normal(shape=(240, 120, 3)).astype(np.uint8)
79    # no expand when ratio <= 1
80    out, _ = transforms.image.random_expand(image, max_ratio=0.1, keep_ratio=True)
81    np.testing.assert_allclose(out.asnumpy(), image.asnumpy())
82    # check ratio
83    out, _ = transforms.image.random_expand(image, 4, keep_ratio=True)
84    np.testing.assert_allclose(out.shape[0] / out.shape[1], image.shape[0] / image.shape[1], rtol=1e-2, atol=1e-3)
85    # #
86    out, _ = transforms.image.random_expand(image, 4, keep_ratio=False)
87    np.testing.assert_((np.array(out.shape[:2]) - np.array(image.shape[:2]) + 1).all())
88
89def test_image_random_flip():
90    image = mx.random.normal(shape=(240, 120, 3)).astype(np.uint8)
91    # no flip
92    out, f = transforms.image.random_flip(image, 0, 0)
93    np.testing.assert_allclose(image.asnumpy(), out.asnumpy())
94    assert(f == (False, False))
95    # horizontal
96    out, f = transforms.image.random_flip(image, 1, 0)
97    np.testing.assert_allclose(image.asnumpy()[:, ::-1, :], out.asnumpy())
98    assert(f == (True, False))
99    # vertical
100    out, f = transforms.image.random_flip(image, 0, 1)
101    np.testing.assert_allclose(image.asnumpy()[::-1, :, :], out.asnumpy())
102    assert(f == (False, True))
103    # both
104    out, f = transforms.image.random_flip(image, 1, 1)
105    np.testing.assert_allclose(image.asnumpy()[::-1, ::-1, :], out.asnumpy())
106    assert(f == (True, True))
107
108def test_image_resize_contain():
109    image = mx.random.normal(shape=(240, 120, 3)).astype(np.uint8)
110    width = 123
111    height = 321
112    out, _ = transforms.image.resize_contain(image, (width, height))
113    np.testing.assert_allclose(out.shape, (height, width, 3))
114
115    width = 120
116    height = 20
117    out, _ = transforms.image.resize_contain(image, (width, height))
118    np.testing.assert_allclose(out.shape, (height, width, 3))
119
120def test_image_ten_crop():
121    image = mx.random.normal(shape=(240, 120, 3)).astype(np.uint8)
122    size = (24, 24)
123    crops = transforms.image.ten_crop(image, size)
124    assert len(crops) == 10
125    im = image.asnumpy()[:24, :24, :]
126    np.testing.assert_allclose(im, crops[1].asnumpy())
127    np.testing.assert_allclose(im[:, ::-1, :], crops[6].asnumpy())
128
129def test_experimental_bbox_random_crop_with_constraints():
130    bbox = np.array([[10, 20, 200, 500], [150, 200, 400, 300]])
131    size = (640, 480)
132    for _ in range(10):
133        min_scale = np.random.uniform(0, 0.9)
134        max_scale = np.random.uniform(min_scale, 1)
135        max_aspect_ratio = np.random.uniform(1, 3)
136        out, crop = transforms.experimental.bbox.random_crop_with_constraints(
137            bbox, size, min_scale=min_scale, max_scale=max_scale,
138            max_aspect_ratio=max_aspect_ratio, max_trial=20)
139    assert out.size >= 4
140
141def test_experimental_image_random_color_distort():
142    image = mx.random.normal(shape=(240, 120, 3)).astype(np.float32)
143    for _ in range(10):
144        brightness_delta = np.random.randint(0, 64)
145        contrast_low = np.random.uniform(0, 1)
146        contrast_high = np.random.uniform(1, 2)
147        saturation_low = np.random.uniform(0, 1)
148        saturation_high = np.random.uniform(1, 2)
149        hue_delta = np.random.randint(0, 36)
150        out = transforms.experimental.image.random_color_distort(
151            image, brightness_delta=brightness_delta, contrast_low=contrast_low,
152            contrast_high=contrast_high, saturation_low=saturation_low,
153            saturation_high=saturation_high, hue_delta=hue_delta)
154        np.testing.assert_allclose(out.shape, image.shape)
155
156def test_transforms_presets_ssd():
157    im_fname = gcv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
158                                  'gluoncv/detection/biking.jpg?raw=true', path='biking.jpg')
159    x, orig_img = ssd.load_test(im_fname, short=512)
160    x1, orig_img1 = ssd.transform_test(mx.image.imread(im_fname), short=512)
161    np.testing.assert_allclose(x.asnumpy(), x1.asnumpy())
162    np.testing.assert_allclose(orig_img, orig_img1)
163    if not osp.isdir(osp.expanduser('~/.mxnet/datasets/voc')):
164        return
165    train_dataset = VOCDetectionTiny()
166    val_dataset = VOCDetectionTiny(splits=[('tiny_motorbike', 'test')])
167    width, height = (512, 512)
168    net = gcv.model_zoo.get_model('ssd_512_resnet50_v1_voc', pretrained=False, pretrained_base=False)
169    net.initialize()
170    num_workers = 0
171    batch_size = 4
172    with autograd.train_mode():
173        _, _, anchors = net(mx.nd.zeros((1, 3, height, width)))
174    batchify_fn = Tuple(Stack(), Stack(), Stack())  # stack image, cls_targets, box_targets
175    train_loader = gluon.data.DataLoader(
176        train_dataset.transform(ssd.SSDDefaultTrainTransform(width, height, anchors)),
177        batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=num_workers)
178    val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
179    val_loader = gluon.data.DataLoader(
180        val_dataset.transform(ssd.SSDDefaultValTransform(width, height)),
181        batch_size, False, batchify_fn=val_batchify_fn, last_batch='keep', num_workers=num_workers)
182    train_loader2 = gluon.data.DataLoader(
183        train_dataset.transform(ssd.SSDDefaultTrainTransform(width, height)),
184        batch_size, True, batchify_fn=val_batchify_fn, last_batch='rollover', num_workers=num_workers)
185
186    for loader in [train_loader, val_loader, train_loader2]:
187        for i, batch in enumerate(loader):
188            if i > 1:
189                break
190            pass
191
192def test_transforms_presets_rcnn():
193    im_fname = gcv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
194                                  'gluoncv/detection/biking.jpg?raw=true', path='biking.jpg')
195    x, orig_img = rcnn.load_test(im_fname, short=600, max_size=1000)
196    x1, orig_img1 = rcnn.transform_test(mx.image.imread(im_fname), short=600, max_size=1000)
197    np.testing.assert_allclose(x.asnumpy(), x1.asnumpy())
198    np.testing.assert_allclose(orig_img, orig_img1)
199    if not osp.isdir(osp.expanduser('~/.mxnet/datasets/voc')):
200        return
201    train_dataset = VOCDetectionTiny()
202    val_dataset = VOCDetectionTiny(splits=[('tiny_motorbike', 'test')])
203    width, height = (512, 512)
204    net = gcv.model_zoo.get_model('faster_rcnn_resnet50_v1b_voc', pretrained=False, pretrained_base=False)
205    net.initialize()
206    num_workers = 0
207    short, max_size = 600, 1000
208    batch_size = 4
209    train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
210    train_loader = mx.gluon.data.DataLoader(
211        train_dataset.transform(rcnn.FasterRCNNDefaultTrainTransform(short, max_size, net, flip_p=0.5)),
212        batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers)
213    val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
214    val_loader = mx.gluon.data.DataLoader(
215        val_dataset.transform(rcnn.FasterRCNNDefaultValTransform(short, max_size)),
216        batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
217    train_loader2 = gluon.data.DataLoader(
218        train_dataset.transform(rcnn.FasterRCNNDefaultTrainTransform(short, max_size)),
219        batch_size, True, batchify_fn=batchify.Tuple(*[batchify.Append() for _ in range(2)]),
220        last_batch='rollover', num_workers=num_workers)
221
222    for loader in [train_loader, val_loader, train_loader2]:
223        for i, batch in enumerate(loader):
224            if i > 1:
225                break
226            pass
227
228def test_transforms_presets_mask_rcnn():
229    # use valid only, loading training split is very slow
230    train_dataset = COCOInstanceTiny(root=osp.join('~', '.mxnet', 'datasets', 'tiny_coco'),
231                                          splits=('instances_val2017_tiny',), skip_empty=True)
232    val_dataset = COCOInstanceTiny(root=osp.join('~', '.mxnet', 'datasets', 'tiny_coco'),
233                                        splits=('instances_val2017_tiny',))
234    net = gcv.model_zoo.get_model('mask_rcnn_resnet50_v1b_coco', pretrained=False, pretrained_base=False)
235    net.initialize()
236    num_workers = 0
237    short, max_size = 800, 1333
238    batch_size = 8
239    train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(6)])
240    train_loader = mx.gluon.data.DataLoader(
241        train_dataset.transform(rcnn.MaskRCNNDefaultTrainTransform(short, max_size, net)),
242        batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers)
243    val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
244    val_loader = mx.gluon.data.DataLoader(
245        val_dataset.transform(rcnn.MaskRCNNDefaultValTransform(short, max_size)),
246        batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
247
248    for loader in [train_loader, val_loader]:
249        for i, batch in enumerate(loader):
250            if i > 1:
251                break
252            pass
253
254def test_transforms_presets_yolo():
255    im_fname = gcv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
256                                  'gluoncv/detection/biking.jpg?raw=true', path='biking.jpg')
257    x, orig_img = yolo.load_test(im_fname, short=512)
258    x1, orig_img1 = yolo.transform_test(mx.image.imread(im_fname), short=512)
259    np.testing.assert_allclose(x.asnumpy(), x1.asnumpy())
260    np.testing.assert_allclose(orig_img, orig_img1)
261    if not osp.isdir(osp.expanduser('~/.mxnet/datasets/voc')):
262        return
263    train_dataset = VOCDetectionTiny()
264    val_dataset = VOCDetectionTiny(splits=[('tiny_motorbike', 'test')])
265    width, height = (512, 512)
266    net = gcv.model_zoo.get_model('yolo3_darknet53_voc', pretrained=False, pretrained_base=False)
267    net.initialize()
268    num_workers = 0
269    batch_size = 4
270    batchify_fn = Tuple(*([Stack() for _ in range(6)] + [Pad(axis=0, pad_val=-1) for _ in range(1)]))
271    train_loader = gluon.data.DataLoader(
272        train_dataset.transform(yolo.YOLO3DefaultTrainTransform(width, height, net)),
273        batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=num_workers)
274    val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
275    val_loader = gluon.data.DataLoader(
276        val_dataset.transform(yolo.YOLO3DefaultValTransform(width, height)),
277        batch_size, False, batchify_fn=val_batchify_fn, last_batch='keep', num_workers=num_workers)
278    train_loader2 = gluon.data.DataLoader(
279        train_dataset.transform(yolo.YOLO3DefaultTrainTransform(width, height)),
280        batch_size, True, batchify_fn=val_batchify_fn, last_batch='rollover', num_workers=num_workers)
281
282    for loader in [train_loader, val_loader, train_loader2]:
283        for i, batch in enumerate(loader):
284            if i > 1:
285                break
286            pass
287
288def test_transforms_presets_center_net():
289    im_fname = gcv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
290                                  'gluoncv/detection/biking.jpg?raw=true', path='biking.jpg')
291    x, orig_img = center_net.load_test(im_fname, short=512)
292    x1, orig_img1 = center_net.transform_test(mx.image.imread(im_fname), short=512)
293    np.testing.assert_allclose(x.asnumpy(), x1.asnumpy())
294    np.testing.assert_allclose(orig_img, orig_img1)
295    if not osp.isdir(osp.expanduser('~/.mxnet/datasets/voc')):
296        return
297    train_dataset = VOCDetectionTiny()
298    val_dataset = VOCDetectionTiny(splits=[('tiny_motorbike', 'test')])
299    width, height = (512, 512)
300    net = gcv.model_zoo.get_model('center_net_resnet18_v1b_voc', pretrained=False, pretrained_base=False)
301    net.initialize()
302    num_workers = 0
303    batch_size = 4
304    batchify_fn = Tuple([Stack() for _ in range(6)])
305    train_loader = gluon.data.DataLoader(
306        train_dataset.transform(center_net.CenterNetDefaultTrainTransform(width, height, num_class=len(train_dataset.classes), scale_factor=net.scale)),
307        batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=num_workers)
308    val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
309    val_loader = gluon.data.DataLoader(
310        val_dataset.transform(center_net.CenterNetDefaultValTransform(width, height)),
311        batch_size, False, batchify_fn=val_batchify_fn, last_batch='keep', num_workers=num_workers)
312
313    for loader in [train_loader, val_loader]:
314        for i, batch in enumerate(loader):
315            if i > 1:
316                break
317            pass
318
319if __name__ == '__main__':
320    import nose
321    nose.runmodule()
322