json dict bug fixes and speed improvements

This commit is contained in:
Glenn Jocher 2020-03-13 10:35:58 -07:00
parent 5362e8254e
commit 731305142b
1 changed files with 4 additions and 4 deletions

View File

@ -124,11 +124,11 @@ def test(cfg,
scale_coords(imgs[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for di, d in enumerate(pred):
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({'image_id': image_id,
'category_id': coco91class[int(d[5])],
'bbox': [round(x, 3) for x in box[di].tolist()],
'score': round(d[4].item(), 5)})
'category_id': coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)