From 1bd1ce5b86bd09dd20a56ca07534c6f4b694cdcc Mon Sep 17 00:00:00 2001
From: wangyingming <wangyingming@megvii.com>
Date: Mon, 11 Apr 2022 15:38:13 +0800
Subject: [PATCH] add the test2017 set

---
 datasets/coco.py | 3 ++-
 main.py          | 4 +++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/datasets/coco.py b/datasets/coco.py
index b392b62..41078ed 100644
--- a/datasets/coco.py
+++ b/datasets/coco.py
@@ -145,7 +145,7 @@ def make_coco_transforms(image_set):
             normalize,
         ])
 
-    if image_set == 'val':
+    if image_set == 'val' or image_set == 'test':
         return T.Compose([
             T.RandomResize([800], max_size=1333),
             normalize,
@@ -161,6 +161,7 @@ def build(image_set, args):
     PATHS = {
         "train": (root / "train2017", root / "annotations" / f'{mode}_train2017.json'),
         "val": (root / "val2017", root / "annotations" / f'{mode}_val2017.json'),
+        "test": (root / "test2017", root / "annotations" / 'image_info_test-dev2017.json'),
     }
 
     img_folder, ann_file = PATHS[image_set]
diff --git a/main.py b/main.py
index 07af29b..05d9490 100644
--- a/main.py
+++ b/main.py
@@ -103,6 +103,8 @@ def get_args_parser():
 
     # dataset parameters
     parser.add_argument('--dataset_file', default='coco')
+    parser.add_argument('--eval_set', default='val', choices=['val', 'test'],
+                        type=str,help="dataset to evaluate")
     parser.add_argument('--coco_path', default='/data/coco', type=str)
     parser.add_argument('--coco_panoptic_path', type=str)
     parser.add_argument('--remove_difficult', action='store_true')
@@ -154,7 +156,7 @@ def main(args):
     print('number of params:', n_parameters)
 
     dataset_train = build_dataset(image_set='train', args=args)
-    dataset_val = build_dataset(image_set='val', args=args)
+    dataset_val = build_dataset(image_set=args.eval_set, args=args)
 
     if args.distributed:
         if args.cache_mode: