diff --git a/auto_label.py b/auto_label.py new file mode 100644 index 0000000..be6b60f --- /dev/null +++ b/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() diff --git a/dataset-process.py b/dataset-process.py index 6a30ecd..2e8c81d 100644 --- a/dataset-process.py +++ b/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") diff --git a/dlinfer/detector/b_tensorrt.py b/dlinfer/detector/b_tensorrt.py index 21ceded..3eafa94 100644 --- a/dlinfer/detector/b_tensorrt.py +++ b/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: diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index 8f8e7e8..eb57f95 100644 --- a/docs/.vitepress/config.mts +++ b/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 diff --git a/docs/dataset.md b/docs/dataset.md index dffd029..bf323dd 100644 --- a/docs/dataset.md +++ b/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,12 +50,12 @@ 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 └─ ... ``` @@ -52,8 +63,8 @@ DATASET_DIR=~/data/yolodataset ```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 ``` \ No newline at end of file diff --git a/docs/download-pretrian.md b/docs/download-pretrian.md index c5bf325..4f75806 100644 --- a/docs/download-pretrian.md +++ b/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 # 下载所有模型 ``` ::: diff --git a/docs/install.md b/docs/install.md index 4a8a7e9..7507254 100644 --- a/docs/install.md +++ b/docs/install.md @@ -57,7 +57,9 @@ git clone https://github.com/ultralytics/yolov5.git projects/yolov5 - ✅ Ubuntu 22.04 jammy (CPU & GPU) - ✅ MacOS (CPU) -> 项目不支持 Windows 系统 ❌ ，如果需要在 Windows 系统上运行，可以使用 WSL2 或者根据提供的脚本手动执行；虽然已经测试通过，但是不保证所有功能都能正常运行，因此不接受 Windows 系统的问题反馈 +::: warning +项目不支持 Windows 系统 ❌ ，如果需要在 Windows 系统上运行，可以使用 WSL2 或者根据提供的脚本手动执行；虽然已经测试通过，但是不保证所有功能都能正常运行，因此不接受 Windows 系统的问题反馈 +::: ### GPU @@ -84,7 +86,7 @@ MacOS 系统不支持 CUDA Toolkit，可以使用 CPU 训练模型 (Yolov5 项 提供两种方式安装， venv 或 conda -- **venv** : 如果没有安装，请安装 +- **venv** : 嵌入式设备的部署建议使用这种方案，以确保链接到系统的库，如果没有安装，请安装 ::: code-group @@ -115,94 +117,63 @@ zsh Miniconda3-latest-MacOSX-arm64.sh ::: - -### 方法一：使用提供的脚本 - - -提供的安装脚本依赖于基本环境变量 `scripts/variables.sh` ，可以复制一份到项目目录下进行自定义修改（推荐），如果不需要修改，可以直接执行 - -```shell -cp scripts/variables.sh scripts/variables.custom.sh -``` -- `CACHE_DIR`: 用于存放一些缓存文件，例如 `yolov5/requirements.txt`，默认为项目目录下的 `.cache` -- 安装过程会自动检测 `CUDA_VERSION` 以安装对应的 PyTorch 版本，否则默认安装 CPU 版本的 PyTorch；如果电脑有 NVIDIA GPU 但是不想安装 CUDA Toolkit 到全局系统（需要 sudo）可以取消注释 `export CUDA_VERSION=12.1` 以安装对应的 PyTorch 版本 - -运行会自动检测是否存在用户自定义的环境变量 `scripts/variables.custom.sh` ，如果存在则使用自定义的环境变量，否则使用默认的环境变量 `scripts/variables.sh` - -执行命令自动创建并且激活虚拟环境，默认使用 `venv`，**可以重复执行该脚本获取激活环境的提示信息或者安装依赖** - -::: code-group - -```shell [使用 venv 创建虚拟环境] -bash scripts/create-python-env.sh -i # -i 自动安装依赖 -#zsh scripts/create-python-env.sh -i # zsh -``` - -```shell [使用 conda 创建虚拟环境] -bash scripts/create-python-env.sh -e conda -i # -i 自动安装依赖 -#zsh scripts/create-python-env.sh -e conda -i # zsh -``` - -::: - -- 该脚本会复制 `yolov5/requirements.txt` 到 `.cache/yolov5/requirements.txt`，可以自行修改 `.cache/yolov5/requirements.txt` 文件安装相关依赖，例如取消 `onnx` 的注释以支持 ONNX 格式的模型导出；可以修改后再次执行脚本以重新安装依赖 - -### 方法二：手动安装 +### 方法一：手动安装 创建虚拟环境 ::: code-group ```shell [在项目内安装环境(推荐)] -export ENV_NAME=deep-object-detect-track -conda create -p .env/$ENV_NAME python=3.10 -y -conda activate ./.env/$ENV_NAME +conda create -p .env/deep-object-detect-track python=3.10 -y +conda activate ./.env/deep-object-detect-track ``` ```shell [全局安装环境] -export ENV_NAME=deep-object-detect-track -conda create -n $ENV_NAME python=3.10 -y -conda activate $ENV_NAME +conda create -n deep-object-detect-track python=3.10 -y +conda activate deep-object-detect-track ``` ::: > Python 版本选择 3.10 是因为 Ubuntu 22.04 默认安装的 Python 版本是 3.10 -1. 安装 PyTorch -参考官网 [*INSTALL PYTORCH*](https://pytorch.org/get-started/locally/) 选择配置安装 PyTorch +- 如果电脑有 NVIDIA GPU，可以直接安装 [PyTorch]((https://pytorch.org/get-started/locally/)) 和其他依赖 +```shell +pip install -r requirements.txt +``` +- 如果电脑没有 NVIDIA GPU，可以安装 CPU 版本的 PyTorch ```shell -pip install torch torchvision \ - --index-url https://download.pytorch.org/whl/cu121 +pip install -r requirements/requirements-cpu.txt ``` -> 链接最后的 `cu121` 是需要根据系统的 CUDA 版本进行选择 -接下来安装其他依赖 + +### 方法二：使用提供的脚本 + +提供的安装脚本依赖于基本环境变量 `scripts/variables.sh` ，可以复制一份到项目目录下进行自定义修改（推荐），如果不需要修改，可以直接执行 ```shell -pip install -r projects/yolov5/requirements.txt -pip install -r requirements/requirements.train.txt +cp scripts/variables.sh scripts/variables.custom.sh ``` +- `CACHE_DIR`: 用于存放一些缓存文件，例如 `yolov5/requirements.txt`，默认为项目目录下的 `.cache` +- 安装过程会自动检测 `CUDA_VERSION` 以安装对应的 PyTorch 版本，否则默认安装 CPU 版本的 PyTorch；如果电脑有 NVIDIA GPU 但是不想安装 CUDA Toolkit 到全局系统（需要 sudo）可以取消注释 `export CUDA_VERSION=12.1` 以安装对应的 PyTorch 版本 -如果涉及部署流程，需要自行修改 `requirements.txt` 文件，将下列依赖取消注释掉，然后重新执行上述命令 +运行会自动检测是否存在用户自定义的环境变量 `scripts/variables.custom.sh` ，如果存在则使用自定义的环境变量，否则使用默认的环境变量 `scripts/variables.sh` + +执行命令自动创建并且激活虚拟环境，默认使用 `venv`，**可以重复执行该脚本获取激活环境的提示信息或者安装依赖** -```txt -# Export ---------------------------------------------------------------------- -# coremltools>=6.0 # CoreML export -# onnx>=1.10.0 # ONNX export -# onnx-simplifier>=0.4.1 # ONNX simplifier -# nvidia-pyindex # TensorRT export -# nvidia-tensorrt # TensorRT export -# scikit-learn<=1.1.2 # CoreML quantization -# tensorflow>=2.4.0,<=2.13.1 # TF exports (-cpu, -aarch64, -macos) -# tensorflowjs>=3.9.0 # TF.js export -# openvino-dev>=2023.0 # OpenVINO export +::: code-group + +```shell [使用 venv 创建虚拟环境] +bash scripts/create-python-env.sh -i # -i 自动安装依赖 +#zsh scripts/create-python-env.sh -i # zsh ``` -- `onnx`: ONNX 格式的模型导出支持任意设备，需要取消注释，并且其他导出依赖于 ONNX 模型 -- `coremltools`: 必须依赖于 MacOS 系统 -- `nvidia-*`: 确保硬件支持 NVIDIA GPU +```shell [使用 conda 创建虚拟环境] +bash scripts/create-python-env.sh -e conda -i # -i 自动安装依赖 +#zsh scripts/create-python-env.sh -e conda -i # zsh +``` +::: diff --git a/docs/train.md b/docs/train.md index f6d233e..53f6d4c 100644 --- a/docs/train.md +++ b/docs/train.md @@ -12,19 +12,39 @@ outline: deep -## 训练模型 -提供一个训练脚本 `scripts/train.sh`，复制一份到项目目录下进行自定义修改（推荐） +## 训练模型（手动设置参数） + ```shell -cp scripts/train.sh scripts/train.custom.sh +python train.py \ + --data ~/data/drink-organized/dataset.yaml \ + --cfg models/yolov5s.yaml --weights ../../.cache/yolov5/yolov5s.pt \ + --epochs 10 --batch 2 \ + --device 0 --workers 8 --cache disk \ + --project ../../tmp/train ``` +- `--data`: 数据集配置文件，yaml 格式 +- `--cfg`: 模型配置文件，yaml 格式，需要和预训练模型参数 `--weights` 对应 +- `--device`: 如果有 GPU 可以设置为 `0` 或者 `1`；多个 GPU 用逗号分隔，如果单 GPU 能够满足需求，建议使用单 GPU +- `--batch`: batch size，最重要的参数，建议从 1/2 开始尝试，以避免内存不足，而不要一开始就设置一个很大的值导致系统崩溃 +- `--cache`: 缓存类型，`ram` 或者 `disk`，可以加快训练速度，如果内存足够，建议使用 `ram` ，也可以不使用该参数而不缓存数据 + +如果已经有了训练好的模型，可以继续训练，使用 `--resume` 参数从上次训练的 `last.pt` 模型继续训练 + +```shell +## 训练模型（使用提供的脚本） +提供一个训练脚本 `scripts/train.sh`，复制一份到项目目录下进行自定义修改（推荐） + +```shell +cp scripts/train.sh scripts/train.custom.sh +``` 修改训练脚本 `train.custom.sh` 中的参数，在 `Set Training Variables` 中设置 - `TRAIN_DATA_DIR`: (✅重要) 训练数据路径，例如 `$HOME/data/bottle-organized` -- `TRAIN_DEVICE`: (✅重要) 这里默认设置为 `cpu`；如果有 GPU 可以设置为 `0` 或者 `1`；多个 GPU 用逗号分隔，如果单 GPU 能够满足需求，建议使用单 GPU +- `TRAIN_DEVICE`: (✅重要) 如果有 GPU 可以设置为 `0` 或者 `1`；多个 GPU 用逗号分隔，如果单 GPU 能够满足需求，建议使用单 GPU - `MODEL_NAME`: 训练模型，用以加载 `--cfg` 参数，默认为 `yolov5s` - `PRETRAINED_MODEL`: 上面下载好的预训练模型 - `EPOCHS`: epoch diff --git a/infer.py b/infer.py index 4cf60b2..a1d0715 100644 --- a/infer.py +++ b/infer.py @@ -25,14 +25,14 @@ def main() -> int: # =============== Choose backend to Infer =============== # ------ Choose one and comment out the others ------ ## ------ ONNX ------ - onnx_backend = backends.ONNXBackend - print("-- Available devices:", providers := onnx_backend.SUPPORTED_DEVICES) - detector = onnx_backend(device=providers, inputs=["images"], outputs=["output0"]) + # 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="GPU.0") + ov_backend = backends.OpenVINOBackend + print("-- Available devices:", ov_backend.query_device()) + detector = ov_backend(device="AUTO") ## ------ TensorRT ------ # detector = backends.TensorRTBackend() @@ -41,7 +41,9 @@ def main() -> int: detector.load_model(args.model, verbose=True) with open(".cache/yolov5/coco.yaml", "r") as f: - label_map: Dict[int, str] = yaml.load(f, Loader=yaml.FullLoader)["names"] + label_map: Dict[int, str] = yaml.load(f, Loader=yaml.FullLoader)[ + "names" + ] label_list = list(label_map.values()) # print(label_list) @@ -71,7 +73,9 @@ def main() -> int: total_time = (end_time - start_time) / cv2.getTickFrequency() * 1000 total_sum_time += total_time _cnt += 1 - pbar.set_description(f"Time infer/total: {infer_time:.2f}/{total_time:.2f} ms") + pbar.set_description( + f"Time infer/total: {infer_time:.2f}/{total_time:.2f} ms" + ) # -- mark print(f"-- Average time: {total_sum_time / _cnt:.2f} ms") Process.mark(img, preds, label_list, scale_h, scale_w) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..401b3b8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,15 @@ +-r requirements/torch.txt +-r requirements/yolov5.txt + +--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple +--extra-index-url https://pypi.org/simple + +# onnxruntime +onnxruntime==1.18.0 + +# openvino +openvino==2024.1.0 + +# tensorrt +tensorrt==10.0.1 +cuda-python==12.1.0 \ No newline at end of file diff --git a/requirements/requirements-cpu.txt b/requirements/requirements-cpu.txt new file mode 100644 index 0000000..3e565a3 --- /dev/null +++ b/requirements/requirements-cpu.txt @@ -0,0 +1,17 @@ +torch==2.3.1 +torchvision==0.18.1 + +-r yolov5.txt + +--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple +--extra-index-url https://pypi.org/simple + +# onnxruntime +onnxruntime==1.18.0 + +# openvino +openvino==2024.1.0 + +# tensorrt +tensorrt==10.0.1 +cuda-python==12.1.0 \ No newline at end of file diff --git a/requirements/requirements.train.txt b/requirements/requirements.train.txt deleted file mode 100644 index 3f72d6c..0000000 --- a/requirements/requirements.train.txt +++ /dev/null @@ -1,4 +0,0 @@ --r requirements.txt - -# -- dataset -- -labelImg \ No newline at end of file diff --git a/requirements/torch.txt b/requirements/torch.txt new file mode 100644 index 0000000..08d04fd --- /dev/null +++ b/requirements/torch.txt @@ -0,0 +1,3 @@ +--index-url https://download.pytorch.org/whl/cu121 +torch==2.3.1 +torchvision==0.18.1 \ No newline at end of file diff --git a/requirements/requirements.txt b/requirements/track.txt similarity index 100% rename from requirements/requirements.txt rename to requirements/track.txt diff --git a/requirements/yolov5.txt b/requirements/yolov5.txt new file mode 100644 index 0000000..e97817b --- /dev/null +++ b/requirements/yolov5.txt @@ -0,0 +1,8 @@ +--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple +--extra-index-url https://pypi.org/simple + +opencv-python +PyYAML +requests +tqdm +pandas diff --git a/scripts/base.sh b/scripts/base.sh index 005086a..99ac165 100644 --- a/scripts/base.sh +++ b/scripts/base.sh @@ -22,6 +22,19 @@ function run_script { source $1 } +function copy_file_to_custom { + source_file=$1 + f_without_suffix=$(echo $source_file | sed 's/\.sh//') + custom_file="$f_without_suffix.custom.sh" + if [ ! -f "$custom_file" ] && [ -f "$source_file" ] ; then + cp $source_file $custom_file + print_success "Copied '$source_file' to '$custom_file'" + fi +} + +# copy_file_to_custom "scripts/variables.sh" +# copy_file_to_custom "scripts/train.sh" +# copy_file_to_custom "scripts/export-yolov5.sh" if [ -f "scripts/variables.custom.sh" ]; then run_script "scripts/variables.custom.sh" diff --git a/scripts/download-yolov5.sh b/scripts/download-yolov5.sh index 46fd74f..b860021 100644 --- a/scripts/download-yolov5.sh +++ b/scripts/download-yolov5.sh @@ -7,11 +7,15 @@ source scripts/base.sh tag_name=v7.0 all_model_list=( yolov5n + yolov5n6 yolov5s + yolov5s6 yolov5m + yolov5m6 yolov5l + yolov5l6 yolov5x - yolov5n + yolov5x6 ) weights_dir=$CACHE_DIR/yolov5 diff --git a/scripts/export-yolov5.sh b/scripts/export-yolov5.sh index 2890e9a..cc1d3b6 100644 --- a/scripts/export-yolov5.sh +++ b/scripts/export-yolov5.sh @@ -24,7 +24,7 @@ EXPORTED_MODEL_PATH=$PROJECT_HOME/.cache/yolov5/yolov5s.pt # CUDA:1 (NVIDIA GeForce RTX 4090, 24217MiB) # 电脑有多少显卡就得给多少 TRT_EXPORTED_DEVICE="0,1" # Multiple GPUs -TRT_EXPORTED_DEVICE="0" # Single GPU +# TRT_EXPORTED_DEVICE="0" # Single GPU # ======================================================= diff --git a/scripts/train.sh b/scripts/train.sh index d6a240c..ff97681 100644 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -5,21 +5,37 @@ fi source scripts/base.sh +# =============== Set Device ================ +# -- if GPU memory is enough, recommended training with Single but not Multiple GPU +# -- see issue if error: https://github.com/pytorch/pytorch/issues/110000 + +# -- Multiple GPUs (use all) +# export CUDA_VISIBLE_DEVICES="0,1" +# TRAIN_DEVICE="0,1" + +# -- Multiple GPUs (use only one) +# export CUDA_VISIBLE_DEVICES=0 +# export WORLD_SIZE=1 +# TRAIN_DEVICE=1 + +# -- Single GPU +# TRAIN_DEVICE=0 + +# -- CPU +TRAIN_DEVICE="cpu" + +# ======================================================= + # =============== Set Training Variables ================ -DATASET_DIR=~/data/yolodataset +DATASET_DIR=~/data/drink TRAIN_DATA_DIR=$DATASET_DIR-organized -# if GPU memory is enough, recommended training with Single but not Multiple GPU -TRAIN_DEVICE="0,1" # Multiple GPUs -TRAIN_DEVICE="0" # Single GPU -TRAIN_DEVICE="cpu" # default, MacOS or without GPU - MODEL_NAME=yolov5s PRETRAINED_MODEL=$CACHE_DIR/yolov5/yolov5s.pt -BATCH_SIZE=4 -EPOCHS=4 +BATCH_SIZE=2 +EPOCHS=80 # ======================================================= @@ -44,10 +60,17 @@ python3 train.py \ --data $TRAIN_DATA_DIR/dataset.yaml \ --cfg models/$MODEL_NAME.yaml \ --weights $PRETRAINED_MODEL \ - --epochs $EPOCHS \ - --batch $BATCH_SIZE \ - --device $TRAIN_DEVICE \ - --workers 8 \ - --project $PROJECT_HOME/tmp/runs/train + --epochs $EPOCHS --batch $BATCH_SIZE \ + --device $TRAIN_DEVICE --workers 8 --cache disk \ + --project $PROJECT_HOME/tmp/train + +# python3 train.py \ +# --data $TRAIN_DATA_DIR/dataset.yaml \ +# --cfg models/$MODEL_NAME.yaml \ +# --weights $PRETRAINED_MODEL \ +# --rect --img-size 1280 \ +# --epochs $EPOCHS --batch $BATCH_SIZE \ +# --device $TRAIN_DEVICE --workers 128 --cache ram \ +# --project $PROJECT_HOME/tmp/train # python train.py --data $env:USERPROFILE/data/bottle-organized/dataset.yaml --cfg models/yolov5s.yaml --weights "../../.cache/yolov5/yolov5s.pt" --epochs 10 --batch 4 --device 0 \ No newline at end of file diff --git a/scripts/utils/cuda.py b/scripts/utils/cuda.py new file mode 100644 index 0000000..a5358f0 --- /dev/null +++ b/scripts/utils/cuda.py @@ -0,0 +1,12 @@ +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "1" + +# os.environ["WORLD_SIZE"] = "1" + +import torch + +print(torch.__version__) +print(torch.cuda.is_available()) +print(torch.cuda.device_count()) +print(torch.cuda.current_device()) diff --git a/scripts/utils/python-install-requirements.sh b/scripts/utils/python-install-requirements.sh index 9455580..63f8f86 100644 --- a/scripts/utils/python-install-requirements.sh +++ b/scripts/utils/python-install-requirements.sh @@ -28,12 +28,7 @@ fi # ================== requirements.txt ================== print_info "Installing other requirements..." -mkdir -p $CACHE_DIR/yolov5 -if [ ! -f "$CACHE_DIR/yolov5/requirements.txt" ]; then - cp projects/yolov5/requirements.txt $CACHE_DIR/yolov5/requirements.txt -fi -python3 -m pip install $PIPQ -r $CACHE_DIR/yolov5/requirements.txt -python3 -m pip install $PIPQ -r requirements/requirements.train.txt +python3 -m pip install $PIPQ -r requirements/torch.txt # ================== ONNX Runtime ================== diff --git a/scripts/create-python-env.ps1 b/scripts/win-create-python-env.ps1 similarity index 100% rename from scripts/create-python-env.ps1 rename to scripts/win-create-python-env.ps1 diff --git a/test.py b/test.py new file mode 100644 index 0000000..3fb3c5e --- /dev/null +++ b/test.py @@ -0,0 +1,62 @@ +import os +import argparse +from typing import List + +import yaml + + +class TestArgs: + + @staticmethod + def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "-d", + "--dataset-config", + type=str, + default="~/data/drink-organized/dataset.yaml", + ) + return parser.parse_args() + + def __init__(self) -> None: + # fmt: off + args = self.get_args() + 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}") + # fmt: on + + +def main(): + args = TestArgs() + + dataset_config_file = args.dataset_config_file + with open(dataset_config_file, "r") as f: + data_config = yaml.load(f, Loader=yaml.FullLoader) + print(data_config) + valid_list_file = os.path.join(data_config["path"], data_config["val"]) + train_list_file = os.path.join(data_config["path"], data_config["train"]) + test_list: List[str] = None + + def read_file(file_path: str) -> List[str]: + fl: List[str] = [] + with open(file_path, "r") as f: + for line in f: + fl.append(line.strip()) + return fl + + if os.path.exists(valid_list_file): + test_list = read_file(valid_list_file) + elif os.path.exists(train_list_file): + test_list = read_file(valid_list_file) + else: + raise FileNotFoundError( + f"No valid('{valid_list_file}') or train('{train_list_file}') files found" + ) + + print(test_list) + + + +if __name__ == "__main__": + main() diff --git a/utils/dataset/variables.py b/utils/dataset/variables.py index b5aac8f..47229aa 100644 --- a/utils/dataset/variables.py +++ b/utils/dataset/variables.py @@ -12,4 +12,4 @@ SUPPORTED_IMAGE_TYPES = [ *BASE_SUPPORTED_IMAGE_TYPES, *[str(t).upper() for t in BASE_SUPPORTED_IMAGE_TYPES], -] +] \ No newline at end of file