@@ -136,7 +136,9 @@ def make_grid(anchors, nx=20, ny=20):
def correct_bbox(result, anchors, stride, cls_num, out):
result = torch.tensor(result)
- bs, _, ny, nx, _ = result.shape
+ bs, _, ny, nx = result.shape
+ result = result.reshape(-1, 3, 85, ny, nx)
+ result = result.permute(0, 1, 3, 4, 2)
grid, anchor_grid = make_grid(anchors, nx, ny)
y = result.float().sigmoid()
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid) * stride # xy
@@ -20,7 +20,7 @@ from pathlib import Path
from common.util.dataset import coco80_to_coco91_class, correct_bbox, save_coco_json
try:
- from utils.general import non_max_suppression, scale_coords # tag > 2.0
+ from utils.general import non_max_suppression, scale_boxes # tag > 2.0
except:
from utils.utils import non_max_suppression, scale_coords # tag = 2.0
@@ -83,7 +83,7 @@ def forward_nms_script(model, dataloader, cfg):
boxout = nms(box_out, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"])
for idx, pred in enumerate(boxout):
try:
- scale_coords(img[idx].shape[1:], pred[:, :4], shapes[idx][0], shapes[idx][1]) # native-space pred
+ scale_boxes(img[idx].shape[1:], pred[:, :4], shapes[idx][0], shapes[idx][1]) # native-space pred
except:
pred = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
# append to COCO-JSON dictionary