✨tool: 新增数据集的半自动化标注

auto_label.py

@@ -0,0 +1,164 @@
+import os
+import argparse
+from typing import Dict, List
+
+import cv2
+import yaml
+
+from dlinfer.detector.backend import DetectorInferBackends
+from dlinfer.detector import Process
+
+from utils.dataset.types import BBOX_XYXY, ObjectLabel_BBOX_XYXY
+from utils.dataset.variables import SUPPORTED_IMAGE_TYPES
+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",
+        )
+        parser.add_argument("-t", "--conf-threshold", type=float, default=0.5)
+        return parser.parse_args()
+
+    def __init__(self) -> None:
+        # fmt: off
+        args = self.get_args()
+        self.image_dir: str = os.path.expandvars(os.path.expanduser(args.image_dir))
+        if not os.path.exists(self.image_dir): # check if the directory exists
+            raise FileNotFoundError(f"Dataset directory not found: {self.image_dir}")
+        self.dataset_config_file: str = os.path.expandvars(os.path.expanduser(args.dataset_config))
+        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
+        self.conf_t: float = args.conf_threshold
+        # fmt: on
+
+
+def main():
+    args = AutoLabelArgs()
+
+    # =============== Load dataset configuration ===============
+    with open(args.dataset_config_file, "r") as f:
+        data_config = yaml.load(f, Loader=yaml.FullLoader)
+    class_map: Dict[int, str] = data_config["names"]  # TODO: data type check
+    print("-- class map :")
+    for i, c in class_map.items():
+        print(f"  {i}: {c}")
+
+    # =============== Choose backend to Infer ===============
+    backends = DetectorInferBackends()
+    ## ------ ONNX ------
+    # onnx_backend = backends.ONNXBackend
+    # print("-- Available devices:", providers := onnx_backend.SUPPORTED_DEVICES)
+    # detector = onnx_backend(
+    #     device=providers, inputs=["images"], outputs=["output0"]
+    # )
+
+    ## ------ OpenVINO ------
+    ov_backend = backends.OpenVINOBackend
+    print("-- Available devices:", ov_backend.query_device())
+    detector = ov_backend(device="AUTO")
+
+    ## ------ TensorRT ------
+    # detector = backends.TensorRTBackend()
+
+    # =======================================================
+    detector.load_model(args.model, verbose=True)
+
+    image_dir = args.image_dir
+
+    for file in os.listdir(image_dir):
+        # 获取文件后缀，查看是否是图片文件
+        suffix = os.path.splitext(file)[-1]
+        if suffix not in [f".{ext}" for ext in SUPPORTED_IMAGE_TYPES]:
+            continue
+        file_name = os.path.splitext(file)[0]
+        xml_file = os.path.join(image_dir, f"{file_name}.xml")
+        if os.path.exists(xml_file):
+            print(
+                f"File {xml_file} already exists. "
+                f"If you want to re-label, please delete it by 'rm {xml_file}'"
+            )
+            continue
+
+        # =============== 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
+        output_t = detector.infer(input_t)
+        preds = Process.postprocess(output_t)
+        end_time = cv2.getTickCount()
+        infer_time = (end_time - start_time) / cv2.getTickFrequency() * 1000
+
+        # print(f"File: {file}")
+        # print(preds)
+
+        bboxes: List[ObjectLabel_BBOX_XYXY] = []
+        cls_cnt = 0
+        for pred in preds:
+            x1 = int(scale_w * pred[0])
+            y1 = int(scale_h * pred[1])
+            x2 = int(scale_w * pred[2])
+            y2 = int(scale_h * pred[3])
+            conf = pred[4]
+            clsid = int(pred[5])
+            if conf < args.conf_t:
+                continue
+            bbox = BBOX_XYXY(int(x1), int(y1), int(x2), int(y2))
+            cls = class_map[clsid]
+            bboxes.append(ObjectLabel_BBOX_XYXY(cls, bbox))
+            cls_cnt += 1
+
+        # =============== Save to xml ===============
+        size = (img.shape[1], img.shape[0], img.shape[2])
+        root = ET.Element("annotation")
+        filename = ET.SubElement(root, "filename")
+        filename.text = file
+        size_node = ET.SubElement(root, "size")
+        width = ET.SubElement(size_node, "width")
+        width.text = str(size[0])
+        height = ET.SubElement(size_node, "height")
+        height.text = str(size[1])
+        depth = ET.SubElement(size_node, "depth")
+        depth.text = str(size[2])
+        for obj in bboxes:
+            object_node = ET.SubElement(root, "object")
+            name = ET.SubElement(object_node, "name")
+            name.text = obj.cls
+            bndbox = ET.SubElement(object_node, "bndbox")
+            xmin = ET.SubElement(bndbox, "xmin")
+            xmin.text = str(obj.bbox.xmin)
+            ymin = ET.SubElement(bndbox, "ymin")
+            ymin.text = str(obj.bbox.ymin)
+            xmax = ET.SubElement(bndbox, "xmax")
+            xmax.text = str(obj.bbox.xmax)
+            ymax = ET.SubElement(bndbox, "ymax")
+            ymax.text = str(obj.bbox.ymax)
+
+        tree = ET.ElementTree(root)
+        tree.write(xml_file, encoding="utf-8")
+        print(
+            f"Infer {infer_time:.3f} ms, File: {file}, {cls_cnt} objects saved to {xml_file} ({[b.cls for b in bboxes]})"
+        )
+
+
+if __name__ == "__main__":
+    main()

dataset-process.py

@@ -21,7 +21,7 @@ def __init__(self) -> None:
     @staticmethod
     def get_args():
         parser = argparse.ArgumentParser()
-        parser.add_argument("--datadir", type=str, default="~/data/bottle")
+        parser.add_argument("-d","--datadir", type=str, default="~/data/drink")
         return parser.parse_args()
 
 
@@ -37,8 +37,11 @@ def main():
         os.makedirs(organized_datadir, exist_ok=False)
     else:
         raise FileExistsError(
-            f"Directory '{organized_datadir}' already exists, "
-            f"delete by 'rm -rf {organized_datadir}'"
+            f"Directory '{organized_datadir}' already exists."
+            f"\033[00;33m To avoid overwriting, please manually delete by\033[0m"
+            f"\033[00;32m 'rm -rf {organized_datadir}'\033[0m"
+            f"\033[00;33m and run this script again.\033[0m"
+
         )
 
     images_dir = os.path.join(organized_datadir, "images")

dlinfer/detector/b_tensorrt.py

@@ -16,12 +16,14 @@
     from .b_tensorrt_10_0 import TensorRTDetectorBackend_10_0
 
     TensorRTDetectorBackend = TensorRTDetectorBackend_10_0
-elif f"{tv_major}.{tv_minor}" in ["8.6"]:
+elif f"{tv_major}.{tv_minor}" in ["8.4", "8.6"]:
     from .b_tensorrt_8_6 import TensorRTDetectorBackend_8_6
 
     TensorRTDetectorBackend = TensorRTDetectorBackend_8_6
 elif f"{tv_major}.{tv_minor}" in ["8.2"]:
-    raise NotImplementedError("TensorRT 8.2 is to be implemented in the future only for Jetson Nano device.")
+    raise NotImplementedError(
+        "TensorRT 8.2 is to be implemented in the future only for Jetson Nano device."
+    )
     # from .b_tensorrt_8_2 import TensorRTDetectorBackend_8_2
     # TensorRTDetectorBackend = TensorRTDetectorBackend_8_2
 else:

docs/.vitepress/config.mts

@@ -41,6 +41,7 @@ export default defineConfig({
     outline: {
       label: '页面导航'
     },
+
     lastUpdated: {
       text: '最后更新于',
       formatOptions: {
@@ -59,6 +60,7 @@ export default defineConfig({
       { icon: 'github', link: 'https://github.com/HenryZhuHR/deep-object-detect-track' }
     ]
   },
+  lastUpdated: true,
   markdown: {
     math: true,
     lineNumbers: false

docs/dataset.md

@@ -9,6 +9,10 @@ outline: deep
 
 ## 数据集采集和归档
 
+::: tip
+数据是比代码更重要的资产，因此不要放置在项目内
+:::
+
 将数据集放入如下目录
 
 ```shell
@@ -17,17 +21,24 @@ DATASET_DIR=/path/to/dataset
 
 > 需要注意的是，数据集通常需要放置在项目外的路径，例如 `~/data` 或 `$HOME/data` （推荐）（win 下为 `$env:USERPROFILE/data`）。如果放置在项目内，导致编辑器对于项目的索引过大，会导致编辑器卡顿
 
-这里准备好了一个示例数据集，可以下载
+这里准备好了一个示例数据集，可以下载，并保存在 `~/data/drink` 目录下
+
+```shell
+wget -P ~/data https://github.com/HenryZhuHR/deep-object-detect-track/releases/download/v1.0.0/drink.tar.bz2
+tar -xf ~/data/drink.tar.bz2 -C ~/data
+cp -r ~/data/drink ~/data/drink.unlabel
+rm -rf ~/data/drink.unlabel/**/*.xml
+```
 
+为了方便可以在项目内建立软链接，软链接不会影响编辑器进行索引，但是可以方便查看
+
 ```shell
-wget -P ~/data https://github.com/HenryZhuHR/deep-object-detect-track/releases/download/v1.0.0/bottle.tar.xz
-tar -xvf ~/data/bottle.tar.xz -C ~/data
-mv ~/data/bottle ~/data/yolodataset
+ln -s ~/data resource/data
 ```
 
 随后可以设置数据集目录为
 ```shell
-DATASET_DIR=~/data/yolodataset
+DATASET_DIR=~/data/drink
 ```
 
 参考该目录构建自己的数据集，并且完成标注
@@ -39,21 +50,21 @@ DATASET_DIR=~/data/yolodataset
 ·
 └── /path/to/dataset
     ├── class_A         
-    │   ├─ file_A1.jpg  
-    │   ├─ file_A1.xml     
+    │   ├─ file_A1.jpg
+    │   ├─ file_A1.xml
     │   └─ ...
     └── class_B       
-        ├─ file_B1.jpg   
-        ├─ file_B1.xml   
+        ├─ file_B1.jpg
+        ├─ file_B1.xml
         └─ ...
 ```
 
 不进行类别划分的目录结构参考
 ```shell
 ·
 └── /path/to/dataset    
-    ├─ file_1.jpg  
-    ├─ file_1.xml     
+    ├─ file_1.jpg
+    ├─ file_1.xml
     └─ ...
 ```
 
@@ -103,10 +114,23 @@ labelImg
 
 运行脚本，生成同名目录，但是会带 `-organized` 后缀，例如
 ```shell
-python dataset-process.py --datadir ~/data/yolodataset
+python dataset-process.py --datadir ~/data/drink
 ```
 
-生成的目录 `~/data/yolodataset-organized` 用于数据集训练，并且该目录为 yolov5 中指定的数据集路径
+该脚本会自动递归地扫描目录 `~/data/drink` 下的所有 `.xml` 文件，并查看是否存在对应的 `.jpg` 文件
+
+::: tip
+因此，你可以不必担心目录结构，只需要确保每张图像有对应的标签文件即可，也不必担心没有标注完成的情况，脚本只处理以及标注完成的图像
+:::
+
+生成的目录 `~/data/drink-organized` 用于数据集训练，并且该目录为 yolov5 中指定的数据集路径
+
+
+## 数据自定义处理
+
+::: tip
+通常来说不需要自定义处理，只需要遵循上述的规则即可快速创建数据集，但是如果需要，可以参考下面提供的接口
+:::
 
 如果不需要完全遍历数据集、数据集自定义路径，则在 `get_all_label_files()` 函数中传入自定义的 `custom_get_all_files` 函数，以获取全部文件路径，该自定义函数可以参考 `default_get_all_files()`
 
@@ -128,4 +152,31 @@ label_file_list = get_all_label_files(          # [!code ++]
     args.datadir,                               # [!code ++]
     custom_get_all_files=default_get_all_files  # [!code ++]
 )                                               # [!code ++]
+```
+
+
+
+## 半自动标注
+
+训练完少量数据集后可以使用 `auto_label.py` 脚本进行半自动标注
+
+该脚本需要使用 OpenVINO 的模型进行推理，因此参考 [*导出模型*](./deploy.md#导出模型) 导出 openvino 模型，主要修改 `EXPORTED_MODEL_PATH` 为导出的模型路径和 数据集配置 `DATASET_CONFIG` ，然后注释导出命令 `python3 export.py ... --include onnx ` 和 `python3 export.py ... --include engine `
+
+```shell
+bash scripts/export-yolov5.sh
+```
+
+查看 `auto_label.py` 参数设置，然后执行
+```shell
+python auto_label.py
+```
+
+如果已经被标注，会提示 `If you want to re-label, please delete it by 'rm ~/data/drink.unlabel/cola/cola_0000.xml'` 的信息防止已经被标注的数据被覆盖，如果希望删除全部，可以用正则表达式
+```shell
+rm -rf ~/data/drink.unlabel/**/*.xml
+```
+
+随后可以用 LabelImg 进行**检查**和精细化调整
+```shell
+labelImg ~/data/drink.unlabel
 ```

docs/download-pretrian.md

@@ -4,11 +4,11 @@
 
 ```shell [bash]
 bash scripts/download-yolov5.sh yolov5s # 仅下载单个模型
-# bash scripts/download-yolov5.sh         # 下载所有模型
+# bash scripts/download-yolov5.sh       # 下载所有模型
 ```
 ```shell [zsh]
 zsh scripts/download-yolov5.sh yolov5s # 仅下载单个模型
-# zsh scripts/download-yolov5.sh         # 下载所有模型
+# zsh scripts/download-yolov5.sh       # 下载所有模型
 ```
 
 :::