diff --git a/detect.py b/detect.py index 7d696221..29151bbd 100644 --- a/detect.py +++ b/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)) diff --git a/models.py b/models.py index 07aa8d47..877e6439 100755 --- a/models.py +++ b/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 = [] diff --git a/test.py b/test.py index e78d6be1..f218e7e9 100644 --- a/test.py +++ b/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): diff --git a/utils/datasets.py b/utils/datasets.py index e1f71667..2ebed801 100755 --- a/utils/datasets.py +++ b/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 diff --git a/utils/utils.py b/utils/utils.py index 4ce682a2..f1c0ea8d 100755 --- a/utils/utils.py +++ b/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