🚑 hotfix: 修复 图片尺寸未参数化的错误

auto_label.py

@@ -13,29 +13,18 @@
 import xml.etree.ElementTree as ET
 
 
-
 class AutoLabelArgs:
 
     @staticmethod
     def get_args():
         parser = argparse.ArgumentParser()
-        parser.add_argument(
-            "-d", "--image-dir", type=str, default="~/data/drink.unlabel/cola"
-        )
-        parser.add_argument(
-            "-c",
-            "--dataset-config",
-            type=str,
-            default="~/data/drink-organized/dataset.yaml",
-        )
-        parser.add_argument(
-            "-m",
-            "--model",
-            type=str,
-            # default=".cache/yolov5/yolov5s.onnx",
-            default="temp/drink-yolov5x6/weights/best.onnx",
-        )
+        # fmt: off
+        parser.add_argument("-d", "--image-dir", type=str, default="~/data/drink.unlabel/cola")
+        parser.add_argument("-c", "--dataset-config", type=str, default="~/data/drink-organized/dataset.yaml")
+        parser.add_argument("-m", "--model", type=str, default="temp/drink-yolov5x6/weights/best.onnx")
+        parser.add_argument("-s", "--img-size", nargs="+", type=int, default=[640, 640])
         parser.add_argument("-t", "--conf-threshold", type=float, default=0.5)
+        # fmt: on
         return parser.parse_args()
 
     def __init__(self) -> None:
@@ -48,6 +37,12 @@ def __init__(self) -> None:
         if not os.path.exists(self.dataset_config_file): # check if the directory exists
             raise FileNotFoundError( f"Dataset configuration file not found: {self.dataset_config_file}")
         self.model: str = args.model
+        if len(args.img_size) == 2:
+            self.img_size: List[int] = args.img_size
+        elif len(args.img_size) == 1:
+            self.img_size: List[int] = [args.img_size, args.img_size]
+        else:
+            raise ValueError("Invalid img_size")
         self.conf_t: float = args.conf_threshold
         # fmt: on
 
@@ -84,6 +79,7 @@ def main():
     detector.load_model(args.model, verbose=True)
 
     image_dir = args.image_dir
+    img_size = args.img_size
 
     for file in os.listdir(image_dir):
         # 获取文件后缀，查看是否是图片文件
@@ -102,7 +98,7 @@ def main():
         # =============== Auto label ===============
         start_time = cv2.getTickCount()
         img = cv2.imread(os.path.join(image_dir, file))  # H W C
-        input_t, scale_h, scale_w = Process.preprocess(img)  # B C H W
+        input_t, scale_h, scale_w = Process.preprocess(img, img_size)  # B C H W
         output_t = detector.infer(input_t)
         preds = Process.postprocess(output_t)
         end_time = cv2.getTickCount()

infer-track.py

@@ -32,7 +32,7 @@ def get_args():
         parser.add_argument("-v", "--video", type=str, default=".cache/palace.mp4")
         parser.add_argument("-o", "--outdir", type=str, default="tmp")
         parser.add_argument("-m", "--model", type=str, default=".cache/yolov5/yolov5s_openvino_model/yolov5s.xml")
-        parser.add_argument("-c", "--config", type=str, default=".cache/yolov5/coco.yaml")
+        parser.add_argument("-c", "--dataset-config", type=str, default=".cache/yolov5/coco.yaml")
         parser.add_argument("-s", "--img-size", nargs="+", type=int, default=[640, 640])
         parser.add_argument("--aspect_ratio_thresh", type=float, default=1.6,
             help="threshold for filtering out boxes of which aspect ratio are above the given value." )
@@ -42,11 +42,10 @@ def get_args():
 
     def __init__(self) -> None:
         args = self.get_args()
-
         self.video: str = args.video
         self.output_dir: str = args.outdir
         self.model_path: str = args.model
-        self.config: str = args.config
+        self.dataset_config: str = args.dataset_config
         if len(args.img_size) == 2:
             self.img_size: List[int] = args.img_size
         elif len(args.img_size) == 1:
@@ -84,7 +83,7 @@ def main():
     print("-- Available devices:", detector.query_device())
     detector.load_model(args.model_path, verbose=True)
 
-    with open(args.config, "r") as f:
+    with open(args.dataset_config, "r") as f:
         label_map: Dict[int, str] = yaml.load(f, Loader=yaml.FullLoader)[
             "names"
         ]

infer.py

@@ -15,7 +15,7 @@ def get_args():
         parser = argparse.ArgumentParser(add_help=False)
         # fmt: off
         parser.add_argument("-m", "--model", type=str, default=".cache/yolov5/yolov5s.onnx")
-        parser.add_argument("-c", "--config", type=str, default=".cache/yolov5/coco.yaml")
+        parser.add_argument("-c", "--dataset-config", type=str, default=".cache/yolov5/coco.yaml")
         parser.add_argument("-s", "--img-size", nargs="+", type=int, default=[640, 640])
         parser.add_argument("-i", "--input", type=str, default="images/bus.jpg")
         # fmt: on
@@ -24,7 +24,7 @@ def get_args():
     def __init__(self) -> None:
         args = self.get_args()
         self.model: str = args.model
-        self.config: str = args.config
+        self.dataset_config: str = args.dataset_config
         if len(args.img_size) == 2:
             self.img_size: List[int] = args.img_size
         elif len(args.img_size) == 1:
@@ -55,7 +55,7 @@ def main() -> int:
 
     detector.load_model(args.model, verbose=True)
 
-    with open(args.config, "r") as f:
+    with open(args.dataset_config, "r") as f:
         file_content = yaml.load(f, Loader=yaml.FullLoader)
     label_map: Dict[int, str] = file_content["names"]
     label_list = list(label_map.values())