Add eval json save in coco result format

datasets/coco_eval.py

@@ -32,13 +32,15 @@ def __init__(self, coco_gt, iou_types):
 
         self.img_ids = []
         self.eval_imgs = {k: [] for k in iou_types}
+        self.coco_predictions = []
 
     def update(self, predictions):
         img_ids = list(np.unique(list(predictions.keys())))
         self.img_ids.extend(img_ids)
 
         for iou_type in self.iou_types:
             results = self.prepare(predictions, iou_type)
+            self.coco_predictions.extend(results)
 
             # suppress pycocotools prints
             with open(os.devnull, 'w') as devnull:

main.py

@@ -186,6 +186,7 @@ def main(args):
                                               data_loader_val, base_ds, device, args.output_dir)
         if args.output_dir:
             utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
+            utils.save_json(coco_evaluator.coco_json_predictions, output_dir / "results.json")
         return
 
     print("Start training")

util/misc.py

@@ -10,6 +10,7 @@
 from collections import defaultdict, deque
 import datetime
 import pickle
+import json
 from packaging import version
 from typing import Optional, List
 
@@ -404,6 +405,11 @@ def save_on_master(*args, **kwargs):
         torch.save(*args, **kwargs)
 
 
+def save_json(results, save_dir):
+    with open(save_dir, 'w') as f:
+        f.write(json.dumps(results))
+
+
 def init_distributed_mode(args):
     if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
         args.rank = int(os.environ["RANK"])