π [Merge] branch 'DATASET' into INFERENCE
Browse files
yolo/tools/data_augmentation.py
CHANGED
@@ -9,6 +9,7 @@ class AugmentationComposer:
|
|
9 |
|
10 |
def __init__(self, transforms, image_size: int = [640, 640]):
|
11 |
self.transforms = transforms
|
|
|
12 |
self.image_size = image_size[0]
|
13 |
self.pad_resize = PadAndResize(self.image_size)
|
14 |
|
@@ -38,10 +39,10 @@ class PadAndResize:
|
|
38 |
|
39 |
resized_img = square_img.resize((self.image_size, self.image_size))
|
40 |
|
41 |
-
boxes[:, 1] = (boxes[:, 1] + left) * scale
|
42 |
-
boxes[:, 2] = (boxes[:, 2] + top) * scale
|
43 |
-
boxes[:, 3] = (boxes[:, 3] + left) * scale
|
44 |
-
boxes[:, 4] = (boxes[:, 4] + top) * scale
|
45 |
|
46 |
return resized_img, boxes
|
47 |
|
|
|
9 |
|
10 |
def __init__(self, transforms, image_size: int = [640, 640]):
|
11 |
self.transforms = transforms
|
12 |
+
# TODO: handle List of image_size [640, 640]
|
13 |
self.image_size = image_size[0]
|
14 |
self.pad_resize = PadAndResize(self.image_size)
|
15 |
|
|
|
39 |
|
40 |
resized_img = square_img.resize((self.image_size, self.image_size))
|
41 |
|
42 |
+
boxes[:, 1] = (boxes[:, 1] * image.width + left) / self.image_size * scale
|
43 |
+
boxes[:, 2] = (boxes[:, 2] * image.height + top) / self.image_size * scale
|
44 |
+
boxes[:, 3] = (boxes[:, 3] * image.width + left) / self.image_size * scale
|
45 |
+
boxes[:, 4] = (boxes[:, 4] * image.height + top) / self.image_size * scale
|
46 |
|
47 |
return resized_img, boxes
|
48 |
|