From 79f5b58a4e19f947c9eef2b720c9ca8fb9e930bd Mon Sep 17 00:00:00 2001
From: Zhiqiang Wang <zhiqwang@foxmail.com>
Date: Sat, 26 Mar 2022 13:34:31 +0800
Subject: [PATCH] Add CLI tools for converting yolo to mscoco

---
 tools/convert_yolo_to_coco.py | 29 +++++++++++++++++++++++++++++
 yolort/utils/yolo2coco.py     | 13 ++++---------
 2 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 tools/convert_yolo_to_coco.py

diff --git a/tools/convert_yolo_to_coco.py b/tools/convert_yolo_to_coco.py
new file mode 100644
index 00000000..fc7f5437
--- /dev/null
+++ b/tools/convert_yolo_to_coco.py
@@ -0,0 +1,29 @@
+# Copyright (c) 2022, yolort team. All rights reserved.
+
+import argparse
+
+from yolort.utils.yolo2coco import YOLO2COCO
+
+
+def get_parser():
+    parser = argparse.ArgumentParser("Datasets converter from yolo to coco", add_help=True)
+
+    parser.add_argument("--data_source", default="./coco128", help="Dataset root path")
+    parser.add_argument("--class_names", default="./coco.name", help="Path of the label names")
+    parser.add_argument("--split", default="train", choices=["train", "val"], help="Dataset split part")
+    parser.add_argument("--year", default=2017, type=int, help="Year of the dataset")
+
+    return parser
+
+
+def cli_main():
+    parser = get_parser()
+    args = parser.parse_args()
+    print(f"Command Line Args: {args}")
+
+    converter = YOLO2COCO(args.data_source, args.class_names, split=args.split, year=args.year)
+    converter.generate()
+
+
+if __name__ == "__main__":
+    cli_main()
diff --git a/yolort/utils/yolo2coco.py b/yolort/utils/yolo2coco.py
index 5c3cd401..ecb7b0b3 100644
--- a/yolort/utils/yolo2coco.py
+++ b/yolort/utils/yolo2coco.py
@@ -16,7 +16,7 @@ class YOLO2COCO:
         root (string): Root directory of the Sintel Dataset
         metalabels (string | List[string]): Concrete label names of different classes
         split (string, optional): The dataset split, either "train" (default) or "test"
-        year (int, optional): The year of the dataset. Default: 2021
+        year (int, optional): The year of the dataset. Default: 2017
     """
 
     def __init__(
@@ -24,12 +24,12 @@ def __init__(
         root: str,
         metalabels: Union[str, List[str]],
         split: str = "train",
-        year: int = 2021,
+        year: int = 2017,
     ) -> None:
 
         self._year = year
         self.type = "instances"
-        self.split = split
+        self.split = f"{split}{year}"
         self.root_path = Path(root)
         self.label_path = self.root_path / "labels"
         self.annotation_root = self.root_path / "annotations"
@@ -50,12 +50,7 @@ def __init__(
             }
         ]
         self.categories = [
-            {
-                "id": coco_category["id"],
-                "name": coco_category["name"],
-                "supercategory": coco_category["supercategory"],
-            }
-            for coco_category in self.metadata
+            {"id": coco_category["id"], "name": coco_category["name"]} for coco_category in self.metadata
         ]
 
     @staticmethod