Remove is_coco argument from test() (ultralytics#3553)

fhkiel-mlaip · Jun 9, 2021 · 7b687c5 · 7b687c5
1 parent a708b90
commit 7b687c5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/test.py b/test.py
@@ -39,7 +39,6 @@ def test(data,
          wandb_logger=None,
          compute_loss=None,
          half_precision=True,
-         is_coco=False,
          opt=None):
     # Initialize/load model and set device
     training = model is not None
@@ -71,10 +70,10 @@ def test(data,
     # Configure
     model.eval()
     if isinstance(data, str):
-        is_coco = data.endswith('coco.yaml')
         with open(data) as f:
             data = yaml.safe_load(f)
     check_dataset(data)  # check
+    is_coco = data['val'].endswith('coco/val2017.txt')  # COCO dataset
     nc = 1 if single_cls else int(data['nc'])  # number of classes
     iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()

diff --git a/train.py b/train.py
@@ -365,8 +365,7 @@ def train(hyp, opt, device, tb_writer=None):
                                                  verbose=nc < 50 and final_epoch,
                                                  plots=plots and final_epoch,
                                                  wandb_logger=wandb_logger,
-                                                 compute_loss=compute_loss,
-                                                 is_coco=is_coco)
+                                                 compute_loss=compute_loss)
 
             # Write
             with open(results_file, 'a') as f:
@@ -434,8 +433,7 @@ def train(hyp, opt, device, tb_writer=None):
                                               dataloader=testloader,
                                               save_dir=save_dir,
                                               save_json=True,
-                                              plots=False,
-                                              is_coco=is_coco)
+                                              plots=False)
 
             # Strip optimizers
             for f in last, best: