-
Notifications
You must be signed in to change notification settings - Fork 2.6k
/
transforms.py
119 lines (89 loc) · 2.93 KB
/
transforms.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
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
import torch
import torch.nn.functional as F
import numpy as np
import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
from .utils import xywh2xyxy_np
import torchvision.transforms as transforms
class ImgAug(object):
def __init__(self, augmentations=[]):
self.augmentations = augmentations
def __call__(self, data):
# Unpack data
img, boxes = data
# Convert xywh to xyxy
boxes = np.array(boxes)
boxes[:, 1:] = xywh2xyxy_np(boxes[:, 1:])
# Convert bounding boxes to imgaug
bounding_boxes = BoundingBoxesOnImage(
[BoundingBox(*box[1:], label=box[0]) for box in boxes],
shape=img.shape)
# Apply augmentations
img, bounding_boxes = self.augmentations(
image=img,
bounding_boxes=bounding_boxes)
# Clip out of image boxes
bounding_boxes = bounding_boxes.clip_out_of_image()
# Convert bounding boxes back to numpy
boxes = np.zeros((len(bounding_boxes), 5))
for box_idx, box in enumerate(bounding_boxes):
# Extract coordinates for unpadded + unscaled image
x1 = box.x1
y1 = box.y1
x2 = box.x2
y2 = box.y2
# Returns (x, y, w, h)
boxes[box_idx, 0] = box.label
boxes[box_idx, 1] = ((x1 + x2) / 2)
boxes[box_idx, 2] = ((y1 + y2) / 2)
boxes[box_idx, 3] = (x2 - x1)
boxes[box_idx, 4] = (y2 - y1)
return img, boxes
class RelativeLabels(object):
def __init__(self, ):
pass
def __call__(self, data):
img, boxes = data
w, h, _ = img.shape
boxes[:, [1, 3]] /= h
boxes[:, [2, 4]] /= w
return img, boxes
class AbsoluteLabels(object):
def __init__(self, ):
pass
def __call__(self, data):
img, boxes = data
w, h, _ = img.shape
boxes[:, [1, 3]] *= h
boxes[:, [2, 4]] *= w
return img, boxes
class PadSquare(ImgAug):
def __init__(self, ):
self.augmentations = iaa.Sequential([
iaa.PadToAspectRatio(
1.0,
position="center-center").to_deterministic()
])
class ToTensor(object):
def __init__(self, ):
pass
def __call__(self, data):
img, boxes = data
# Extract image as PyTorch tensor
img = transforms.ToTensor()(img)
bb_targets = torch.zeros((len(boxes), 6))
bb_targets[:, 1:] = transforms.ToTensor()(boxes)
return img, bb_targets
class Resize(object):
def __init__(self, size):
self.size = size
def __call__(self, data):
img, boxes = data
img = F.interpolate(img.unsqueeze(0), size=self.size, mode="nearest").squeeze(0)
return img, boxes
DEFAULT_TRANSFORMS = transforms.Compose([
AbsoluteLabels(),
PadSquare(),
RelativeLabels(),
ToTensor(),
])