From 782a65977c303b0033e4e12a7920c29a71dd05be Mon Sep 17 00:00:00 2001
From: SWHL <liekkaskono@163.com>
Date: Sat, 30 Mar 2024 15:52:28 +0800
Subject: [PATCH] Update input parameter

---
 label_convert/coco_to_labelImg.py      |  3 +++
 label_convert/darknet_to_coco.py       |  3 +++
 label_convert/labelImg_to_publaynet.py | 16 ++++++++++------
 label_convert/labelImg_to_yolov5.py    |  3 +++
 label_convert/labelme_to_coco.py       |  6 ++++--
 label_convert/yolov5_to_coco.py        |  4 +++-
 6 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/label_convert/coco_to_labelImg.py b/label_convert/coco_to_labelImg.py
index d18cfc8..4c6d36b 100644
--- a/label_convert/coco_to_labelImg.py
+++ b/label_convert/coco_to_labelImg.py
@@ -15,6 +15,9 @@
 
 class COCOTolabelImg:
     def __init__(self, data_dir: ValueType = None, save_dir: ValueType = None):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
+
         self.data_dir = Path(data_dir)
         self.verify_exists(self.data_dir)
 
diff --git a/label_convert/darknet_to_coco.py b/label_convert/darknet_to_coco.py
index e80abe2..347a171 100644
--- a/label_convert/darknet_to_coco.py
+++ b/label_convert/darknet_to_coco.py
@@ -20,6 +20,9 @@ def __init__(
         data_dir: ValueType = None,
         save_dir: ValueType = None,
     ):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
+
         self.data_dir = Path(data_dir)
         self.verify_exists(self.data_dir)
 
diff --git a/label_convert/labelImg_to_publaynet.py b/label_convert/labelImg_to_publaynet.py
index dafe78b..c11b4b1 100644
--- a/label_convert/labelImg_to_publaynet.py
+++ b/label_convert/labelImg_to_publaynet.py
@@ -23,6 +23,9 @@ def __init__(
         have_test: bool = True,
         test_ratio: float = 0.2,
     ):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
+
         self.data_dir = Path(data_dir)
         self.verify_exists(data_dir)
 
@@ -132,7 +135,7 @@ def generate_json(
             label_path = img_path.with_name(f"{img_path.stem}.txt")
             try:
                 label_datas = self.read_txt(label_path)
-            except Exception as e:
+            except Exception:
                 print(f"{label_path} meets error.")
                 continue
 
@@ -204,16 +207,17 @@ def mkdir(dir_path: Union[Path, str]) -> None:
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--data_dir", type=str, required=True, help="The directory from labelImg."
-    )
+    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 = LabelImgToPubLayNet(args.val_ratio, args.have_test, args.test_ratio)
-    converter(args.data_dir)
+    converter = LabelImgToPubLayNet(
+        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}")
 
 
diff --git a/label_convert/labelImg_to_yolov5.py b/label_convert/labelImg_to_yolov5.py
index 5235f69..7ab7f53 100644
--- a/label_convert/labelImg_to_yolov5.py
+++ b/label_convert/labelImg_to_yolov5.py
@@ -22,6 +22,9 @@ def __init__(
         have_test: bool = False,
         test_ratio: float = 0.2,
     ):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
+
         self.data_dir = Path(data_dir)
         self.verify_exists(self.data_dir)
 
diff --git a/label_convert/labelme_to_coco.py b/label_convert/labelme_to_coco.py
index 8df5847..a29a0d0 100644
--- a/label_convert/labelme_to_coco.py
+++ b/label_convert/labelme_to_coco.py
@@ -7,7 +7,7 @@
 import shutil
 import time
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Union
 
 import numpy as np
 from tqdm import tqdm
@@ -16,7 +16,6 @@
 
 
 class LabelmeToCOCO:
-
     def __init__(
         self,
         data_dir: ValueType = None,
@@ -25,7 +24,10 @@ def __init__(
         have_test: bool = False,
         test_ratio: float = 0.2,
     ):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
         self.data_dir = Path(data_dir)
+        self.verify_exists(self.data_dir)
 
         self.val_ratio = val_ratio
         self.test_ratio = test_ratio
diff --git a/label_convert/yolov5_to_coco.py b/label_convert/yolov5_to_coco.py
index f61dba1..80abf94 100644
--- a/label_convert/yolov5_to_coco.py
+++ b/label_convert/yolov5_to_coco.py
@@ -18,8 +18,10 @@
 
 class YOLOV5ToCOCO:
     def __init__(self, data_dir: ValueType = None, save_dir: ValueType = None):
+        if data_dir is None:
+            raise ValueError("data_dir must not be None")
         self.data_dir = Path(data_dir)
-
+        self.verify_exists(self.data_dir)
         self.verify_exists(self.data_dir / "images")
         self.verify_exists(self.data_dir / "labels")