✨ [Add] a RandomCrop transform, align paper
Browse files
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
|