change padding

modules/dataset_loader.py

@@ -345,29 +345,20 @@ def resize_and_pad(image, target, new_size=(1333, 800)):
     scale = min(new_size[0] / original_size[0], new_size[1] / original_size[1])
     new_scaled_size = (int(original_size[0] * scale), int(original_size[1] * scale))
 
+    target['boxes'] = resize_boxes(target['boxes'], (image.size[0],image.size[1]), (new_scaled_size))
+    if 'area' in target:
+        target['area'] = (target['boxes'][:, 3] - target['boxes'][:, 1]) * (target['boxes'][:, 2] - target['boxes'][:, 0])
+
+    if 'keypoints' in target:
+        for i in range(len(target['keypoints'])):
+            target['keypoints'][i] = resize_keypoints(target['keypoints'][i], (image.size[0],image.size[1]), (new_scaled_size))
+
     # Resize image to new scaled size
     image = F.resize(image, (new_scaled_size[1], new_scaled_size[0]))
 
-    # Calculate padding to center the image
-    pad_left = (new_size[0] - new_scaled_size[0]) // 2
-    pad_top = (new_size[1] - new_scaled_size[1]) // 2
-    pad_right = new_size[0] - new_scaled_size[0] - pad_left
-    pad_bottom = new_size[1] - new_scaled_size[1] - pad_top
-
     # Pad the resized image to make it exactly the desired size
-    image = F.pad(image, (pad_left, pad_top, pad_right, pad_bottom), fill=250, padding_mode='constant')
-
-    # Adjust bounding boxes
-    target['boxes'] = resize_boxes(target['boxes'], original_size, new_scaled_size)
-    target['boxes'][:, 0::2] += pad_left
-    target['boxes'][:, 1::2] += pad_top
-
-    # Adjust keypoints if they exist in the target
-    if 'keypoints' in target:
-        for i in range(len(target['keypoints'])):
-            target['keypoints'][i] = resize_keypoints(target['keypoints'][i], original_size, new_scaled_size)
-            target['keypoints'][i][:, 0] += pad_left
-            target['keypoints'][i][:, 1] += pad_top
+    padding = [0, 0, new_size[0] - new_scaled_size[0], new_size[1] - new_scaled_size[1]]
+    image = F.pad(image, padding, fill=200, padding_mode='edge')
 
     return image, target