Optim parameters

RapidAI · Mar 30, 2024 · 762b359 · 762b359
1 parent 782a659
commit 762b359
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 8 deletions.
diff --git a/label_convert/darknet_to_coco.py b/label_convert/darknet_to_coco.py
@@ -216,14 +216,21 @@ def read_txt(txt_path: str) -> List:
         return data
 
 
-if __name__ == "__main__":
+def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--data_dir",
         default="dataset/darknet_dataset",
         help="Dataset root path",
     )
+    parser.add_argument(
+        "--save_dir", type=str, default=None, help="Path to save the converted dataset."
+    )
     args = parser.parse_args()
 
-    converter = DarknetToCOCO(args.data_dir)
+    converter = DarknetToCOCO(args.data_dir, args.save_dir)
     converter()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/label_convert/labelImg_to_publaynet.py b/label_convert/labelImg_to_publaynet.py
@@ -218,7 +218,6 @@ def main():
         args.data_dir, args.save_dir, args.val_ratio, args.have_test, args.test_ratio
     )
     converter()
-    print(f"Successfully output to the {args.out_dir}")
 
 
 if __name__ == "__main__":

diff --git a/label_convert/labelImg_to_yolov5.py b/label_convert/labelImg_to_yolov5.py
@@ -143,18 +143,17 @@ def cp_file(self, file_path: Path, dst_dir: Path):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--src_dir", type=str, default=None)
+    parser.add_argument("--data_dir", type=str, default=None)
     parser.add_argument("--save_dir", type=str, default=None)
     parser.add_argument("--val_ratio", type=float, default=0.2)
     parser.add_argument("--have_test", action="store_true", default=False)
     parser.add_argument("--test_ratio", type=float, default=0.2)
     args = parser.parse_args()
 
     converter = LabelImgToYOLOV5(
-        args.src_dir, args.save_dir, args.val_ratio, args.have_test, args.test_ratio
+        args.data_dir, args.save_dir, args.val_ratio, args.have_test, args.test_ratio
     )
     converter()
-    print(f"Successfully output to the {args.save_dir}")
 
 
 if __name__ == "__main__":

diff --git a/label_convert/yolov5_to_coco.py b/label_convert/yolov5_to_coco.py
@@ -255,12 +255,13 @@ def main():
         default="tests/test_files/yolov5_dataset",
         help="Dataset root path",
     )
+    parser.add_argument("--save_dir", type=str, default=None)
     parser.add_argument(
         "--mode_list", type=str, default="train,val", help="generate which mode"
     )
     args = parser.parse_args()
 
-    converter = YOLOV5ToCOCO(args.data_dir)
+    converter = YOLOV5ToCOCO(args.data_dir, args.save_dir)
     converter(mode_list=args.mode_list.split(","))
 
 

diff --git a/label_convert/yolov5_yaml_to_coco.py b/label_convert/yolov5_yaml_to_coco.py
@@ -263,9 +263,10 @@ def main():
         default="tests/test_files/crack.v1i.yolov7pytorch/data.yaml",
         help="Dataset cfg file",
     )
+    parser.add_argument("--save_dir", type=str, default=None)
     args = parser.parse_args()
 
-    converter = YOLOV5CfgToCOCO(args.yaml_path)
+    converter = YOLOV5CfgToCOCO(args.yaml_path, args.save_dir)
     converter()