This commit is contained in:
Glenn Jocher 2019-04-26 14:14:28 +02:00
parent 75f08c1cd1
commit 3b134e3a84
1 changed files with 7 additions and 1 deletions

View File

@ -60,6 +60,7 @@ def test(
for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc='Computing mAP')): for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc='Computing mAP')):
targets = targets.to(device) targets = targets.to(device)
imgs = imgs.to(device) imgs = imgs.to(device)
_, _, height, width = imgs.shape # batch size, channels, height, width
# Plot images with bounding boxes # Plot images with bounding boxes
if batch_i == 0 and not os.path.exists('test_batch0.jpg'): if batch_i == 0 and not os.path.exists('test_batch0.jpg'):
@ -108,7 +109,12 @@ def test(
correct = [0] * len(pred) correct = [0] * len(pred)
if nl: if nl:
detected = [] detected = []
tbox = xywh2xyxy(labels[:, 1:5]) * img_size # target boxes tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5])
tbox[[0, 2]] *= width
tbox[[1, 3]] *= height
# Search for correct predictions # Search for correct predictions
for i, (*pbox, pconf, pcls_conf, pcls) in enumerate(pred): for i, (*pbox, pconf, pcls_conf, pcls) in enumerate(pred):