From 0e38aec4ce47733bc3791debcdb4eec30885902b Mon Sep 17 00:00:00 2001 From: HenryZhuHR <296506195@qq.com> Date: Sun, 9 Jun 2024 14:45:03 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=91=20hotfix:=20=E4=BF=AE=E5=A4=8D=20?= =?UTF-8?q?=E5=9B=BE=E7=89=87=E5=B0=BA=E5=AF=B8=E6=9C=AA=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E5=8C=96=E7=9A=84=E9=94=99=E8=AF=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- auto_label.py | 32 ++++++++++++++------------------ infer-track.py | 7 +++---- infer.py | 6 +++--- 3 files changed, 20 insertions(+), 25 deletions(-) diff --git a/auto_label.py b/auto_label.py index be6b60f..e13e960 100644 --- a/auto_label.py +++ b/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() diff --git a/infer-track.py b/infer-track.py index 5489ab0..2b8d802 100644 --- a/infer-track.py +++ b/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" ] diff --git a/infer.py b/infer.py index 21137fe..59e9a58 100644 --- a/infer.py +++ b/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())