From b73c268fd0988dcdbeb025569a03447ea7a3822a Mon Sep 17 00:00:00 2001 From: cvhub Date: Sat, 25 May 2024 19:42:45 +0800 Subject: [PATCH] [Enhancement] Support yolov10 model --- anylabeling/app_info.py | 2 +- anylabeling/configs/auto_labeling/models.yaml | 12 ++ .../configs/auto_labeling/yolov10b.yaml | 86 +++++++++ .../configs/auto_labeling/yolov10l.yaml | 86 +++++++++ .../configs/auto_labeling/yolov10m.yaml | 86 +++++++++ .../configs/auto_labeling/yolov10n.yaml | 86 +++++++++ .../configs/auto_labeling/yolov10s.yaml | 86 +++++++++ .../configs/auto_labeling/yolov10x.yaml | 86 +++++++++ .../services/auto_labeling/__base__/yolo.py | 33 +++- .../services/auto_labeling/model_manager.py | 23 +++ anylabeling/services/auto_labeling/yolov10.py | 5 + docs/en/get_started.md | 2 +- docs/en/model_zoo.md | 6 + docs/zh_cn/get_started.md | 2 +- docs/zh_cn/model_zoo.md | 7 +- tools/export_yolov10_onnx.py | 181 ++++++++++++++++++ 16 files changed, 784 insertions(+), 5 deletions(-) create mode 100644 anylabeling/configs/auto_labeling/yolov10b.yaml create mode 100644 anylabeling/configs/auto_labeling/yolov10l.yaml create mode 100644 anylabeling/configs/auto_labeling/yolov10m.yaml create mode 100644 anylabeling/configs/auto_labeling/yolov10n.yaml create mode 100644 anylabeling/configs/auto_labeling/yolov10s.yaml create mode 100644 anylabeling/configs/auto_labeling/yolov10x.yaml create mode 100644 anylabeling/services/auto_labeling/yolov10.py create mode 100644 tools/export_yolov10_onnx.py diff --git a/anylabeling/app_info.py b/anylabeling/app_info.py index 5fe01653..93f830da 100644 --- a/anylabeling/app_info.py +++ b/anylabeling/app_info.py @@ -1,4 +1,4 @@ __appname__ = "X-AnyLabeling" __appdescription__ = "Advanced Auto Labeling Solution with Added Features" -__version__ = "2.3.5" +__version__ = "2.3.6" __preferred_device__ = "CPU" # GPU or CPU diff --git a/anylabeling/configs/auto_labeling/models.yaml b/anylabeling/configs/auto_labeling/models.yaml index 68769f56..82182fb1 100644 --- a/anylabeling/configs/auto_labeling/models.yaml +++ b/anylabeling/configs/auto_labeling/models.yaml @@ -180,6 +180,18 @@ config_file: ":/yolov9c.yaml" - model_name: "yolov9e-r20240224" config_file: ":/yolov9e.yaml" +- model_name: "yolov10b-r20240525" + config_file: ":/yolov10b.yaml" +- model_name: "yolov10l-r20240525" + config_file: ":/yolov10l.yaml" +- model_name: "yolov10m-r20240525" + config_file: ":/yolov10m.yaml" +- model_name: "yolov10n-r20240525" + config_file: ":/yolov10n.yaml" +- model_name: "yolov10s-r20240525" + config_file: ":/yolov10s.yaml" +- model_name: "yolov10x-r20240525" + config_file: ":/yolov10x.yaml" - model_name: "yolow_l-r20240227" config_file: ":/yolow_l.yaml" - model_name: "yolox_l_dwpose_ucoco-r20230820" diff --git a/anylabeling/configs/auto_labeling/yolov10b.yaml b/anylabeling/configs/auto_labeling/yolov10b.yaml new file mode 100644 index 00000000..1260cae4 --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10b.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10b-r20240525 +display_name: YOLOv10b THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10b.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/auto_labeling/yolov10l.yaml b/anylabeling/configs/auto_labeling/yolov10l.yaml new file mode 100644 index 00000000..91a5916b --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10l.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10l-r20240525 +display_name: YOLOv10l THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10l.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/auto_labeling/yolov10m.yaml b/anylabeling/configs/auto_labeling/yolov10m.yaml new file mode 100644 index 00000000..9f69b6e5 --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10m.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10m-r20240525 +display_name: YOLOv10m THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10m.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/auto_labeling/yolov10n.yaml b/anylabeling/configs/auto_labeling/yolov10n.yaml new file mode 100644 index 00000000..c4b9f4ce --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10n.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10n-r20240525 +display_name: YOLOv10n THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10n.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/auto_labeling/yolov10s.yaml b/anylabeling/configs/auto_labeling/yolov10s.yaml new file mode 100644 index 00000000..e44169a4 --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10s.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10s-r20240525 +display_name: YOLOv10s THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10s.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/auto_labeling/yolov10x.yaml b/anylabeling/configs/auto_labeling/yolov10x.yaml new file mode 100644 index 00000000..4d5f7d5f --- /dev/null +++ b/anylabeling/configs/auto_labeling/yolov10x.yaml @@ -0,0 +1,86 @@ +type: yolov10 +name: yolov10x-r20240525 +display_name: YOLOv10x THU-MIG +model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10x.onnx +confidence_threshold: 0.25 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/services/auto_labeling/__base__/yolo.py b/anylabeling/services/auto_labeling/__base__/yolo.py index b6387ce8..e5745cb6 100644 --- a/anylabeling/services/auto_labeling/__base__/yolo.py +++ b/anylabeling/services/auto_labeling/__base__/yolo.py @@ -128,6 +128,7 @@ def __init__(self, model_config, on_message) -> None: "yolov7", "yolov8", "yolov9", + "yolov10", "gold_yolo", ]: self.task = "det" @@ -222,7 +223,12 @@ def postprocess(self, preds): multi_label=False, nc=self.nc, ) - + elif self.model_type == "yolov10": + p = self.postprocess_v10( + preds[0][0], + conf_thres=self.conf_thres, + classes=self.filter_classes + ) masks = None img_shape = (self.img_height, self.img_width) if self.task == "seg": @@ -470,6 +476,13 @@ def process_mask(self, protos, masks_in, bboxes, shape, upsample=False): return masks + def postprocess_v10(self, prediction, task="det", conf_thres=0.25, classes=None): + x = prediction[prediction[:, 4] >= conf_thres] + x[:, -1] = x[:, -1].astype(int) + if classes is not None: + x = x[np.isin(x[:, -1], classes)] + return [x] + @staticmethod def crop_mask_np(masks, boxes): """ @@ -489,5 +502,23 @@ def crop_mask_np(masks, boxes): return masks * ((r >= x1) & (r < x2) & (c >= y1) & (c < y2)) + @staticmethod + def rescale_coords_v10(boxes, image_shape, input_shape): + image_height, image_width = image_shape + input_height, input_width = input_shape + + scale = min(input_width / image_width, input_height / image_height) + + pad_w = (input_width - image_width * scale) / 2 + pad_h = (input_height - image_height * scale) / 2 + + boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale + boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale + + boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, image_width) + boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, image_height) + + return boxes + def unload(self): del self.net diff --git a/anylabeling/services/auto_labeling/model_manager.py b/anylabeling/services/auto_labeling/model_manager.py index 22da178f..35d3e40f 100644 --- a/anylabeling/services/auto_labeling/model_manager.py +++ b/anylabeling/services/auto_labeling/model_manager.py @@ -207,6 +207,7 @@ def load_custom_model(self, config_file): "depth_anything", "yolov9", "yolow", + "yolov10", ] ): self.new_model_status.emit( @@ -440,6 +441,28 @@ def _load_model(self, model_id): ) ) return + elif model_config["type"] == "yolov10": + from .yolov10 import YOLOv10 + + try: + model_config["model"] = YOLOv10( + model_config, on_message=self.new_model_status.emit + ) + self.auto_segmentation_model_unselected.emit() + except Exception as e: # noqa + self.new_model_status.emit( + self.tr( + "Error in loading model: {error_message}".format( + error_message=str(e) + ) + ) + ) + print( + "Error in loading model: {error_message}".format( + error_message=str(e) + ) + ) + return elif model_config["type"] == "yolow": from .yolow import YOLOW diff --git a/anylabeling/services/auto_labeling/yolov10.py b/anylabeling/services/auto_labeling/yolov10.py new file mode 100644 index 00000000..5be26a7b --- /dev/null +++ b/anylabeling/services/auto_labeling/yolov10.py @@ -0,0 +1,5 @@ +from .__base__.yolo import YOLO + + +class YOLOv10(YOLO): + pass diff --git a/docs/en/get_started.md b/docs/en/get_started.md index 7a8e0278..624d0292 100644 --- a/docs/en/get_started.md +++ b/docs/en/get_started.md @@ -61,7 +61,7 @@ To ensure access to the latest features and stable performance, it is highly rec Running in the GUI environment is convenient but may have limitations compared to running from the source code. Consider the pros and cons based on your specific needs and preferences. -Download Link: [Release](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v2.3.5) | [Baidu Disk](https://pan.baidu.com/s/1n3vfXZpBUG9s12NZjMxhlw?pwd=odyi) +Download Link: [Release](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v2.3.6) | [Baidu Disk](https://pan.baidu.com/s/1rtw_UY_qTOopKNqFfXEfzA?pwd=itwe) Note: - For MacOS: diff --git a/docs/en/model_zoo.md b/docs/en/model_zoo.md index e5619ecb..b2154717 100644 --- a/docs/en/model_zoo.md +++ b/docs/en/model_zoo.md @@ -90,6 +90,12 @@ | yolov9e.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [yolov9e.yaml](../../anylabeling/configs/auto_labeling/yolov9e.yaml) | 265.43MB | [baidu](https://pan.baidu.com/s/1oJJQ1L5UfKi43kT96tauSA?pwd=vpa1) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/yolov9e.onnx) | | gelan-c.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [gelan-c.yaml](../../anylabeling/configs/auto_labeling/yolov9_gelan_c.yaml) | 97.43MB | [baidu](https://pan.baidu.com/s/1cM0Tc056ICuA5jvesacCng?pwd=whb3) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/gelan-c.onnx) | | gelan-e.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [gelan-e.yaml](../../anylabeling/configs/auto_labeling/yolov9_gelan_e.yaml) | 221.94MB | [baidu](https://pan.baidu.com/s/1Iy0w5Cgga0-GmOLNSvYflg?pwd=s1p2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/gelan-c.onnx) | +| yolov10n.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10n.yaml](../../anylabeling/configs/auto_labeling/yolov10n.yaml) | 8.98MB | [baidu](https://pan.baidu.com/s/12h4-UBbCdbgAYCTBroZh_g?pwd=mo5k) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10n.onnx) | +| yolov10s.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10s.yaml](../../anylabeling/configs/auto_labeling/yolov10s.yaml) | 27.86MB | [baidu](https://pan.baidu.com/s/1cRZk-eSdaqmsQLdKqX_-Iw?pwd=pwa2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10s.onnx) | +| yolov10m.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10m.yaml](../../anylabeling/configs/auto_labeling/yolov10m.yaml) | 58.81MB | [baidu](https://pan.baidu.com/s/19iKuQ9DhJqzPqN7zBuZQ9Q?pwd=3hn2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10m.onnx) | +| yolov10b.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10b.yaml](../../anylabeling/configs/auto_labeling/yolov10b.yaml) | 72.95MB | [baidu](https://pan.baidu.com/s/15rabV7mFyy6HQVgTE1Wv6Q?pwd=dol1) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10b.onnx) | +| yolov10l.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10l.yaml](../../anylabeling/configs/auto_labeling/yolov10l.yaml) | 93.20MB | [baidu](https://pan.baidu.com/s/1dJGVhtPRC31ssbnhQWpztw?pwd=13ax) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10l.onnx) | +| yolov10x.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10x.yaml](../../anylabeling/configs/auto_labeling/yolov10x.yaml) | 112.68MB | [baidu](https://pan.baidu.com/s/1o6rCcq-NyIBe4ftcOJlKzQ?pwd=c73x) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10x.onnx) | - Oriented Bounding Box diff --git a/docs/zh_cn/get_started.md b/docs/zh_cn/get_started.md index a2acb404..694e55df 100644 --- a/docs/zh_cn/get_started.md +++ b/docs/zh_cn/get_started.md @@ -74,7 +74,7 @@ python anylabeling/app.py 为了在选择运行方式时能够更好地权衡利弊，建议用户根据具体需求和偏好，灵活选择源码运行或GUI环境运行，以达到最佳的使用体验。 -下载链接：[Release](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v5) | [百度网盘](https://pan.baidu.com/s/1n3vfXZpBUG9s12NZjMxhlw?pwd=odyi) +下载链接：[Release](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v2.3.6) | [百度网盘](https://pan.baidu.com/s/1rtw_UY_qTOopKNqFfXEfzA?pwd=itwe) ### 文件导入 diff --git a/docs/zh_cn/model_zoo.md b/docs/zh_cn/model_zoo.md index 6e087ccd..c5bc4051 100644 --- a/docs/zh_cn/model_zoo.md +++ b/docs/zh_cn/model_zoo.md @@ -88,7 +88,12 @@ | yolov9e.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [yolov9e.yaml](../../anylabeling/configs/auto_labeling/yolov9e.yaml) | 265.43MB | [百度网盘](https://pan.baidu.com/s/1oJJQ1L5UfKi43kT96tauSA?pwd=vpa1) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/yolov9e.onnx) | | gelan-c.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [gelan-c.yaml](../../anylabeling/configs/auto_labeling/yolov9_gelan_c.yaml) | 97.43MB | [百度网盘](https://pan.baidu.com/s/1cM0Tc056ICuA5jvesacCng?pwd=whb3) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/gelan-c.onnx) | | gelan-e.onnx | [YOLOv9](https://github.com/WongKinYiu/yolov9)-COCO | [gelan-e.yaml](../../anylabeling/configs/auto_labeling/yolov9_gelan_e.yaml) | 221.94MB | [百度网盘](https://pan.baidu.com/s/1Iy0w5Cgga0-GmOLNSvYflg?pwd=s1p2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.2/gelan-c.onnx) | - +| yolov10n.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10n.yaml](../../anylabeling/configs/auto_labeling/yolov10n.yaml) | 8.98MB | [百度网盘](https://pan.baidu.com/s/12h4-UBbCdbgAYCTBroZh_g?pwd=mo5k) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10n.onnx) | +| yolov10s.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10s.yaml](../../anylabeling/configs/auto_labeling/yolov10s.yaml) | 27.86MB | [百度网盘](https://pan.baidu.com/s/1cRZk-eSdaqmsQLdKqX_-Iw?pwd=pwa2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10s.onnx) | +| yolov10m.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10m.yaml](../../anylabeling/configs/auto_labeling/yolov10m.yaml) | 58.81MB | [百度网盘](https://pan.baidu.com/s/19iKuQ9DhJqzPqN7zBuZQ9Q?pwd=3hn2) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10m.onnx) | +| yolov10b.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10b.yaml](../../anylabeling/configs/auto_labeling/yolov10b.yaml) | 72.95MB | [百度网盘](https://pan.baidu.com/s/15rabV7mFyy6HQVgTE1Wv6Q?pwd=dol1) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10b.onnx) | +| yolov10l.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10l.yaml](../../anylabeling/configs/auto_labeling/yolov10l.yaml) | 93.20MB | [百度网盘](https://pan.baidu.com/s/1dJGVhtPRC31ssbnhQWpztw?pwd=13ax) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10l.onnx) | +| yolov10x.onnx | [YOLOv10](https://github.com/THU-MIG/yolov10)-COCO | [yolov10x.yaml](../../anylabeling/configs/auto_labeling/yolov10x.yaml) | 112.68MB | [百度网盘](https://pan.baidu.com/s/1o6rCcq-NyIBe4ftcOJlKzQ?pwd=c73x) \| [github](https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.3.6/yolov10x.onnx) | - 有向目标检测 diff --git a/tools/export_yolov10_onnx.py b/tools/export_yolov10_onnx.py new file mode 100644 index 00000000..18039b42 --- /dev/null +++ b/tools/export_yolov10_onnx.py @@ -0,0 +1,181 @@ +import os +import cv2 +import random +import numpy as np +import onnxruntime as ort + +""" +The onnxruntime demo of the YOLOv10 +Written by Wei Wang (CVHub) + Usage: + 1. git clone https://github.com/THU-MIG/yolov10 + 2. cd yolov10 and install the package + 3. export PYTHONPATH=/path/to/yolov10 + 4. Download the corresponding weights + 5. Export the model to onnx format: + ```bash + yolo export model=/path/to/yolov10n/s/m/b/l/x.pt format=onnx opset=13 simplify + ``` + 7. Modified the paramters and run this script + ```bash + python ${export_yolov10_onnx.py} + ``` +""" + + +random.seed(10086) +CLASS_NAMES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', + 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', + 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', + 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', + 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', + 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', + 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', + 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', + 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', + 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', + 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', + 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', + 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', + 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush') +CLASS_COLORS = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(CLASS_NAMES))] + + +def letterbox( + im, + new_shape, + color=(114, 114, 114), + auto=False, + scaleFill=False, + scaleup=True, + stride=32, +): + """ + Resize and pad image while meeting stride-multiple constraints + Returns: + im (array): (height, width, 3) + ratio (array): [w_ratio, h_ratio] + (dw, dh) (array): [w_padding h_padding] + """ + shape = im.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): # [h_rect, w_rect] + new_shape = (new_shape, new_shape) + + # Scale ratio (new / old) + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scaleup: # only scale down, do not scale up (for better val mAP) + r = min(r, 1.0) + + # Compute padding + ratio = r, r # wh ratios + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) # w h + dw, dh = ( + new_shape[1] - new_unpad[0], + new_shape[0] - new_unpad[1], + ) # wh padding + + if auto: # minimum rectangle + dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding + elif scaleFill: # stretch + dw, dh = 0.0, 0.0 + new_unpad = (new_shape[1], new_shape[0]) # [w h] + ratio = ( + new_shape[1] / shape[1], + new_shape[0] / shape[0], + ) # [w_ratio, h_ratio] + + dw /= 2 # divide padding into 2 sides + dh /= 2 + if shape[::-1] != new_unpad: # resize + im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + im = cv2.copyMakeBorder( + im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color + ) + return im, ratio, (dw, dh) + +def rescale_coords(boxes, image_shape, input_shape): + image_height, image_width = image_shape + input_height, input_width = input_shape + + scale = min(input_width / image_width, input_height / image_height) + + pad_w = (input_width - image_width * scale) / 2 + pad_h = (input_height - image_height * scale) / 2 + + boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale + boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale + + boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, image_width) + boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, image_height) + + return boxes.astype(int) + +def preprocess(image, input_shape): + # Resize + input_img = letterbox(image, input_shape)[0] + # Transpose + input_img = input_img[..., ::-1].transpose(2, 0, 1) + # Expand + input_img = input_img[np.newaxis, :, :, :].astype(np.float32) + # Contiguous + input_img = np.ascontiguousarray(input_img) + # Norm + blob = input_img / 255.0 + return blob + +def postprocess(outs, conf_thres, image_shape, input_shape): + + # Filtered by conf + outs = outs[outs[:, 4] >= conf_thres] + + # Extract + boxes = outs[:, :4] + scores = outs[:, -2] + labels = outs[:, -1].astype(int) + + # Rescale + boxes = rescale_coords(boxes, image_shape, input_shape) + + return boxes, scores, labels + + +def main(): + + conf_thres = 0.25 + input_shape = (640, 640) + image_path = '/path/to/bus.jpg' + save_path = '' + model_path = '/path/to/yolov10n/s/m/b/l/x.onnx' + + ort_model = ort.InferenceSession(model_path) + + # Preprocess + im0 = cv2.imread(image_path) + image_shape = im0.shape[:2] + blob = preprocess(im0, input_shape) + + # Inference + outs = ort_model.run(None, {'images': blob})[0][0] + + # Postprocess + boxes, scores, labels = postprocess(outs, conf_thres, image_shape, input_shape) + + # 保存结果 + for label, score, box in zip(labels, scores, boxes): + label_text = f'{CLASS_NAMES[label]}: {score:.2f}' + cv2.rectangle(im0, (box[0], box[1]), (box[2], box[3]), CLASS_COLORS[label], 2) + cv2.putText(im0, label_text, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) + + try: + cv2.imshow('image', im0) + cv2.waitKey(0) + except: + if not os.path.exists(os.path.dirname(save_path)): + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, im0) + + +if __name__ == '__main__': + main() \ No newline at end of file