henry000 commited on
Commit
860b0a5
Β·
2 Parent(s): 41ef509 92f4383

πŸ”€ [Merge] branch 'DATASET' into INFERENCE

Browse files
Files changed (1) hide show
  1. yolo/tools/data_augmentation.py +5 -4
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