forked from uoip/transforms
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_bounding_boxes.py
65 lines (55 loc) · 1.86 KB
/
test_bounding_boxes.py
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
from numpy.random import RandomState
# import imageio
def test_bboxes():
PRNG = RandomState()
PRNG2 = RandomState()
if args.seed > 0:
PRNG.seed(args.seed)
PRNG2.seed(args.seed)
transform = Compose([
[ColorJitter(prob=0.5)], # or write [ColorJitter(), None]
BoxesToCoords(),
HorizontalFlip(),
Expand((1, 4), prob=0.5),
ObjectRandomCrop(),
Resize(300),
CoordsToBoxes(),
#[SubtractMean(mean=VOC.MEAN)],
],
PRNG,
mode=None,
fillval=VOC.MEAN,
outside_points='clamp')
viz = Viz()
voc_dataset = VOCDetection(
root=args.root,
image_set=[('2007', 'trainval')],
keep_difficult=True,
transform=transform)
results = []
count = 0
i = PRNG2.choice(len(voc_dataset))
for _ in range(100):
img, boxes, labels = voc_dataset[i]
if len(labels) == 0:
continue
img = viz.draw_bbox(img, boxes, labels, True)
results.append(img)
cv2.imshow('0', img[:, :, ::-1])
c = cv2.waitKey(500)
if c == 27 or c == ord('q'): # ESC / 'q'
break
elif c == ord('c') or count >= 5:
count = 0
i = PRNG2.choice(len(voc_dataset))
count += 1
# imageio.mimsave('bboxes.gif', results, duration=0.5)
if __name__ == '__main__':
from transforms import *
from pascal_voc import VOC, VOCDetection, Viz
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, help='voc dataset root path', default='path/to/your/VOCdevkit')
parser.add_argument('--seed', type=int, help='random seed', default=0)
args = parser.parse_args()
test_bboxes()