From 0c26c4e83157895ab7cc7492745573bd0213a331 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 14 Nov 2020 13:53:58 +0100 Subject: [PATCH] Increase pycocotools robustness (#1396) --- test.py | 5 ++--- train.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/test.py b/test.py index 10af85b..61bf5d5 100644 --- a/test.py +++ b/test.py @@ -86,8 +86,7 @@ def test(data, img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images - dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, - hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0] + dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True)[0] seen = 0 names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)} @@ -166,7 +165,7 @@ def test(data, box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner for p, b in zip(pred.tolist(), box.tolist()): jdict.append({'image_id': int(image_id) if image_id.isnumeric() else image_id, - 'category_id': coco91class[int(p[5])], + 'category_id': coco91class[int(p[5])] if is_coco else int(p[5]), 'bbox': [round(x, 3) for x in b], 'score': round(p[4], 5)}) diff --git a/train.py b/train.py index 536c8f1..4eee5f2 100644 --- a/train.py +++ b/train.py @@ -183,7 +183,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None): if rank in [-1, 0]: ema.updates = start_epoch * nb // accumulate # set EMA updates testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, - hyp=hyp, augment=False, cache=opt.cache_images and not opt.notest, rect=True, + hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1, world_size=opt.world_size, workers=opt.workers)[0] # testloader if not opt.resume: