henry000 commited on
Commit
c627401
·
1 Parent(s): 0bd2faa

✨ [Add] a RandomCrop transform, align paper

Browse files
Files changed (1) hide show
  1. yolo/tools/data_augmentation.py +31 -0
yolo/tools/data_augmentation.py CHANGED
@@ -155,3 +155,34 @@ class MixUp:
155
  mixed_boxes = torch.cat([lam * boxes, (1 - lam) * boxes2])
156
 
157
  return TF.to_pil_image(mixed_image), mixed_boxes
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
  mixed_boxes = torch.cat([lam * boxes, (1 - lam) * boxes2])
156
 
157
  return TF.to_pil_image(mixed_image), mixed_boxes
158
+
159
+
160
+ class RandomCrop:
161
+ """Randomly crops the image to half its size along with adjusting the bounding boxes."""
162
+
163
+ def __init__(self, prob=0.5):
164
+ """
165
+ Args:
166
+ prob (float): Probability of applying the crop.
167
+ """
168
+ self.prob = prob
169
+
170
+ def __call__(self, image, boxes):
171
+ if torch.rand(1) < self.prob:
172
+ original_width, original_height = image.size
173
+ crop_height, crop_width = original_height // 2, original_width // 2
174
+ top = torch.randint(0, original_height - crop_height + 1, (1,)).item()
175
+ left = torch.randint(0, original_width - crop_width + 1, (1,)).item()
176
+
177
+ image = TF.crop(image, top, left, crop_height, crop_width)
178
+
179
+ boxes[:, [1, 3]] = boxes[:, [1, 3]] * original_width - left
180
+ boxes[:, [2, 4]] = boxes[:, [2, 4]] * original_height - top
181
+
182
+ boxes[:, [1, 3]] = boxes[:, [1, 3]].clamp(0, crop_width)
183
+ boxes[:, [2, 4]] = boxes[:, [2, 4]].clamp(0, crop_height)
184
+
185
+ boxes[:, [1, 3]] /= crop_width
186
+ boxes[:, [2, 4]] /= crop_height
187
+
188
+ return image, boxes