updates

detect.py

@@ -63,19 +63,19 @@ def detect(
         # Get detections
         img = torch.from_numpy(img).unsqueeze(0).to(device)
         pred, _ = model(img)
-        detections = non_max_suppression(pred, conf_thres, nms_thres)[0]
+        det = non_max_suppression(pred, conf_thres, nms_thres)[0]
 
-        if detections is not None and len(detections) > 0:
+        if det is not None and len(det) > 0:
             # Rescale boxes from 416 to true image size
-            detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape).round()
+            det[:, :4] = scale_coords(img.shape, det[:, :4], im0.shape).round()
 
             # Print results to screen
-            for c in detections[:, -1].unique():
-                n = (detections[:, -1] == c).sum()
+            for c in det[:, -1].unique():
+                n = (det[:, -1] == c).sum()
                 print('%g %ss' % (n, classes[int(c)]), end=', ')
 
             # Draw bounding boxes and labels of detections
-            for *xyxy, conf, cls_conf, cls in detections:
+            for *xyxy, conf, cls_conf, cls in det:
                 if save_txt:  # Write to file
                     with open(save_path + '.txt', 'a') as file:
                         file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))

models.py

@@ -182,7 +182,7 @@ class Darknet(nn.Module):
         self.yolo_layers = get_yolo_layers(self)
 
     def forward(self, x, var=None):
-        img_size = x.shape[-1]
+        img_size = max(x.shape[-2:])
         layer_outputs = []
         output = []
 

test.py

@@ -93,7 +93,7 @@ def test(
                 # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
                 image_id = int(Path(paths[si]).stem.split('_')[-1])
                 box = pred[:, :4].clone()  # xyxy
-                scale_coords(img_size, box, shapes[si])  # to original shape
+                scale_coords(imgs[si].shape, box, shapes[si])  # to original shape
                 box = xyxy2xywh(box)  # xywh
                 box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
                 for di, d in enumerate(pred):

utils/datasets.py

@@ -74,7 +74,7 @@ class LoadImages:  # for inference
             print('image %g/%g %s: ' % (self.count, self.nF, path), end='')
 
         # Padded resize
-        img, _, _, _ = letterbox(img0, height=self.height)
+        img, _, _, _ = letterbox_rect(img0, height=self.height)
         print('%gx%g ' % img.shape[:2], end='')  # print image size
 
         # Normalize RGB

utils/utils.py

@@ -102,13 +102,9 @@ def xywh2xyxy(x):
 
 def scale_coords(img1_shape, coords, img0_shape):
     # Rescale coords1 (xyxy) from img1_shape to img0_shape
-    gain = img1_shape / max(img0_shape[:2])  # gain  = old / new
-    # pad_x = np.mod(img1_shape - img0_shape[1] * gain, 32) / 2  # width padding
-    # pad_y = np.mod(img1_shape - img0_shape[0] * gain, 32) / 2  # height padding
-    pad_x = (img1_shape - img0_shape[1] * gain) / 2  # width padding
-    pad_y = (img1_shape - img0_shape[0] * gain) / 2  # height padding
-    coords[:, [0, 2]] -= pad_x
-    coords[:, [1, 3]] -= pad_y
+    gain = max(img1_shape) / max(img0_shape)  # gain  = old / new
+    coords[:, [0, 2]] -= (img1_shape[3] - img0_shape[1] * gain) / 2  # x padding
+    coords[:, [1, 3]] -= (img1_shape[2] - img0_shape[0] * gain) / 2  # y padding
     coords[:, :4] /= gain
     coords[:, :4] = coords[:, :4].clamp(min=0)
     return coords