From f3508150d5b9ef5ffd98cfed597bcef3b66743d1 Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Mon, 9 Mar 2026 21:33:43 +0800 Subject: [PATCH 1/6] =?UTF-8?q?feat:=20MindNLP=E5=85=BC=E5=AE=B9Ultralytic?= =?UTF-8?q?s=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/mindnlp/models/ultralytics/__init__.py | 0 src/mindnlp/models/ultralytics/cfg/hyp.yaml | 60 ++ .../ultralytics/cfg/models/11/yolo11-cls.yaml | 33 ++ .../cfg/models/11/yolo11-pose.yaml | 51 ++ .../ultralytics/cfg/models/11/yolo11-seg.yaml | 50 ++ .../ultralytics/cfg/models/11/yolo11.yaml | 50 ++ .../models/ultralytics/configuration_yolo.py | 83 +++ .../models/ultralytics/engine/__init__.py | 0 .../models/ultralytics/engine/predictor.py | 148 +++++ .../models/ultralytics/engine/trainer.py | 171 ++++++ .../models/ultralytics/engine/validator.py | 123 ++++ .../examples/yolo/classify/inference.py | 84 +++ .../examples/yolo/classify/run_train.py | 37 ++ .../examples/yolo/classify/run_val.py | 77 +++ .../examples/yolo/detect/inference.py | 57 ++ .../examples/yolo/detect/run_train.py | 47 ++ .../examples/yolo/detect/run_val.py | 78 +++ .../examples/yolo/pose/inference.py | 71 +++ .../examples/yolo/pose/run_train.py | 36 ++ .../ultralytics/examples/yolo/pose/run_val.py | 77 +++ .../examples/yolo/segment/inference.py | 59 ++ .../examples/yolo/segment/run_train.py | 43 ++ .../examples/yolo/segment/run_val.py | 88 +++ .../models/ultralytics/modeling_yolo.py | 238 ++++++++ .../models/ultralytics/models/__init__.py | 0 .../models/yolo/classify/__init__.py | 0 .../models/yolo/classify/predict.py | 121 ++++ .../ultralytics/models/yolo/classify/train.py | 69 +++ .../ultralytics/models/yolo/classify/val.py | 88 +++ .../models/yolo/detect/__init__.py | 0 .../ultralytics/models/yolo/detect/predict.py | 92 +++ .../ultralytics/models/yolo/detect/train.py | 121 ++++ .../ultralytics/models/yolo/detect/val.py | 112 ++++ .../ultralytics/models/yolo/pose/__init__.py | 0 .../ultralytics/models/yolo/pose/predict.py | 108 ++++ .../ultralytics/models/yolo/pose/train.py | 118 ++++ .../ultralytics/models/yolo/pose/val.py | 219 +++++++ .../models/yolo/segment/__init__.py | 0 .../models/yolo/segment/predict.py | 152 +++++ .../ultralytics/models/yolo/segment/train.py | 103 ++++ .../ultralytics/models/yolo/segment/val.py | 205 +++++++ src/mindnlp/models/ultralytics/modules.py | 469 +++++++++++++++ src/mindnlp/models/ultralytics/readme.md | 79 +++ .../models/ultralytics/tools/convert.py | 158 ++++++ .../models/ultralytics/utils/__init__.py | 0 src/mindnlp/models/ultralytics/utils/ema.py | 46 ++ src/mindnlp/models/ultralytics/utils/loss.py | 535 ++++++++++++++++++ .../models/ultralytics/utils/metrics.py | 315 +++++++++++ src/mindnlp/models/ultralytics/utils/ops.py | 401 +++++++++++++ .../models/ultralytics/utils/optimizer.py | 89 +++ src/mindnlp/models/ultralytics/utils/tal.py | 140 +++++ 51 files changed, 5501 insertions(+) create mode 100644 src/mindnlp/models/ultralytics/__init__.py create mode 100644 src/mindnlp/models/ultralytics/cfg/hyp.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml create mode 100644 src/mindnlp/models/ultralytics/configuration_yolo.py create mode 100644 src/mindnlp/models/ultralytics/engine/__init__.py create mode 100644 src/mindnlp/models/ultralytics/engine/predictor.py create mode 100644 src/mindnlp/models/ultralytics/engine/trainer.py create mode 100644 src/mindnlp/models/ultralytics/engine/validator.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py create mode 100644 src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py create mode 100644 src/mindnlp/models/ultralytics/modeling_yolo.py create mode 100644 src/mindnlp/models/ultralytics/models/__init__.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/classify/__init__.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/classify/predict.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/classify/train.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/classify/val.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/detect/__init__.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/detect/predict.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/detect/train.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/detect/val.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/pose/__init__.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/pose/predict.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/pose/train.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/pose/val.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/segment/predict.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/segment/train.py create mode 100644 src/mindnlp/models/ultralytics/models/yolo/segment/val.py create mode 100644 src/mindnlp/models/ultralytics/modules.py create mode 100644 src/mindnlp/models/ultralytics/readme.md create mode 100644 src/mindnlp/models/ultralytics/tools/convert.py create mode 100644 src/mindnlp/models/ultralytics/utils/__init__.py create mode 100644 src/mindnlp/models/ultralytics/utils/ema.py create mode 100644 src/mindnlp/models/ultralytics/utils/loss.py create mode 100644 src/mindnlp/models/ultralytics/utils/metrics.py create mode 100644 src/mindnlp/models/ultralytics/utils/ops.py create mode 100644 src/mindnlp/models/ultralytics/utils/optimizer.py create mode 100644 src/mindnlp/models/ultralytics/utils/tal.py diff --git a/src/mindnlp/models/ultralytics/__init__.py b/src/mindnlp/models/ultralytics/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/cfg/hyp.yaml b/src/mindnlp/models/ultralytics/cfg/hyp.yaml new file mode 100644 index 000000000..efeb0c145 --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/hyp.yaml @@ -0,0 +1,60 @@ +# ========================================== +# YOLO11 统一超参数配置 (适用于 Classify, Detect, Segment, Pose) +# ========================================== + +# 1. 优化器与学习率 (Optimizer & LR) + +optimizer: 'SGD' # 优化器选择 +lr0: 0.01 # 初始学习率（微调pose任务,使用coco8-pose数据集时,推荐使用0.001） +lrf: 0.01 # 最终学习率比例 (lr0 * lrf = 0.0001) +momentum: 0.937 # 动量 (SGD momentum / Adam beta1) +weight_decay: 0.0005 # 权重衰减系数 (L2 正则化) +warmup_epochs: 3.0 # 预热轮数 (前几个 Epoch 学习率从低爬升) +warmup_momentum: 0.8 # 预热阶段的初始动量 +warmup_bias_lr: 0.1 # 预热阶段偏置项(bias)的学习率 + + +# 2. 任务损失权重 (Loss Weights) +# 说明：各个任务会自动读取自己需要的权重，无关的会被忽略 + +box: 7.5 # 边界框回归损失 (GIoU/CIoU) +cls: 0.5 # 类别分类损失 (BCE / CE) +dfl: 1.5 # 分布焦点损失 (DFL) +pose: 12.0 # 姿态估计：关键点坐标损失 +kobj: 1.0 # 姿态估计：关键点可见度损失 +seg: 2.5 # 实例分割：掩码 BCE 损失 +label_smoothing: 0.0 # 标签平滑系数 (防止过拟合，通常为0) + + +# 3. 基础数据增强 (Data Augmentation) + +hsv_h: 0.015 # 图像 HSV 色调增强比例 +hsv_s: 0.7 # 图像 HSV 饱和度增强比例 +hsv_v: 0.4 # 图像 HSV 明度增强比例 +degrees: 0.0 # 图像旋转角度 (+/- deg) +translate: 0.1 # 图像平移比例 (+/- fraction) +scale: 0.5 # 图像缩放比例 (+/- gain) +shear: 0.0 # 图像剪切角度 (+/- deg) +perspective: 0.0 # 透视变换系数 (0-0.001) +flipud: 0.0 # 图像上下翻转概率 +fliplr: 0.5 # 图像左右翻转概率 (分类和检测常用) + + +# 4. 高级图像拼接增强 (Mosaic & MixUp) +# 注意：分类任务通常不使用马赛克增强，代码里需做判断 + +mosaic: 1.0 # 马赛克数据增强概率 (检测/分割/姿态 核心) +mixup: 0.0 # MixUp 增强概率 +copy_paste: 0.0 # Copy-Paste 分割增强概率 + + +# 5. 验证与推理超参数 (Val & Inference Settings) +# 注意：在验证(Val)计算 mAP 时，系统会自动将 conf 降为 0.001； +# 下面填写的 conf 主要供常规推理(Predict)或最终展示使用。 + +conf: 0.25 # 目标置信度阈值 (预测时使用 0.25，验证算 mAP 时代码内部会强转为 0.001) +iou: 0.6 # NMS (非极大值抑制) 的 IoU 阈值 +max_det: 300 # 每张图像最多保留的预测框数量 +half: False # 是否使用 FP16 半精度推理 (Ascend/GPU 推荐开启以提速) +dnn: False # 是否使用 OpenCV DNN 进行 ONNX 推理 +plots: True # 是否在验证期间保存画好框的结果图和各种评估图表 \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml new file mode 100644 index 000000000..753e27b6d --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml @@ -0,0 +1,33 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +# Ultralytics YOLO11-cls image classification model +# Model docs: https://docs.ultralytics.com/models/yolo11 +# Task docs: https://docs.ultralytics.com/tasks/classify + +# Parameters +nc: 1000 # number of classes +scales: # model compound scaling constants, i.e. 'model=yolo11n-cls.yaml' will call yolo11-cls.yaml with scale 'n' + # [depth, width, max_channels] + n: [0.50, 0.25, 1024] # summary: 86 layers, 1633584 parameters, 1633584 gradients, 0.5 GFLOPs + s: [0.50, 0.50, 1024] # summary: 86 layers, 5545488 parameters, 5545488 gradients, 1.6 GFLOPs + m: [0.50, 1.00, 512] # summary: 106 layers, 10455696 parameters, 10455696 gradients, 5.0 GFLOPs + l: [1.00, 1.00, 512] # summary: 176 layers, 12937104 parameters, 12937104 gradients, 6.2 GFLOPs + x: [1.00, 1.50, 512] # summary: 176 layers, 28458544 parameters, 28458544 gradients, 13.7 GFLOPs + +# YOLO11n backbone +backbone: + # [from, repeats, module, args] + - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 + - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 + - [-1, 2, C3k2, [256, False, 0.25]] + - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 + - [-1, 2, C3k2, [512, False, 0.25]] + - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 + - [-1, 2, C3k2, [512, True]] + - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 + - [-1, 2, C3k2, [1024, True]] + - [-1, 2, C2PSA, [1024]] # 9 + +# YOLO11n head +head: + - [-1, 1, Classify, [nc]] # Classify diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml new file mode 100644 index 000000000..c687f8fb5 --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml @@ -0,0 +1,51 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +# Ultralytics YOLO11-pose keypoints/pose estimation model with P3/8 - P5/32 outputs +# Model docs: https://docs.ultralytics.com/models/yolo11 +# Task docs: https://docs.ultralytics.com/tasks/pose + +# Parameters +nc: 80 # number of classes +kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) +scales: # model compound scaling constants, i.e. 'model=yolo11n-pose.yaml' will call yolo11-pose.yaml with scale 'n' + # [depth, width, max_channels] + n: [0.50, 0.25, 1024] # summary: 196 layers, 2908507 parameters, 2908491 gradients, 7.7 GFLOPs + s: [0.50, 0.50, 1024] # summary: 196 layers, 9948811 parameters, 9948795 gradients, 23.5 GFLOPs + m: [0.50, 1.00, 512] # summary: 246 layers, 20973273 parameters, 20973257 gradients, 72.3 GFLOPs + l: [1.00, 1.00, 512] # summary: 372 layers, 26230745 parameters, 26230729 gradients, 91.4 GFLOPs + x: [1.00, 1.50, 512] # summary: 372 layers, 58889881 parameters, 58889865 gradients, 204.3 GFLOPs + +# YOLO11n backbone +backbone: + # [from, repeats, module, args] + - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 + - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 + - [-1, 2, C3k2, [256, False, 0.25]] + - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 + - [-1, 2, C3k2, [512, False, 0.25]] + - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 + - [-1, 2, C3k2, [512, True]] + - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 + - [-1, 2, C3k2, [1024, True]] + - [-1, 1, SPPF, [1024, 5]] # 9 + - [-1, 2, C2PSA, [1024]] # 10 + +# YOLO11n head +head: + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 6], 1, Concat, [1]] # cat backbone P4 + - [-1, 2, C3k2, [512, False]] # 13 + + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 4], 1, Concat, [1]] # cat backbone P3 + - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small) + + - [-1, 1, Conv, [256, 3, 2]] + - [[-1, 13], 1, Concat, [1]] # cat head P4 + - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium) + + - [-1, 1, Conv, [512, 3, 2]] + - [[-1, 10], 1, Concat, [1]] # cat head P5 + - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large) + + - [[16, 19, 22], 1, Pose, [nc, kpt_shape]] # Detect(P3, P4, P5) diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml new file mode 100644 index 000000000..1186666c6 --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml @@ -0,0 +1,50 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +# Ultralytics YOLO11-seg instance segmentation model with P3/8 - P5/32 outputs +# Model docs: https://docs.ultralytics.com/models/yolo11 +# Task docs: https://docs.ultralytics.com/tasks/segment + +# Parameters +nc: 80 # number of classes +scales: # model compound scaling constants, i.e. 'model=yolo11n-seg.yaml' will call yolo11-seg.yaml with scale 'n' + # [depth, width, max_channels] + n: [0.50, 0.25, 1024] # summary: 203 layers, 2876848 parameters, 2876832 gradients, 10.5 GFLOPs + s: [0.50, 0.50, 1024] # summary: 203 layers, 10113248 parameters, 10113232 gradients, 35.8 GFLOPs + m: [0.50, 1.00, 512] # summary: 253 layers, 22420896 parameters, 22420880 gradients, 123.9 GFLOPs + l: [1.00, 1.00, 512] # summary: 379 layers, 27678368 parameters, 27678352 gradients, 143.0 GFLOPs + x: [1.00, 1.50, 512] # summary: 379 layers, 62142656 parameters, 62142640 gradients, 320.2 GFLOPs + +# YOLO11n backbone +backbone: + # [from, repeats, module, args] + - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 + - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 + - [-1, 2, C3k2, [256, False, 0.25]] + - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 + - [-1, 2, C3k2, [512, False, 0.25]] + - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 + - [-1, 2, C3k2, [512, True]] + - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 + - [-1, 2, C3k2, [1024, True]] + - [-1, 1, SPPF, [1024, 5]] # 9 + - [-1, 2, C2PSA, [1024]] # 10 + +# YOLO11n head +head: + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 6], 1, Concat, [1]] # cat backbone P4 + - [-1, 2, C3k2, [512, False]] # 13 + + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 4], 1, Concat, [1]] # cat backbone P3 + - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small) + + - [-1, 1, Conv, [256, 3, 2]] + - [[-1, 13], 1, Concat, [1]] # cat head P4 + - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium) + + - [-1, 1, Conv, [512, 3, 2]] + - [[-1, 10], 1, Concat, [1]] # cat head P5 + - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large) + + - [[16, 19, 22], 1, Segment, [nc, 32, 256]] # Detect(P3, P4, P5) diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml new file mode 100644 index 000000000..c90c44485 --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml @@ -0,0 +1,50 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +# Ultralytics YOLO11 object detection model with P3/8 - P5/32 outputs +# Model docs: https://docs.ultralytics.com/models/yolo11 +# Task docs: https://docs.ultralytics.com/tasks/detect + +# Parameters +nc: 80 # number of classes +scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n' + # [depth, width, max_channels] + n: [0.50, 0.25, 1024] # summary: 181 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs + s: [0.50, 0.50, 1024] # summary: 181 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs + m: [0.50, 1.00, 512] # summary: 231 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs + l: [1.00, 1.00, 512] # summary: 357 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs + x: [1.00, 1.50, 512] # summary: 357 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs + +# YOLO11n backbone +backbone: + # [from, repeats, module, args] + - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 + - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 + - [-1, 2, C3k2, [256, False, 0.25]] + - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 + - [-1, 2, C3k2, [512, False, 0.25]] + - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 + - [-1, 2, C3k2, [512, True]] + - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 + - [-1, 2, C3k2, [1024, True]] + - [-1, 1, SPPF, [1024, 5]] # 9 + - [-1, 2, C2PSA, [1024]] # 10 + +# YOLO11n head +head: + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 6], 1, Concat, [1]] # cat backbone P4 + - [-1, 2, C3k2, [512, False]] # 13 + + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 4], 1, Concat, [1]] # cat backbone P3 + - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small) + + - [-1, 1, Conv, [256, 3, 2]] + - [[-1, 13], 1, Concat, [1]] # cat head P4 + - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium) + + - [-1, 1, Conv, [512, 3, 2]] + - [[-1, 10], 1, Concat, [1]] # cat head P5 + - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large) + + - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5) diff --git a/src/mindnlp/models/ultralytics/configuration_yolo.py b/src/mindnlp/models/ultralytics/configuration_yolo.py new file mode 100644 index 000000000..f8cf92f29 --- /dev/null +++ b/src/mindnlp/models/ultralytics/configuration_yolo.py @@ -0,0 +1,83 @@ +import os +import yaml + +#from mindnlp.models.utils import PretrainedConfig +# 尝试获取 MindNLP 基类，增强套件兼容性 + +try: + + from mindnlp.models.utils import PretrainedConfig + +except Exception: + + # 兼容性防御：如果环境未完全安装，使用基础 Config 类防止崩溃 + + class PretrainedConfig: + + def __init__(self, **kwargs): + + for k, v in kwargs.items(): + + setattr(self, k, v) + +class YOLOConfig(PretrainedConfig): + """ + YOLO11 全任务通用配置类 + 支持：分类 (Classify)、检测 (Detect)、分割 (Segment)、姿态 (Pose) + """ + model_type = "yolo11" + + def __init__( + self, + yaml_path=None, + scale='n', + nc=80, + kpt_shape=None, + reg_max=16, + nm=32, + npr=256, + **kwargs + ): + # 修复 Python 可变默认参数陷阱 + if kpt_shape is None: + kpt_shape = [17, 3] + + self.yaml_path = yaml_path + self.scale = scale + self.reg_max = reg_max + self.nm = nm + self.npr = npr + self.yaml_dict = {} + + # 先设置基础兜底参数 + self.nc = nc + self.kpt_shape = kpt_shape + self.depth_multiple = 1.0 + self.width_multiple = 1.0 + self.max_channels = 1024 + self.backbone = [] + self.head = [] + + # 读取并解析 YAML + if yaml_path: + if os.path.exists(yaml_path): + with open(yaml_path, 'r', encoding='utf-8') as f: + self.yaml_dict = yaml.safe_load(f) + else: + raise FileNotFoundError(f"错误：找不到 YAML 配置文件 '{yaml_path}'。请检查相对路径！") + + # 从 YAML 中覆盖参数 + if self.yaml_dict: + self.nc = self.yaml_dict.get('nc', self.nc) + self.kpt_shape = self.yaml_dict.get('kpt_shape', self.kpt_shape) + self.backbone = self.yaml_dict.get('backbone', []) + self.head = self.yaml_dict.get('head', []) + + # 根据 scale 动态解析 depth, width, max_channels + scales = self.yaml_dict.get('scales', {}) + if scale in scales: + self.depth_multiple, self.width_multiple, self.max_channels = scales[scale] + else: + print(f"警告: YAML 中未找到 scale '{scale}'，将使用默认值 1.0") + + super().__init__(**kwargs) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/engine/__init__.py b/src/mindnlp/models/ultralytics/engine/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/engine/predictor.py b/src/mindnlp/models/ultralytics/engine/predictor.py new file mode 100644 index 000000000..749296837 --- /dev/null +++ b/src/mindnlp/models/ultralytics/engine/predictor.py @@ -0,0 +1,148 @@ +import os +import time +import logging +import yaml +import numpy as np +import cv2 +import mindspore as ms +from mindspore import Tensor + +# 统一日志配置 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +LOGGER = logging.getLogger(__name__) + +class BasePredictor: + """ + MindNLP 推理引擎基类 + 负责端到端的推理：数据加载 -> 预处理 -> 前向推理 -> 后处理 + """ + + def __init__(self, cfg=None): + """ + 初始化预测器配置 + """ + self.cfg = cfg or {} + + # 解析超参数配置 (支持从 hyp.yaml 读取 conf, iou 等推理参数) + self.hyp = {} + if hasattr(self.cfg, 'hyp') and self.cfg.hyp and os.path.exists(self.cfg.hyp): + with open(self.cfg.hyp, "r", encoding="utf-8") as f: + self.hyp = yaml.safe_load(f) + + self.conf_thres = self.hyp.get('conf', getattr(self.cfg, 'conf', 0.25)) + self.iou_thres = self.hyp.get('iou', getattr(self.cfg, 'iou', 0.45)) + self.imgsz = getattr(self.cfg, 'imgsz', 640) + + self.model = None + self.dataset = [] + self.results = [] + + def setup_model(self, model, ckpt_path=None): + """加载模型权重并设置推理模式""" + self.model = model + if ckpt_path and os.path.exists(ckpt_path): + param_dict = ms.load_checkpoint(ckpt_path) + ms.load_param_into_net(self.model, param_dict, strict_load=False) + LOGGER.info(f"成功加载检查点权重: {ckpt_path}") + + # 锁定为评估模式，冻结 BatchNorm 等动态层 + self.model.set_train(False) + + def setup_source(self, source): + """解析输入数据源 (支持单张图像或目录遍历)""" + self.dataset = [] + if isinstance(source, str): + if os.path.isdir(source): + for f in os.listdir(source): + if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')): + self.dataset.append(os.path.join(source, f)) + elif os.path.isfile(source): + self.dataset.append(source) + elif isinstance(source, np.ndarray): + self.dataset.append(source) + + if not self.dataset: + raise ValueError(f"[ERROR] 无法从指定数据源解析到有效图像: {source}") + + def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)): + """ + 图像仿射变换：等比例缩放并对边缘进行 Padding 填充 + 返回处理后的图像以及变换参数 (供坐标还原使用) + """ + shape = img.shape[:2] + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] + dw, dh = dw / 2, dh / 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, 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)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + + return img, (r, dw, dh) + + def preprocess(self, img_source): + """ + 基类图像预处理流水线：BGR -> RGB -> LetterBox -> CHW -> Normalize -> Tensor + """ + if isinstance(img_source, str): + orig_img = cv2.imread(img_source) + else: + orig_img = img_source.copy() + + img, (ratio, pad_w, pad_h) = self.letterbox(orig_img, new_shape=(self.imgsz, self.imgsz)) + + img = img[:, :, ::-1].transpose(2, 0, 1) + img = np.ascontiguousarray(img) + + im_tensor = Tensor(img, ms.float32) / 255.0 + im_tensor = im_tensor.expand_dims(0) + + return im_tensor, orig_img, (ratio, pad_w, pad_h) + + def inference(self, im_tensor): + """执行模型前向传播""" + return self.model(im_tensor) + + def postprocess(self, preds, orig_img, preprocess_info): + """后处理抽象方法，需由各任务子类具体实现""" + raise NotImplementedError("BasePredictor 不执行特定的解析逻辑，请在子类中重写 postprocess 方法。") + + def __call__(self, source, model=None, ckpt_path=None): + """推理流水线主调度入口""" + if model is not None: + self.setup_model(model, ckpt_path) + + if self.model is None: + raise RuntimeError("[ERROR] 模型未初始化，无法启动推理流水线。") + + self.setup_source(source) + self.results = [] + + LOGGER.info(f"推理引擎启动，共探测到 {len(self.dataset)} 份输入样本。") + + for img_path_or_arr in self.dataset: + # 1. 预处理 + t1 = time.time() + im_tensor, orig_img, prep_info = self.preprocess(img_path_or_arr) + + # 2. 推理 + t2 = time.time() + preds = self.inference(im_tensor) + + # 3. 后处理 + t3 = time.time() + result = self.postprocess(preds, orig_img, prep_info) + self.results.append(result) + + # 4. 耗时统计 + inf_time = (t3 - t2) * 1000 + post_time = (time.time() - t3) * 1000 + name = img_path_or_arr if isinstance(img_path_or_arr, str) else "numpy_array" + LOGGER.info(f"处理完成 [{os.path.basename(name)}] | 前向推理: {inf_time:.1f}ms | 后处理: {post_time:.1f}ms") + + return self.results \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/engine/trainer.py b/src/mindnlp/models/ultralytics/engine/trainer.py new file mode 100644 index 000000000..64beb6b3c --- /dev/null +++ b/src/mindnlp/models/ultralytics/engine/trainer.py @@ -0,0 +1,171 @@ +import os +import yaml +from pathlib import Path +import mindspore as ms +import numpy as np +from mindspore import nn, ops +from utils.ema import ModelEMA + +class TrainStepWithClip(nn.TrainOneStepCell): + """支持梯度裁剪的单步训练封装类""" + def __init__(self, network, optimizer, clip_val=10.0): + super(TrainStepWithClip, self).__init__(network, optimizer) + self.clip_val = clip_val + self.grad_fn = ops.value_and_grad(self.network, grad_position=None, weights=self.weights) + + def construct(self, *inputs): + loss, grads = self.grad_fn(*inputs) + # 限制全局梯度范数，防止梯度爆炸 + grads = ops.clip_by_global_norm(grads, self.clip_val) + loss = ops.depend(loss, self.optimizer(grads)) + return loss + +class YOLOWithLossCell(nn.Cell): + """适配检测任务的多输入 Loss 封装，负责连接模型输出与损失函数""" + def __init__(self, backbone, loss_fn): + super(YOLOWithLossCell, self).__init__(auto_prefix=False) + self._backbone = backbone + self._loss_fn = loss_fn + + def construct(self, x, labels): + # 前向传播获取预测结果 + pred = self._backbone(x) + # 将预测值与完整的标签字典传给损失函数，防止维度丢失 + return self._loss_fn(pred, labels) + +class BaseTrainer: + """训练基类：负责调度全局流程、解析配置、管理模型与日志保存""" + + def __init__(self, args): + self.args = args + self.device = ms.get_context("device_target") + self.save_dir = Path(args.save_dir) + self.save_dir.mkdir(parents=True, exist_ok=True) + + self.best_fitness = 0.0 + self.start_epoch = 0 + self.epochs = args.epochs + + # 统一加载超参数配置 (hyp.yaml) + if hasattr(args, 'hyp') and os.path.exists(args.hyp): + with open(args.hyp, "r", encoding="utf-8") as f: + self.hyp = yaml.safe_load(f) + else: + raise FileNotFoundError(f"[ERROR] 未找到超参数配置文件: {getattr(args, 'hyp', '未指定')}") + + # 调度子类构建核心组件 + self.train_loader = self.get_dataloader(is_training=True) + self.val_loader = self.get_dataloader(is_training=False) + self.model = self.get_model(cfg=self.args.model_cfg, weights=self.args.weights) + self.optimizer = self.build_optimizer(self.model) + self.ema = ModelEMA(self.model) + + # 构建训练计算图 + self.loss_fn = self.get_loss_fn() + self.net_with_loss = YOLOWithLossCell(self.model, self.loss_fn) + + self.train_step = TrainStepWithClip(self.net_with_loss, self.optimizer, clip_val=10.0) + + def train(self): + """核心训练迭代逻辑""" + print(f"[INFO] 训练任务启动，总轮数: {self.epochs} epochs") + steps_per_epoch = self.train_loader.get_dataset_size() + + for epoch in range(self.start_epoch, self.epochs): + self.model.set_train(True) + self.train_step.set_train(True) + + for step, batch in enumerate(self.train_loader.create_dict_iterator()): + # 预处理批次数据 + batch = self.preprocess_batch(batch) + + # 执行单步训练，batch 字典包含所需的所有标签信息 + imgs = batch["image"] + loss = self.train_step(imgs, batch) + self.ema.update(self.model) + + # 日志打印逻辑 + if step % 10 == 0: + if isinstance(loss, (tuple, list)): + loss_items = [float(x.asnumpy()) for x in loss] + num_loss = len(loss_items) + + if num_loss == 4: # 分割任务 + l_box, l_cls, l_dfl, l_mask = loss_items + total_loss = sum(loss_items) + loss_str = f"Box: {l_box:.4f} | Cls: {l_cls:.4f} | DFL: {l_dfl:.4f} | Mask: {l_mask:.4f}" + elif num_loss == 3: # 检测任务 + l_box, l_cls, l_dfl = loss_items + total_loss = sum(loss_items) + loss_str = f"Box: {l_box:.4f} | Cls: {l_cls:.4f} | DFL: {l_dfl:.4f}" + else: + total_loss = sum(loss_items) + loss_str = f"Loss Items: {[f'{x:.4f}' for x in loss_items]}" + + print(f"Epoch [{epoch}/{self.epochs-1}] Step [{step}/{steps_per_epoch}] | Total Loss: {total_loss:.4f} | {loss_str}") + else: + loss_val = float(loss.asnumpy()) if hasattr(loss, "asnumpy") else float(loss) + print(f"Epoch [{epoch}/{self.epochs-1}] Step [{step}/{steps_per_epoch}] | Loss: {loss_val:.4f}") + + # 验证与保存阶段 + if (epoch + 1) % self.args.val_interval == 0 or epoch == self.epochs - 1: + print(f"\n[INFO] 开始执行 Epoch {epoch} 验证程序...") + self.model.set_train(False) + + validator = self.get_validator() + stats = validator(self.ema.ema_model) + + print("-" * 50) + print(f"[评估报告] Epoch {epoch}") + for k, v in stats.items(): + if k.startswith('metrics/'): + metric_name = k.replace('metrics/', '') + print(f" - {metric_name:<15} : {float(v):.5f}") + + fitness_f = float(stats.get('fitness', 0.0)) + print(f"[INFO] 当前模型综合评价指标 (Fitness): {fitness_f:.5f}") + print("-" * 50 + "\n") + + self._save_checkpoint(epoch, fitness_f) + + self.model.set_train(True) + self.train_step.set_train(True) + + def _save_checkpoint(self, epoch, fitness): + """权重序列化保存逻辑""" + ms.save_checkpoint(self.ema.ema_model, str(self.save_dir / "last.ckpt")) + if fitness > self.best_fitness: + self.best_fitness = fitness + ms.save_checkpoint(self.ema.ema_model, str(self.save_dir / "best.ckpt")) + print(f"[INFO] 已更新最佳模型权重 (best.ckpt)，当前最高精度: {self.best_fitness:.4f}") + + # 抽象接口声明，子类需实现具体的业务逻辑 + def get_dataloader(self, is_training): raise NotImplementedError + def get_model(self, cfg, weights): raise NotImplementedError + def build_optimizer(self, model): raise NotImplementedError + def get_validator(self): raise NotImplementedError + def get_loss_fn(self): raise NotImplementedError + + def preprocess_batch(self, batch): + """ + 数据预处理：执行图像归一化与标签 Tensor 转换 + """ + # 图像预处理：强制转为 Tensor 并执行像素归一化 + img = batch["image"] + if not isinstance(img, ms.Tensor): + img = ms.Tensor(img, ms.float32) + + # 仅在数据为原始像素值 (0-255) 时进行归一化 + if img.max() > 1.0: + img = ops.cast(img, ms.float32) / 255.0 + batch["image"] = img + + # 标签处理：确保 bboxes 和 batch_idx 类型符合 Loss 计算要求 + if "bboxes" in batch and "batch_idx" in batch: + if not isinstance(batch["bboxes"], ms.Tensor): + batch["bboxes"] = ms.Tensor(batch["bboxes"], ms.float32) + + if not isinstance(batch["batch_idx"], ms.Tensor): + batch["batch_idx"] = ms.Tensor(batch["batch_idx"], ms.int32) + + return batch \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/engine/validator.py b/src/mindnlp/models/ultralytics/engine/validator.py new file mode 100644 index 000000000..5872d767a --- /dev/null +++ b/src/mindnlp/models/ultralytics/engine/validator.py @@ -0,0 +1,123 @@ +import time +import logging +import yaml +from pathlib import Path + +import mindspore as ms +from mindspore import ops +from tqdm import tqdm + +from utils.ops import non_max_suppression + +# 统一日志配置 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +LOGGER = logging.getLogger(__name__) + +class BaseValidator: + """ + 验证任务基类 + 负责管理验证数据集的加载、前向推理耗时统计，以及指标计算的评估等通用逻辑。 + """ + def __init__(self, dataloader=None, save_dir=None, args=None): + self.args = args + self.dataloader = dataloader + self.save_dir = Path(save_dir) if save_dir else Path("./runs/val") + self.save_dir.mkdir(parents=True, exist_ok=True) + + self.device = ms.get_context("device_target") + self.speed = {"preprocess": 0.0, "inference": 0.0, "postprocess": 0.0} + self.seen = 0 + self.training = False + + # 解析超参数配置 (支持从 hyp.yaml 读取 conf, iou, half 等验证参数) + self.hyp = {} + if hasattr(args, 'hyp') and args.hyp and Path(args.hyp).exists(): + with open(args.hyp, "r", encoding="utf-8") as f: + self.hyp = yaml.safe_load(f) + + def get_dataloader(self, dataset_path, batch_size): + """抽象接口：获取验证数据集的 DataLoader""" + raise NotImplementedError("子类必须实现 get_dataloader 方法") + + def __call__(self, model): + """验证流程主入口""" + self.training = model.training if hasattr(model, 'training') else False + + # 确保模型处于推理模式 + model.set_train(False) + self.init_metrics(model) + + bar = tqdm(self.dataloader.create_dict_iterator(), + desc="Validating", + total=self.dataloader.get_dataset_size()) + + for batch_i, batch in enumerate(bar): + # 1. 数据预处理 + t0 = time.time() + batch = self.preprocess(batch) + self.speed["preprocess"] += time.time() - t0 + + # 2. 模型前向推理 + t1 = time.time() + preds = model(batch["image"]) + self.speed["inference"] += time.time() - t1 + + # 3. 后处理与指标收集 + t2 = time.time() + preds = self.postprocess(preds) + self.update_metrics(preds, batch) + self.speed["postprocess"] += time.time() - t2 + + self.seen += batch["image"].shape[0] + + # 4. 指标汇总与结果打印 + self.finalize_metrics() + stats = self.get_stats() + self.print_results() + + return stats + + def preprocess(self, batch): + """ + 通用图像预处理逻辑：类型转换与归一化 + """ + img = ops.cast(batch["image"], ms.float32) + + # 像素值归一化至 [0.0, 1.0] 域 + if img.max() > 2.0: + img = img / 255.0 + + batch["image"] = img + return batch + + def postprocess(self, preds): + """ + 通用后处理逻辑（默认为目标检测任务的非极大值抑制 NMS） + 注意：分类等特定任务需在子类中重写此方法 + """ + p = preds[0] if isinstance(preds, tuple) else preds + + # 优先从超参数配置读取阈值，若无则回退至命令行参数或默认值 + conf_thres = self.hyp.get('conf', getattr(self.args, 'conf', 0.001)) + iou_thres = self.hyp.get('iou', getattr(self.args, 'iou', 0.6)) + + out = non_max_suppression( + p, + conf_thres=conf_thres, + iou_thres=iou_thres, + nc=getattr(self, 'nc', 80) + ) + return out + + # --- 抽象接口定义 --- + def init_metrics(self, model): pass + def update_metrics(self, preds, batch): pass + def finalize_metrics(self): pass + + def get_stats(self): + """计算平均单张图像的各阶段耗时 (ms)""" + return {"speed": {k: v / self.seen * 1000 for k, v in self.speed.items() if self.seen > 0}} + + def print_results(self): + speed_str = " | ".join([f"{k}: {v:.1f}ms" for k, v in self.get_stats().get('speed', {}).items()]) + LOGGER.info(f"推理测速: {speed_str}") \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py b/src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py new file mode 100644 index 000000000..84da1c62c --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py @@ -0,0 +1,84 @@ +import argparse +import os +import logging +import mindspore as ms + +from models.yolo.classify.predict import ClassificationPredictor +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForClassification + +# 统一日志配置 +logging.basicConfig(level=logging.INFO, format='%(message)s') +LOGGER = logging.getLogger(__name__) + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Classification Inference Pipeline") + + # --- 核心数据与权重参数 --- + parser.add_argument('--source', type=str, default='./datasets/imagenette2-160/val/n01440764', help='待推理的图像文件或目录路径') + parser.add_argument('--weights', type=str, default='./yolo11n-cls.ckpt', help='预训练权重文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-cls.yaml', help='模型拓扑配置文件路径') + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help='超参数配置文件路径 (包含推理阈值配置)') + + # --- 架构超参数 --- + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型参数规模') + parser.add_argument('--nc', type=int, default=10, help='分类任务类别总数') + parser.add_argument('--imgsz', type=int, default=224, help='网络输入分辨率') + + # --- 运行环境设定 --- + parser.add_argument('--device', type=str, default='Ascend', help='计算硬件平台 (Ascend/GPU/CPU)') + parser.add_argument('--save_dir', type=str, default='./runs/cls/predict', help='可视化结果持久化保存目录') + + args = parser.parse_args() + + # 配置 MindSpore 计算图执行模式 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + LOGGER.info(f"[INFO] 正在构建网络拓扑 (规模: {args.scale}, 类别映射空间: {args.nc})") + cfg = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task='classify') + cfg.nc = args.nc + model = YOLO11ForClassification(cfg) + + # 架构载入模块：解决迁移学习阶段的维度冲突 + param_dict = ms.load_checkpoint(args.weights) + model_dict = model.parameters_and_names() + new_param_dict = {} + + for name, param in model_dict: + if name in param_dict: + ckpt_param = param_dict[name] + + # 过滤形状不匹配的参数（如分类头），防止加载权重时报错 + if param.shape == ckpt_param.shape: + new_param_dict[name] = ckpt_param + else: + LOGGER.warning(f"[WARNING] 丢弃层特征权重 [{name}] | 原因: 输出维度不匹配 " + f"(当前网络需求: {param.shape} | 检查点参数: {ckpt_param.shape})") + else: + LOGGER.warning(f"[WARNING] 检查点中未匹配到网络特征层 [{name}]，已默认回退为随机初始化。") + + ms.load_param_into_net(model, new_param_dict, strict_load=False) + LOGGER.info("[INFO] 网络主干权重数据对齐并装载完毕。") + + # 实例化推理引擎 + predictor = ClassificationPredictor(cfg=args) + + # 执行前向推理逻辑 + results = predictor(source=args.source, model=model, ckpt_path=None) + + # 遍历评估结果与图表保存 + LOGGER.info("\n" + "="*50) + for i, res in enumerate(results): + source_item = predictor.dataset[i] + file_name = os.path.basename(source_item) if isinstance(source_item, str) else f"pred_stream_{i}.jpg" + + # 执行结果可视化并保存至本地文件 + res.save(save_dir=args.save_dir, file_name=file_name) + LOGGER.info(f"成功处理图片 [{file_name}] | 预测分类: {res.top1_name:<15} | 相对置信度: {res.top1[1]:.4f}") + + LOGGER.info("="*50) + absolute_save_dir = os.path.abspath(args.save_dir) + LOGGER.info(f"[INFO] 批量推理任务执行完毕，结果文件已保存至: {absolute_save_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py b/src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py new file mode 100644 index 000000000..7614c7c5a --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py @@ -0,0 +1,37 @@ +import argparse +import mindspore as ms +import sys +import os + +# 从核心库中导入纯净的 Trainer +from models.yolo.classify.train import ClassificationTrainer + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Classification Training") + + parser.add_argument('--data', type=str, default='./cfg/datasets/imagenette2-160.yaml') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-cls.yaml') + parser.add_argument('--weights', type=str, default='./yolo11n-cls.ckpt', help="初始化权重路径，留空则从头训练") + parser.add_argument('--scale', type=str, default='n') + parser.add_argument('--imgsz', type=int, default=224) + parser.add_argument('--batch', type=int, default=64) + parser.add_argument('--epochs', type=int, default=100) + parser.add_argument('--workers', type=int, default=8, help="数据加载的线程数") + parser.add_argument('--device', type=str, default='Ascend') + parser.add_argument('--val_interval', type=int, default=1, help="每隔几个 epoch 验证一次") + parser.add_argument('--save_dir', type=str, default="./runs/cls/train") + + + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help="算法调优超参数配置文件路径") + + args = parser.parse_args() + + # 设置运行模式与硬件 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 面向对象启动 + trainer = ClassificationTrainer(args) + trainer.train() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py b/src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py new file mode 100644 index 000000000..8222d9b76 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py @@ -0,0 +1,77 @@ +import argparse +import mindspore as ms + +from models.yolo.classify.val import ClassificationValidator +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForClassification + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Classification Validation Pipeline") + + # 核心路径配置 + parser.add_argument('--data', type=str, default='./datasets/imagenette2-160/val', help="验证集路径") + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-cls.yaml', help="模型结构配置文件") + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help="超参数配置文件路径") + parser.add_argument('--weights', type=str, default='./yolo11n-cls.ckpt', help="验证所需加载的模型权重路径 (例如: best.ckpt)") + #parser.add_argument('--weights', type=str, default='./runs/cls/train_finetune/best.ckpt', help="验证所需加载的模型权重路径 (例如: best.ckpt)") + + # 运行参数 + parser.add_argument('--imgsz', type=int, default=224, help="输入图像尺寸") + parser.add_argument('--batch', type=int, default=128, help="验证批次大小") + parser.add_argument('--workers', type=int, default=8, help="数据加载线程数") + parser.add_argument('--device', type=str, default='Ascend', help="计算硬件类型") + parser.add_argument('--save_dir', type=str, default="./runs/cls/val", help="验证结果保存目录") + + # 用于兼容基类的占位参数 + parser.add_argument('--model', type=str, default='', help="向下兼容参数") + + args = parser.parse_args() + + # 权重路径兼容处理 + if args.weights and not args.model: + args.model = args.weights + + # 设置 MindSpore 运行环境 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + print(f"[INFO] 正在解析模型拓扑配置: {args.model_cfg}") + cfg = YOLOConfig(yaml_path=args.model_cfg, task='classify') + + # 强制覆盖分类数量，实际项目中应从 args.data 对应的 dataset.yaml 中读取 + cfg.nc = getattr(args, 'nc', 10) + model = YOLO11ForClassification(cfg) + + # ---------------- 权重加载与形状对齐 ---------------- + print(f"[INFO] 正在加载检查点权重: {args.weights}") + param_dict = ms.load_checkpoint(args.weights) + model_dict = model.parameters_and_names() + + new_param_dict = {} + for name, param in model_dict: + if name in param_dict: + ckpt_param = param_dict[name] + + # 权重形状匹配检查 + # 说明：当跨数据集评估时（如 ImageNet 的 1000 类权重，在 10 类的 Imagenette 上验证）， + # 最后的线性分类层尺寸将不匹配，此时忽略分类头，仅加载 Backbone 权重。 + if param.shape == ckpt_param.shape: + new_param_dict[name] = ckpt_param + else: + print(f"[WARNING] 尺寸不匹配，跳过权重加载: {name} (当前模型 {param.shape} vs 检查点 {ckpt_param.shape})") + else: + print(f"[WARNING] 检查点中缺失权重项: {name}，当前层将保留随机初始化") + + ms.load_param_into_net(model, new_param_dict, strict_load=False) + print(f"[INFO] 权重对齐并加载完成。") + + # ---------------- 验证流程启动 ---------------- + validator = ClassificationValidator(args=args) + validator.dataloader = validator.get_dataloader(args.data, batch_size=args.batch) + + print(f"[INFO] DataLoader 构建完毕，即将进入评估环节...") + + # 启动 Validator 的 __call__ 流程 + validator(model=model) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py b/src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py new file mode 100644 index 000000000..c3102b93b --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py @@ -0,0 +1,57 @@ +import argparse +import mindspore as ms +import sys +import os + +from models.yolo.detect.predict import DetectionPredictor +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForObjectDetection + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Object Detection Inference Pipeline") + + # --- 核心推理参数 --- + parser.add_argument('--source', type=str, default='./datasets/coco128/coco128/images/train2017', help='待预测的图像或目录路径') + parser.add_argument('--weights', type=str, default='./yolo11n.ckpt', help='预训练权重文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + parser.add_argument('--nc', type=int, default=80, help='分类类别总数') + + # --- 推理后处理参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='网络输入分辨率大小') + parser.add_argument('--conf', type=float, default=0.25, help='预测边界框的置信度阈值') + parser.add_argument('--iou', type=float, default=0.45, help='NMS 交并比阈值') + + # --- 运行环境参数 --- + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='计算硬件目标') + parser.add_argument('--save_dir', type=str, default='./runs/detect/predict', help='预测渲染结果保存目录') + + args = parser.parse_args() + + # 配置 MindSpore 运行环境 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + print(f"[Info] 初始化检测模型架构 (Scale: {args.scale}, Classes: {args.nc})") + cfg = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="detect") + cfg.nc = args.nc + model = YOLO11ForObjectDetection(cfg) + + # 实例化预测器 + predictor = DetectionPredictor(cfg=args) + + # 启动推理流水线 + results = predictor(source=args.source, model=model, ckpt_path=args.weights) + + print("\n" + "-" * 50) + for i, res in enumerate(results): + source_item = predictor.dataset[i] + file_name = os.path.basename(source_item) if isinstance(source_item, str) else f"pred_{i}.jpg" + + res.save(save_dir=args.save_dir, file_name=file_name) + print(f"[Result] 图像 {file_name}: 共检测到 {len(res.det)} 个目标") + + print("-" * 50) + print(f"[Info] 批量推理流水线执行完毕，结果存放于: {args.save_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py b/src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py new file mode 100644 index 000000000..e5ed4f6a8 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py @@ -0,0 +1,47 @@ +import argparse +import os +import yaml +import mindspore as ms + +from models.yolo.detect.train import DetectionTrainer + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Object Detection Training Pipeline") + + # ---------------- 核心执行参数 ---------------- + parser.add_argument('--data', type=str, default='./cfg/datasets/coco128.yaml', help='数据集配置文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11.yaml', help='模型拓扑架构配置文件') + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help='超参数配置文件 (包含学习率、增强策略等)') + parser.add_argument('--weights', type=str, default='', help='预训练 Checkpoint 路径 (留空则随机初始化)') + + # ---------------- 训练规模参数 ---------------- + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模变体') + parser.add_argument('--epochs', type=int, default=100, help='训练总轮数') + parser.add_argument('--batch', type=int, default=16, help='单卡批次大小') + parser.add_argument('--imgsz', type=int, default=640, help='网络输入分辨率') + parser.add_argument('--workers', type=int, default=8, help='数据加载线程数') + + # ---------------- 环境与保存机制 ---------------- + parser.add_argument('--val_interval', type=int, default=10, help='验证评估频率 (单位: Epoch)') + parser.add_argument('--save_dir', type=str, default='./runs/detect/train', help='模型与日志保存目录') + parser.add_argument('--device', type=str, default='Ascend', help='计算硬件 (Ascend/GPU/CPU)') + + args = parser.parse_args() + + # 配置 MindSpore 运行模式与硬件靶标 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 动态解析数据集类别总数 + if not os.path.exists(args.data): + raise FileNotFoundError(f"[ERROR] 无法找到数据集配置文件: {args.data}") + + with open(args.data, 'r', encoding='utf-8') as f: + data_cfg = yaml.safe_load(f) + args.nc = data_cfg.get('nc', 80) + + # 实例化并触发训练调度器 + trainer = DetectionTrainer(args) + trainer.train() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py b/src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py new file mode 100644 index 000000000..9a806cc24 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py @@ -0,0 +1,78 @@ +import argparse +import mindspore as ms +import os +import yaml + +from models.yolo.detect.val import DetectionValidator +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForObjectDetection +from data.loaders import create_dataloader + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Object Detection Standalone Validation") + + # --- 核心验证参数 --- + parser.add_argument('--data', type=str, default='./cfg/datasets/coco128.yaml', help='数据集配置文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--weights', type=str, default='./yolo11n.ckpt', help='待验证的预训练权重文件路径') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + + # --- 运行环境参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='验证测试所用的图像分辨率') + parser.add_argument('--batch', type=int, default=16, help='数据加载过程的批次大小') + parser.add_argument('--workers', type=int, default=8, help='数据加载的并行子线程数') + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='目标计算硬件平台') + parser.add_argument('--save_dir', type=str, default='./runs/detect/val', help='验证日志及性能指标的保存目录') + + args = parser.parse_args() + + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 解析数据集配置 + print(f"[Info] 载入数据集配置信息: {args.data}") + with open(args.data, 'r', encoding='utf-8') as f: + data_cfg = yaml.safe_load(f) + + args.nc = data_cfg.get('nc', 80) + data_path = data_cfg.get('path', '') + val_dir = os.path.join(data_path, data_cfg.get('val', 'val')) + + print("[Info] 构建目标检测验证数据集 DataLoader...") + val_loader = create_dataloader( + path=val_dir, + imgsz=args.imgsz, + batch_size=args.batch, + task='detect', + is_training=False, + num_workers=args.workers + ) + + # 实例化模型架构并加载权重 + print(f"[Info] 部署模型计算图架构 (Scale: {args.scale}, Classes: {args.nc})") + config = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="detect") + config.nc = args.nc + model = YOLO11ForObjectDetection(config) + + if os.path.exists(args.weights): + print(f"[Info] 加载模型参数权重: {args.weights}") + param_dict = ms.load_checkpoint(args.weights) + ms.load_param_into_net(model, param_dict, strict_load=False) + else: + raise FileNotFoundError(f"[Error] 指定的权重文件不存在: {args.weights}，请验证输入路径。") + + # 实例化 Validator 执行验证评估 + print("[Info] 触发验证评估流水线执行过程...") + validator = DetectionValidator(dataloader=val_loader, save_dir=args.save_dir, args=args) + stats = validator(model=model) + + mAP50 = stats.get('metrics/mAP50(B)', 0.0) + mAP50_95 = stats.get('metrics/mAP50-95(B)', 0.0) + + print("\n" + "-" * 50) + print("[Result] 目标检测验证任务结束。全局性能指标如下：") + print(f" - mAP@50: {mAP50:.4f}") + print(f" - mAP@50-95: {mAP50_95:.4f}") + print("-" * 50) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py b/src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py new file mode 100644 index 000000000..e331501d9 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py @@ -0,0 +1,71 @@ +import argparse +import mindspore as ms +import sys +import os + +from models.yolo.pose.predict import PosePredictor +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForPose + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Inference Pipeline") + + # --- 核心推理输入与架构参数 --- + parser.add_argument('--source', type=str, default='./datasets/coco8-pose/images/val', help='待推理图像的路径或文件夹目录') + parser.add_argument('--weights', type=str, default='./yolo11n-pose.ckpt', help='预训练权重文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-pose.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + parser.add_argument('--nc', type=int, default=1, help='预测的类别总数 (姿态估计通常默认为 1 类: person)') + + # --- 推理后处理与阈值参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='网络推理输入分辨率') + parser.add_argument('--conf', type=float, default=0.25, help='预测框置信度阈值') + parser.add_argument('--iou', type=float, default=0.45, help='NMS 交并比 (IoU) 阈值') + + # --- 运行环境配置 --- + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='计算硬件目标') + parser.add_argument('--save_dir', type=str, default='./runs/pose/predict', help='渲染结果的保存目录') + + args = parser.parse_args() + + # 配置计算平台与运行模式 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 实例化模型架构配置 + print(f"[Info] 初始化模型架构 (Scale: {args.scale}, Classes: {args.nc})") + cfg = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="pose") + cfg.nc = args.nc + + # 确保 config 包含关键点形状定义以防越界异常 + if not hasattr(cfg, 'kpt_shape'): + cfg.kpt_shape = [17, 3] + + model = YOLO11ForPose(cfg) + + # 实例化姿态预测器 + predictor = PosePredictor(cfg=args) + + # 执行批量推理流水线 + results = predictor(source=args.source, model=model, ckpt_path=args.weights) + + # 解析并存储预测结果 + print("\n" + "-"*50) + for i, res in enumerate(results): + source_item = predictor.dataset[i] + file_name = os.path.basename(source_item) if isinstance(source_item, str) else f"pose_{i}.jpg" + + # 调用封装的 render 方法保存图像 + res.save(save_dir=args.save_dir, file_name=file_name) + + # 统计检出的人体实例数量 + num_instances = len(res.det) if hasattr(res, 'det') and res.det is not None else 0 + if num_instances == 0 and hasattr(res, 'boxes') and res.boxes is not None: + num_instances = len(res.boxes) + + print(f"[Result] 文件 {file_name}: 检出 {num_instances} 个人体姿态实例") + + print("-"*50) + print(f"[Info] 推理任务结束。所有结果已存入: {args.save_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py b/src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py new file mode 100644 index 000000000..81961cef4 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py @@ -0,0 +1,36 @@ +import argparse +import os +import mindspore as ms + +from models.yolo.pose.train import PoseTrainer + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Training Pipeline") + + # 1. 基础与架构参数 + parser.add_argument('--data', type=str, default='./cfg/datasets/coco8-pose.yaml', help='数据集配置文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-pose.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help='训练超参数配置文件 (学习率、衰减等)') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='网络规模标识') + parser.add_argument('--weights', type=str, default='./yolo11n-pose.ckpt', help='预训练权重文件路径 (.ckpt)') + + # 2. 训练周期与计算超参数 + parser.add_argument('--epochs', type=int, default=100, help='总训练迭代轮数') + parser.add_argument('--batch', type=int, default=4, help='全局批处理大小') + parser.add_argument('--imgsz', type=int, default=640, help='网络输入图像分辨率') + parser.add_argument('--workers', type=int, default=8, help='数据加载并行线程数') + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='目标计算硬件平台') + parser.add_argument('--val_interval', type=int, default=10, help='验证评估频率 (单位: epoch)') + parser.add_argument('--save_dir', type=str, default='./runs/pose/train', help='权重及日志保存输出目录') + + args = parser.parse_args() + + # 配置 MindSpore 运行环境 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 面向对象启动训练流水线 + trainer = PoseTrainer(args) + trainer.train() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py b/src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py new file mode 100644 index 000000000..845b42806 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py @@ -0,0 +1,77 @@ +import argparse +import mindspore as ms +import os +import sys +import yaml + +from models.yolo.pose.val import PoseValidator +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForPose +from data.loaders import create_dataloader + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Standalone Validation") + + # --- 基础配置参数 --- + parser.add_argument('--data', type=str, default='./cfg/datasets/coco8-pose.yaml', help='数据集配置文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-pose.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--weights', type=str, default='./yolo11n-pose.ckpt', help='待评估的预训练权重文件路径') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + + # --- 评估环境参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='网络输入图像分辨率') + parser.add_argument('--batch', type=int, default=4, help='验证过程的批次大小') + parser.add_argument('--workers', type=int, default=8, help='数据加载并行线程数') + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='计算硬件目标') + parser.add_argument('--save_dir', type=str, default='./runs/pose/val', help='验证结果与指标保存目录') + + args = parser.parse_args() + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 解析数据集配置 + with open(args.data, 'r', encoding='utf-8') as f: + data_cfg = yaml.safe_load(f) + + args.nc = data_cfg.get('nc', 1) + kpt_shape = data_cfg.get('kpt_shape', [17, 3]) + names = data_cfg.get('names', {i: f'class_{i}' for i in range(args.nc)}) + val_dir = os.path.join(data_cfg.get('path', ''), data_cfg.get('val', 'val')) + + # 实例化数据加载器 + val_loader = create_dataloader( + path=val_dir, + imgsz=args.imgsz, + batch_size=args.batch, + task='pose', + is_training=False, + num_workers=args.workers + ) + + # 实例化模型架构并加载权重 + config = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="pose") + config.nc = args.nc + config.kpt_shape = kpt_shape + model = YOLO11ForPose(config) + + print(f"[Info] 加载验证权重: {args.weights}") + ms.load_param_into_net(model, ms.load_checkpoint(args.weights), strict_load=False) + + # 执行验证流水线 + validator = PoseValidator( + dataloader=val_loader, + save_dir=args.save_dir, + args=args, + names=names, + kpt_shape=kpt_shape + ) + stats = validator(model=model) + + # 格式化输出核心评估指标 + print("\n" + "-"*50) + print("[Result] 姿态估计验证完成。核心评估指标如下：") + print(f" - [Box] 边界框 mAP@50-95: {stats.get('metrics/mAP50-95(B)', 0.0):.4f}") + print(f" - [Pose] 关键点 mAP@50-95: {stats.get('metrics/mAP50-95(P)', 0.0):.4f}") + print("-"*50 + "\n") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py b/src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py new file mode 100644 index 000000000..d42415f24 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py @@ -0,0 +1,59 @@ +import argparse +import mindspore as ms +import sys +import os + +from models.yolo.segment.predict import SegmentationPredictor +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForSegmentation + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Segmentation Inference Pipeline") + + # --- 核心推理输入与架构参数 --- + parser.add_argument('--source', type=str, default='./datasets/coco128-seg/images/train2017', help='待推理图像的路径或文件夹目录') + parser.add_argument('--weights', type=str, default='./yolo11n-seg.ckpt', help='预训练权重文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-seg.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + parser.add_argument('--nc', type=int, default=80, help='预测的类别总数') + + # --- 推理后处理与阈值参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='网络推理输入分辨率') + parser.add_argument('--conf', type=float, default=0.25, help='预测框置信度阈值') + parser.add_argument('--iou', type=float, default=0.45, help='NMS 交并比 (IoU) 阈值') + + # --- 运行环境配置 --- + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='计算硬件目标') + parser.add_argument('--save_dir', type=str, default='./runs/segment/predict', help='渲染结果的保存目录') + + args = parser.parse_args() + + # 配置计算平台与运行模式 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 实例化模型架构配置 + print(f"[Info] 初始化分割模型架构 (Scale: {args.scale}, Classes: {args.nc})") + cfg = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="segment") + cfg.nc = args.nc + model = YOLO11ForSegmentation(cfg) + + # 实例化分割预测器 + predictor = SegmentationPredictor(cfg=args) + + # 执行批量推理流水线 + results = predictor(source=args.source, model=model, ckpt_path=args.weights) + + # 解析并存储预测结果 + print("\n" + "-"*50) + for i, res in enumerate(results): + source_item = predictor.dataset[i] + file_name = os.path.basename(source_item) if isinstance(source_item, str) else f"pred_{i}.jpg" + + res.save(save_dir=args.save_dir, file_name=file_name) + print(f"[Result] 文件 {file_name}: 检出并分割出 {len(res.det)} 个目标实例") + + print("-" * 50) + print(f"[Info] 批量分割推理流水线执行完毕，结果存放于: {args.save_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py b/src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py new file mode 100644 index 000000000..b306999f4 --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py @@ -0,0 +1,43 @@ +import argparse +import os +import logging +import mindspore as ms + +from models.yolo.segment.train import SegmentationTrainer + +logging.basicConfig(level=logging.INFO, format='%(message)s') +LOGGER = logging.getLogger(__name__) + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Instance Segmentation Training Pipeline") + + # 核心路径与配置 + parser.add_argument('--data', type=str, default='./cfg/datasets/coco128-seg.yaml', help='数据集结构定义 YAML') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-seg.yaml', help='模型拓扑配置文件') + parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help='控制训练周期的超参数文件 (学习率、增强等)') + parser.add_argument('--weights', type=str, default=None, help='预训练检查点路径 (.ckpt)') + + # 模型架构与执行流参数 + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='基础模型复杂度层级') + parser.add_argument('--epochs', type=int, default=100, help='训练总回合数') + parser.add_argument('--batch_size', type=int, default=16, help='单路处理批次') + parser.add_argument('--imgsz', type=int, default=640, help='进入网络的图像分辨率空间') + parser.add_argument('--workers', type=int, default=8, help='多进程数据装载核心数') + + # 硬件与保存配置 + parser.add_argument('--device', type=str, default='Ascend', help='深度学习计算硬件靶标 (Ascend/GPU/CPU)') + parser.add_argument('--val_interval', type=int, default=10, help='两次评估循环之间的 Epoch 间距') + parser.add_argument('--save_dir', type=str, default='./runs/segment/train', help='训练档案与模型文件的输出') + + args = parser.parse_args() + + # 指配计算域 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 启动封装完善的对象生命周期 + LOGGER.info("[INFO] SegmentationTrainer 初始化就绪，即将进入训练大循环。") + trainer = SegmentationTrainer(args) + trainer.train() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py b/src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py new file mode 100644 index 000000000..40bb368fc --- /dev/null +++ b/src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py @@ -0,0 +1,88 @@ +import argparse +import mindspore as ms +import os +import yaml + +from models.yolo.segment.val import SegmentationValidator +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForSegmentation +from data.loaders import create_dataloader + +def main(): + parser = argparse.ArgumentParser(description="YOLO11 Segmentation Standalone Validation") + + # --- 核心验证参数 --- + parser.add_argument('--data', type=str, default='./cfg/datasets/coco128-seg.yaml', help='数据集配置文件路径') + parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-seg.yaml', help='模型架构 YAML 配置文件') + parser.add_argument('--weights', type=str, default='./yolo11n-seg.ckpt', help='待评估的预训练权重文件路径') + parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') + + # --- 运行环境参数 --- + parser.add_argument('--imgsz', type=int, default=640, help='网络推理输入分辨率') + parser.add_argument('--batch', type=int, default=16, help='验证过程的批次大小') + parser.add_argument('--workers', type=int, default=8, help='数据加载并行线程数') + parser.add_argument('--device', type=str, default='Ascend', choices=['Ascend', 'GPU', 'CPU'], help='计算硬件目标') + parser.add_argument('--save_dir', type=str, default='./runs/segment/val', help='验证结果与指标保存目录') + + args = parser.parse_args() + + # 设置 MindSpore 运行环境 + ms.set_context(mode=ms.PYNATIVE_MODE, device_target=args.device) + + # 解析数据集配置，准备 DataLoader + print(f"[Info] 解析数据集配置: {args.data}") + with open(args.data, 'r', encoding='utf-8') as f: + data_cfg = yaml.safe_load(f) + + args.nc = data_cfg.get('nc', 80) + names = data_cfg.get('names', {i: f'class_{i}' for i in range(args.nc)}) + + data_path = data_cfg.get('path', '') + val_dir = os.path.join(data_path, data_cfg.get('val', 'val')) + + print("[Info] 正在构建验证集 DataLoader...") + val_loader = create_dataloader( + path=val_dir, + imgsz=args.imgsz, + batch_size=args.batch, + task='segment', + is_training=False, + num_workers=args.workers + ) + + # 实例化模型并加载权重 + print(f"[Info] 初始化分割模型架构 (Scale: {args.scale}, Classes: {args.nc})") + config = YOLOConfig(yaml_path=args.model_cfg, scale=args.scale, task="segment") + config.nc = args.nc + model = YOLO11ForSegmentation(config) + + if os.path.exists(args.weights): + print(f"[Info] 正在加载验证权重: {args.weights}") + param_dict = ms.load_checkpoint(args.weights) + ms.load_param_into_net(model, param_dict, strict_load=False) + else: + raise FileNotFoundError(f"[Error] 找不到权重文件: {args.weights}，请检查路径。") + + # 实例化 Validator 并执行验证 + print("[Info] 启动分割验证流水线...") + validator = SegmentationValidator(dataloader=val_loader, save_dir=args.save_dir, args=args, names=names) + stats = validator(model=model) + + # 提取双重指标 (Bounding Box & Mask) + box_mAP50 = stats.get('metrics/mAP50(B)', 0.0) + box_mAP50_95 = stats.get('metrics/mAP50-95(B)', 0.0) + mask_mAP50 = stats.get('metrics/mAP50(M)', 0.0) + mask_mAP50_95 = stats.get('metrics/mAP50-95(M)', 0.0) + + print("\n" + "-"*50) + print("[Result] 分割验证任务完成。核心指标如下：") + print(" [Box 目标检测]") + print(f" - mAP@50: {box_mAP50:.4f}") + print(f" - mAP@50-95: {box_mAP50_95:.4f}") + print(" [Mask 实例分割]") + print(f" - mAP@50: {mask_mAP50:.4f}") + print(f" - mAP@50-95: {mask_mAP50_95:.4f}") + print("-" * 50 + "\n") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/modeling_yolo.py b/src/mindnlp/models/ultralytics/modeling_yolo.py new file mode 100644 index 000000000..b5e4f1eb1 --- /dev/null +++ b/src/mindnlp/models/ultralytics/modeling_yolo.py @@ -0,0 +1,238 @@ +import math +import numpy as np +import mindspore as ms +from mindspore import nn, ops + +#from mindnlp.core.nn import PreTrainedModel + +try: + # 尝试导入官方的预训练基类 + from mindnlp.core.nn import PreTrainedModel +except ImportError as e: + print(f"⚠️ 警告: 无法从 mindnlp 导入 PreTrainedModel (可能是 mindtorch 版本冲突)。启动降级模式。错误: {e}") + # 本地降级方案：用一个假的基类糊弄过去，保证 train.py 能正常跑 + class PreTrainedModel(nn.Cell): + config_class = None + def __init__(self, config): + super().__init__() + self.config = config + +from configuration_yolo import YOLOConfig +from modules import ConvNormAct, C3k2, C2PSA, Classify, YOLO11DetectHead, YOLO11Segment, YOLO11Pose, Concat, Identity, SPPF, Upsample + +# 模块映射字典 + +MODULE_MAP = { + 'Conv': ConvNormAct, + 'C3k2': C3k2, + 'C2PSA': C2PSA, + 'SPPF': SPPF, + 'Concat': Concat, + 'nn.Upsample': Upsample, + 'Identity': Identity, + 'Detect': YOLO11DetectHead, + 'Segment': YOLO11Segment, + 'Pose': YOLO11Pose, + 'Classify': Classify +} + +class YOLO11Base(PreTrainedModel): + """ + YOLO11 通用基类，负责解析 YAML 配置、搭建网络拓扑结构以及权重的初始化。 + """ + + # 绑定配置类，使得 from_pretrained 能够自动解析对应的 Config + config_class = YOLOConfig + + def __init__(self, config): + super().__init__(config) + self.config = config + self.nc = config.nc + + # 动态解析网络结构 + self.model, self.save = self._parse_model(config) + self._init_weights() + self._init_logits_weights() + + def _init_weights(self): + """统一的权重初始化，包含针对 DFL 层的特殊跳过逻辑""" + for name, cell in self.cells_and_names(): + # DFL (Distribution Focal Loss) 层的权重是固定的预设值，不能被随机初始化覆盖 + if "dfl" in name.lower(): + continue + + if isinstance(cell, nn.Conv2d): + # 使用 HeUniform 初始化卷积层，适配 ReLU/SiLU 激活函数 + cell.weight.set_data(ms.common.initializer.initializer( + ms.common.initializer.HeUniform(negative_slope=math.sqrt(5)), + cell.weight.shape, cell.weight.dtype)) + elif isinstance(cell, nn.BatchNorm2d): + cell.gamma.set_data(ms.common.initializer.initializer('ones', cell.gamma.shape)) + cell.beta.set_data(ms.common.initializer.initializer('zeros', cell.beta.shape)) + + def _init_logits_weights(self): + """ + 初始化分类/检测头的偏置项 + 防止网络训练初期 Loss 爆炸 + """ + head = self.model[-1] if hasattr(self, 'model') else None + + if head and hasattr(head, 'cv3'): + for conv_seq in head.cv3: + last_conv = None + # 倒序查找包含 bias 的最后一层卷积 + cells = list(conv_seq.cells_and_names()) + for _, cell in reversed(cells): + if isinstance(cell, nn.Conv2d): + last_conv = cell + break + + if last_conv is not None and last_conv.has_bias: + # 使用 -4.59 填充 bias + new_bias = np.full(last_conv.bias.shape, -4.59) + last_conv.bias.set_data(ms.Tensor(new_bias, ms.float32)) + + def _parse_model(self, config): + """根据配置项动态生成网络层和拓扑结构""" + gd = config.depth_multiple + gw = config.width_multiple + max_ch = getattr(config, 'max_channels', 1024) or 1024 + + if not hasattr(config, 'yaml_dict'): + raise ValueError("Config validation failed: `yaml_dict` is missing. Ensure YOLOConfig is loaded correctly.") + + layers_cfg = config.yaml_dict['backbone'] + config.yaml_dict['head'] + layers, ch, save = [], [3], [] + + # 提取各个层的拓扑连接信息，避免在 construct 中动态读取 cell.f 和 cell.i + self.layer_f = [] + self.layer_i = [] + + for i, (f, n, m_name, args) in enumerate(layers_cfg): + m = MODULE_MAP.get(m_name, m_name) if isinstance(m_name, str) else m_name + n = max(round(n * gd), 1) if n > 1 else n + + if isinstance(f, list): + c1 = [ch[x + 1 if x >= 0 else x] for x in f] + save.extend(x for x in f if x != -1) + + if m in (YOLO11DetectHead, YOLO11Segment, YOLO11Pose): + _reg_max = 16 + _stride = [8, 16, 32] + + if m is YOLO11DetectHead: + _reg_max = args[1] if len(args) > 1 and isinstance(args[1], int) else 16 + m_ = m(nc=self.nc, reg_max=_reg_max, stride=_stride, ch=c1) + + elif m is YOLO11Segment: + _nm = args[1] if len(args) > 1 and isinstance(args[1], int) else 32 + _npr = args[2] if len(args) > 2 and isinstance(args[2], int) else 256 + _reg_max = args[3] if len(args) > 3 and isinstance(args[3], int) else 16 + m_ = m(nc=self.nc, reg_max=_reg_max, nm=_nm, npr=_npr, stride=_stride, ch=c1) + + elif m is YOLO11Pose: + _kpt_shape = args[1] if len(args) > 1 and isinstance(args[1], (list, tuple)) else getattr(config, 'kpt_shape', [17, 3]) + _reg_max = args[2] if len(args) > 2 and isinstance(args[2], int) else 16 + m_ = m(nc=self.nc, kpt_shape=_kpt_shape, reg_max=_reg_max, stride=_stride, ch=c1) + + c2 = ch[-1] + elif m is Concat: + m_ = m(*args) + c2 = sum(c1) + else: + m_ = m(*args) + c2 = ch[-1] + elif m is Upsample: + m_ = m(size=args[0], scale_factor=args[1], mode=args[2]) + c2 = ch[-1] + else: + c1 = ch[f] + if m is Classify: + m_ = m(c1, self.nc, *args) + c2 = self.nc + else: + c2 = math.ceil(min(args[0], max_ch) * gw / 8) * 8 + new_args = [c1, c2, n, *args[1:]] if m in (C3k2, C2PSA) else [c1, c2, *args[1:]] + m_ = m(*new_args) if n == 1 else nn.SequentialCell([m(*new_args) for _ in range(n)]) + + # 使用列表记录拓扑关系，对静态图更友好 + self.layer_i.append(i) + self.layer_f.append(f) + + layers.append(m_) + ch.append(c2) + + return nn.CellList(layers), sorted(save) + + def construct(self, x): + """前向传播逻辑""" + y = [] + for i, cell in enumerate(self.model): + f = self.layer_f[i] + + # 处理多输入或者跳连特征获取 + if f != -1: + x = y[f] if isinstance(f, int) else [x if j == -1 else y[j] for j in f] + + x = cell(x) + + y.append(x if self.layer_i[i] in self.save else None) + return x + +class YOLO11ForObjectDetection(YOLO11Base): + def construct(self, x): + res = super().construct(x) + + if self.training: + # 训练模式 + return res + + # 推理模式 + x = res[0] if isinstance(res, (list, tuple)) else res + if x.ndim == 2: + x = ops.expand_dims(x, 0) + if x.shape[1] > x.shape[2]: + x = x.transpose(0, 2, 1) + return x + +class YOLO11ForSegmentation(YOLO11Base): + """YOLO11 实例分割模型""" + def construct(self, x): + # 分割 Head 这里的 x 通常返回 (pred, proto) + x = super().construct(x) + if not self.training: + pred, proto = x[0], x[1] + if pred.shape[1] > pred.shape[2]: + pred = pred.transpose(0, 2, 1) + return pred, proto + return x + +class YOLO11ForClassification(YOLO11Base): + """YOLO11 图像分类模型""" + def construct(self, x, labels=None): + x = super().construct(x) + if labels is not None: + # 增加对 labels 的支持，返回 (Loss, Logits) 以兼容 Hugging Face/MindNLP Trainer + loss = nn.CrossEntropyLoss()(x, labels) + return (loss, x) + return x + +class YOLO11ForPose(YOLO11Base): + """YOLO11 姿态估计模型""" + def construct(self, x): + x = super().construct(x) + + if not self.training: + if isinstance(x, (list, tuple)): + pred = x[0] # 取出主预测结果 + else: + pred = x + + # 确保预测张量是 [Batch, Channels, Anchors] + if pred.ndim == 3 and pred.shape[1] > pred.shape[2]: + pred = pred.transpose(0, 2, 1) + + # 包装成列表返回，适配官方 Validator + return [pred] + + return x \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/__init__.py b/src/mindnlp/models/ultralytics/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/__init__.py b/src/mindnlp/models/ultralytics/models/yolo/classify/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/predict.py b/src/mindnlp/models/ultralytics/models/yolo/classify/predict.py new file mode 100644 index 000000000..8a197cabb --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/classify/predict.py @@ -0,0 +1,121 @@ +import os +import numpy as np +import cv2 +from PIL import Image, ImageDraw +import mindspore as ms +from mindspore import Tensor, ops + +from engine.predictor import BasePredictor + +# 分类结果实体类 +class Results: + """图像分类结果的标准化封装实体，提供置信度解析与可视化保存接口""" + + def __init__(self, orig_img, probs, names=None): + self.orig_img = orig_img + self.probs = probs + self.names = names or {i: f"class_{i}" for i in range(len(probs))} + + @property + def top1(self): + """获取最高置信度类别的索引与概率值""" + idx = int(np.argmax(self.probs)) + return idx, float(self.probs[idx]) + + @property + def top1_name(self): + """获取最高置信度类别的映射名称""" + return self.names.get(self.top1[0], str(self.top1[0])) + + def save(self, save_dir=".", file_name="result.jpg"): + """将分类结果绘制于原图并保存""" + # 类型安全防御：确保图像格式兼容 PIL + if isinstance(self.orig_img, np.ndarray): + img_drawn = Image.fromarray(cv2.cvtColor(self.orig_img, cv2.COLOR_BGR2RGB)) + else: + img_drawn = self.orig_img.copy() + + draw = ImageDraw.Draw(img_drawn) + + idx, conf = self.top1 + text = f"Pred: {self.top1_name}\nConf: {conf:.4f}" + + # 绘制背景标签底板与文字 + draw.rectangle([0, 0, 150, 40], fill="black") + draw.text((5, 5), text, fill="white") + + os.makedirs(save_dir, exist_ok=True) + save_path = os.path.join(save_dir, file_name) + img_drawn.save(save_path) + + +# 分类任务预测器 +class ClassificationPredictor(BasePredictor): + """图像分类任务专用预测器""" + + def __init__(self, cfg=None): + super().__init__(cfg) + self.imgsz = getattr(self.cfg, 'imgsz', 224) + self.names = {} + + def setup_model(self, model, ckpt_path=None): + """加载模型权重，并提取类别映射字典以供可视化使用""" + super().setup_model(model, ckpt_path) + + if hasattr(self.model, 'names'): + self.names = self.model.names + else: + nc = getattr(self.model.config, 'nc', 1000) if hasattr(self.model, 'config') else 1000 + self.names = {i: f"class_{i}" for i in range(nc)} + + def preprocess(self, img_source): + """ + 预处理流水线：应用中心裁剪与 ImageNet 标准化，替代基础的 LetterBox 缩放 + """ + if isinstance(img_source, str): + img = Image.open(img_source).convert('RGB') + else: + # 兼容 ndarray 类型的输入 + img = Image.fromarray(cv2.cvtColor(img_source, cv2.COLOR_BGR2RGB)) + + orig_img = img.copy() + + # 1. 保持宽高比，缩放短边至 256 + w, h = img.size + short_side = 256 + if w < h: + new_w, new_h = short_side, int(short_side * h / w) + else: + new_w, new_h = int(short_side * w / h), short_side + img_resized = img.resize((new_w, new_h), Image.BILINEAR) + + # 2. 中心裁剪至目标输入尺寸 (默认 224) + crop_size = self.imgsz + left = (new_w - crop_size) / 2 + top = (new_h - crop_size) / 2 + img_cropped = img_resized.crop((left, top, left + crop_size, top + crop_size)) + + # 3. 像素值归一化与 ImageNet 数据集统计分布标准化 + img_data = np.array(img_cropped).astype(np.float32) / 255.0 + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + img_data = (img_data - mean) / std + + # 4. 调整通道维度并扩增 Batch 轴 + img_data = img_data.transpose(2, 0, 1) + im_tensor = Tensor(img_data, ms.float32).expand_dims(0) + + return im_tensor, orig_img, None + + def postprocess(self, preds, orig_img, preprocess_info): + """ + 解析分类模型输出特征：提取 Softmax 概率分布并封装结果实体 + """ + if isinstance(preds, (tuple, list)): + preds = preds[0] + + # 提取 Batch 首位的分类 logits 并转化为概率空间 + probs = ops.softmax(preds, axis=1).asnumpy()[0] + + res = Results(orig_img=orig_img, probs=probs, names=self.names) + return res \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/train.py b/src/mindnlp/models/ultralytics/models/yolo/classify/train.py new file mode 100644 index 000000000..02b814538 --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/classify/train.py @@ -0,0 +1,69 @@ +import sys +import os +import yaml +from pathlib import Path +import mindspore as ms +from mindspore import nn + +# --- 导入底层架构和组件 --- +from engine.trainer import BaseTrainer +from models.yolo.classify.val import ClassificationValidator +from configuration_yolo import YOLOConfig +from modeling_yolo import YOLO11ForClassification +from data.loaders import create_dataloader +from utils.loss import YOLOClassificationLoss +from utils.optimizer import build_optimizer, get_lr + +class ClassificationTrainer(BaseTrainer): + """图像分类任务专属 Trainer，实现基类的抽象方法""" + + def __init__(self, args): + # 优先解析数据集配置 YAML，后续 DataLoader 与模型初始化需依赖此类信息 + with open(args.data, 'r', encoding='utf-8') as f: + self.data = yaml.safe_load(f) + + # 启动基类的全局构建流程 + super().__init__(args) + + def get_dataloader(self, is_training): + split_key = 'train' if is_training else 'val' + dataset_path = os.path.join(self.data.get('path', ''), self.data[split_key]) + + return create_dataloader( + dataset_path, + imgsz=self.args.imgsz, + batch_size=self.args.batch, + is_training=is_training, + num_workers=self.args.workers + ) + + def get_model(self, cfg=None, weights=None): + config = YOLOConfig(yaml_path=cfg, scale=self.args.scale, task='classify') + config.nc = self.data.get('nc', 1000) + model = YOLO11ForClassification(config) + + if weights and os.path.exists(weights): + print(f"[INFO] 检测到预训练权重: {weights}，启动微调模式。") + param_dict = ms.load_checkpoint(weights) + + # 分类任务微调时，通常需要剥离原有的分类输出层以适应新的类别数 + pop_keys = [k for k in param_dict.keys() if "model.10" in k] + for key in pop_keys: + param_dict.pop(key) + + ms.load_param_into_net(model, param_dict, strict_load=False) + print("[INFO] 预训练 Backbone 参数加载完毕，已移除旧的分类头。") + + return model + + def build_optimizer(self, model): + """生成学习率策略并构建优化器""" + steps_per_epoch = self.train_loader.get_dataset_size() + lr_list = get_lr(self.args, self.hyp, steps_per_epoch) + return build_optimizer(model, lr_list, self.hyp) + + def get_loss_fn(self): + return YOLOClassificationLoss() + + def get_validator(self): + return ClassificationValidator(dataloader=self.val_loader, args=self.args) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/val.py b/src/mindnlp/models/ultralytics/models/yolo/classify/val.py new file mode 100644 index 000000000..159a030a2 --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/classify/val.py @@ -0,0 +1,88 @@ +import logging +from types import SimpleNamespace + +import mindspore as ms +from mindspore import ops + +from engine.validator import BaseValidator +from data.loaders import create_dataloader +from utils.metrics import ClassifyMetrics + +LOGGER = logging.getLogger(__name__) + +class ClassificationValidator(BaseValidator): + """图像分类任务验证器，继承并实现基类的相关接口""" + + def __init__(self, dataloader=None, save_dir=None, args=None): + super().__init__(dataloader, save_dir, args) + + # 确保任务类型正确声明 + if self.args is None: + self.args = SimpleNamespace(task="classify", half=False) + else: + self.args.task = "classify" + + self.names = None + + def get_dataloader(self, dataset_path, batch_size=16): + """构建分类任务专用 DataLoader""" + return create_dataloader( + path=dataset_path, + imgsz=getattr(self.args, 'imgsz', 224), + batch_size=batch_size, + task='classify', + is_training=False, + num_workers=getattr(self.args, 'workers', 8) + ) + + def init_metrics(self, model): + """初始化分类评价指标统计器""" + self.names = getattr(model, 'names', {}) + self.nc = len(self.names) if self.names else 0 + + self.metrics = ClassifyMetrics() + self.targets = [] + self.preds = [] + + def preprocess(self, batch): + """分类任务的特殊预处理：支持半精度推理""" + batch["image"] = ops.cast(batch["image"], ms.float32) + + # 若开启混合精度/半精度验证，转换数据类型 + use_half = self.hyp.get('half', getattr(self.args, 'half', False)) + if use_half: + batch["image"] = ops.cast(batch["image"], ms.float16) + + return batch + + def postprocess(self, preds): + """分类网络输出剥离，无需进行 NMS""" + return preds[0] if isinstance(preds, (list, tuple)) else preds + + def update_metrics(self, preds, batch): + """提取预测结果与真实标签，存储至列表以供后续指标计算""" + target = batch.get("label", batch.get("cls")) + + # 获取置信度最高的 Top-5 索引 + n5 = min(preds.shape[1], 5) + _, topk_indices = ops.topk(preds, n5, dim=1) + + self.preds.append(topk_indices.asnumpy()) + self.targets.append(target.asnumpy()) + + def get_stats(self): + """触发 Metrics 类计算最终的验证指标，并合并统计结果""" + self.metrics.process(self.targets, self.preds) + metrics_dict = self.metrics.results_dict + + stats = super().get_stats() if hasattr(super(), 'get_stats') else {} + stats.update(metrics_dict) + return stats + + def print_results(self): + super().print_results() + stats = self.get_stats() + LOGGER.info( + f"验证结果 | Top-1 Acc: {stats.get('metrics/accuracy_top1', 0):.4f} | " + f"Top-5 Acc: {stats.get('metrics/accuracy_top5', 0):.4f}" + ) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/__init__.py b/src/mindnlp/models/ultralytics/models/yolo/detect/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/predict.py b/src/mindnlp/models/ultralytics/models/yolo/detect/predict.py new file mode 100644 index 000000000..5d66e73f6 --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/detect/predict.py @@ -0,0 +1,92 @@ +import os +import numpy as np +import cv2 +import mindspore as ms + +from engine.predictor import BasePredictor +from utils.ops import non_max_suppression + +# 结果封装与可视化类 +class Results: + """ + 目标检测结果封装类 + 提供统一的内部接口以供调用方提取检测框、分类结果及可视化 + """ + def __init__(self, orig_img, det, names): + self.orig_img = orig_img + self.det = det # [N, 6] 格式：x1, y1, x2, y2, conf, cls + self.names = names + + def save(self, save_dir=".", file_name="result.jpg"): + """ + 在原始图像上绘制边界框及置信度标签并持久化保存 + """ + res_img = self.orig_img.copy() + + for *xyxy, conf, cls in self.det: + bx1, by1, bx2, by2 = map(int, xyxy) + cls_id = int(cls) + label_text = self.names.get(cls_id, f"ID:{cls_id}") + + # 绘制绿色边界框 + color = (0, 255, 0) + cv2.rectangle(res_img, (bx1, by1), (bx2, by2), color, 2) + + # 绘制带背景底色的文本标签，提升可视化对比度 + text = f"{label_text} {conf:.2f}" + (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) + cv2.rectangle(res_img, (bx1, by1 - 20), (bx1 + tw, by1), color, -1) + cv2.putText(res_img, text, (bx1, by1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) + + os.makedirs(save_dir, exist_ok=True) + save_path = os.path.join(save_dir, file_name) + cv2.imwrite(save_path, res_img) + +# 检测任务专用推理算子 +class DetectionPredictor(BasePredictor): + """ + YOLO11 目标检测专属推理调度器 + 基于 BasePredictor 的预处理流水线 + """ + def setup_model(self, model, ckpt_path=None): + """挂载模型结构并提取类别字典映射""" + super().setup_model(model, ckpt_path) + + if hasattr(self.model, 'names'): + self.names = self.model.names + else: + nc = getattr(self.model.config, 'nc', 80) if hasattr(self.model, 'config') else 80 + self.names = {i: f"class_{i}" for i in range(nc)} + + def postprocess(self, preds, orig_img, prep_info): + """ + 检测任务后处理核心逻辑：原始输出解析 -> NMS 算子 -> 坐标映射还原 -> 结果实体封装 + """ + # 兼容性处理 + if isinstance(preds, (list, tuple)): + preds = preds[0] + + # 执行非极大值抑制，依据超参数配置筛除低置信度及高度重叠框 + preds_nms = non_max_suppression( + preds, + conf_thres=self.conf_thres, + iou_thres=self.iou_thres + ) + + # 提取单一批次的检测结果，形状转化为 NumPy Array 以加速标量计算 + det = preds_nms[0].asnumpy() if len(preds_nms) > 0 else np.zeros((0, 6)) + + # 逆向尺度还原：将在网络输入尺度下的边界框映射回原始图像尺度 + ratio, pad_w, pad_h = prep_info + if len(det) > 0: + det[:, 0] = (det[:, 0] - pad_w) / ratio + det[:, 1] = (det[:, 1] - pad_h) / ratio + det[:, 2] = (det[:, 2] - pad_w) / ratio + det[:, 3] = (det[:, 3] - pad_h) / ratio + + h0, w0 = orig_img.shape[:2] + det[:, [0, 2]] = np.clip(det[:, [0, 2]], 0, w0) + det[:, [1, 3]] = np.clip(det[:, [1, 3]], 0, h0) + + return Results(orig_img, det, self.names) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/train.py b/src/mindnlp/models/ultralytics/models/yolo/detect/train.py new file mode 100644 index 000000000..e961b436b --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/detect/train.py @@ -0,0 +1,121 @@ +import os +import yaml +import logging +import mindspore as ms + +from modeling_yolo import YOLO11ForObjectDetection +from configuration_yolo import YOLOConfig +from data.loaders import create_dataloader +from engine.trainer import BaseTrainer +from utils.optimizer import build_optimizer, get_lr +from utils.loss import v8DetectionLoss + +# 统一日志配置 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +LOGGER = logging.getLogger(__name__) + +class DetectionTrainer(BaseTrainer): + """ + 目标检测任务专用 Trainer，继承并实现基类抽象方法 + """ + def __init__(self, args): + super().__init__(args=args) + # 预定义检测任务的三大损失项，用于日志格式化 + self.loss_names = ["box_loss", "cls_loss", "dfl_loss"] + + def get_dataloader(self, is_training=True): + """ + 构建目标检测专用的数据流水线 + """ + with open(self.args.data, 'r', encoding='utf-8') as f: + data_cfg = yaml.safe_load(f) + + split = 'train' if is_training else 'val' + base_path = data_cfg.get('path', '') + sub_path = data_cfg.get(split, '') + dataset_path = os.path.join(base_path, sub_path) if base_path else sub_path + + return create_dataloader( + path=dataset_path, + imgsz=getattr(self.args, 'imgsz', 640), + batch_size=getattr(self.args, 'batch', 16), + task='detect', + is_training=is_training, + num_workers=getattr(self.args, 'workers', 8), + hyp=self.hyp + ) + + def build_optimizer(self, model): + """ + 基于统一调度策略构建优化器与学习率衰减列表 + """ + steps_per_epoch = self.train_loader.get_dataset_size() + lr_list = get_lr(self.args, self.hyp, steps_per_epoch) + return build_optimizer(model, lr_list, self.hyp) + + def get_model(self, cfg=None, weights=None): + """ + 构建目标检测网络拓扑并对齐预训练权重 + """ + config = YOLOConfig(yaml_path=self.args.model_cfg, scale=self.args.scale, task='detect') + + if hasattr(self.args, 'nc'): + config.nc = self.args.nc + + model = YOLO11ForObjectDetection(config) + + # 为网络注入 Loss 计算所需的核心架构常数 + model.nc = config.nc + model.reg_max = getattr(config, 'reg_max', 16) + model.stride = ms.Tensor([8, 16, 32], dtype=ms.float32) + + if weights and os.path.exists(weights): + LOGGER.info(f"正在加载预训练权重: {weights}") + param_dict = ms.load_checkpoint(weights) + param_not_load, _ = ms.load_param_into_net(model, param_dict, strict_load=False) + + # 生成标准的权重审计日志 + LOGGER.info("-" * 40) + LOGGER.info("[权重加载审计报告]") + LOGGER.info(f"模型参数总量: {len(model.trainable_params())}") + LOGGER.info(f"未匹配/未加载参数数量: {len(param_not_load)}") + if len(param_not_load) > 0: + LOGGER.info("未匹配参数抽样清单:") + for p in param_not_load[:10]: + LOGGER.info(f" -> {p}") + LOGGER.info("-" * 40) + + return model + + def get_loss_fn(self): + """ + 实例化目标检测专用损失计算图 + 提取检测头 (Detect Head) 作为基础组件传递 + """ + model_head = self.model.model[-1] + return v8DetectionLoss(model_head) + + def get_validator(self): + """ + 实例化目标检测验证器 + """ + from models.yolo.detect.val import DetectionValidator + + # 净化传递给验证器的配置字典，剔除不相关或由 yaml 管理的键值 + args_dict = vars(self.args).copy() + custom_keys = ['nc', 'model_cfg', 'scale', 'weights'] + for key in custom_keys: + args_dict.pop(key, None) + + args_dict['mode'] = 'val' + args_dict['task'] = 'detect' + + return DetectionValidator(dataloader=self.val_loader, save_dir=self.save_dir, args=self.args) + + def label_loss_items(self, loss_items=None, prefix="train"): + """格式化输出 Loss 指标""" + keys = [f"{prefix}/{x}" for x in self.loss_names] + if loss_items is not None: + loss_items = [round(float(x), 5) for x in loss_items] + return dict(zip(keys, loss_items)) + return keys \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/val.py b/src/mindnlp/models/ultralytics/models/yolo/detect/val.py new file mode 100644 index 000000000..bd186339f --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/detect/val.py @@ -0,0 +1,112 @@ +import logging +import numpy as np +import mindspore as ms +from mindspore import ops + +from engine.validator import BaseValidator +from utils.metrics import DetMetrics, process_batch +from utils.ops import non_max_suppression, xywh2xyxy_np +from data.loaders import create_dataloader + +LOGGER = logging.getLogger(__name__) + +class DetectionValidator(BaseValidator): + """ + 目标检测任务专用验证器 + 继承 BaseValidator + """ + def __init__(self, dataloader=None, save_dir=None, args=None): + super().__init__(dataloader, save_dir, args) + # 初始化 COCO 标准的 10 个 IoU 评估阈值 (0.50 到 0.95，步长 0.05) + self.iou_v = np.linspace(0.5, 0.95, 10) + self.metrics = None + self.names = None + self.stats = {} + + def get_dataloader(self, dataset_path, batch_size=16): + """ + 构建目标检测任务数据流 + 使用 is_training=False 关闭随机数据增强 + """ + return create_dataloader( + path=dataset_path, + imgsz=getattr(self.args, 'imgsz', 640), + batch_size=batch_size, + task='detect', + is_training=False, + num_workers=getattr(self.args, 'workers', 8) + ) + + def init_metrics(self, model): + """初始化评估指标计算器与类别映射表""" + self.names = getattr(model, 'names', {i: f'class_{i}' for i in range(getattr(self.args, 'nc', 80))}) + self.metrics = DetMetrics(names=self.names) + + def preprocess(self, batch): + """图像预处理：标准化像素值域至 [0, 1]""" + img_key = "img" if "img" in batch else "image" + img = ops.cast(batch[img_key], ms.float32) / 255.0 + batch[img_key] = img + return batch + + def postprocess(self, preds): + """ + 验证期后处理：执行非极大值抑制 (NMS) + 注：为计算完整的 Precision-Recall 曲线，验证阶段的 conf_thres 必须设定为极小值 (如 0.001) + """ + if isinstance(preds, (list, tuple)): + preds = preds[0] + + # IoU 阈值由 hyp.yaml 控制，置信度强制为 0.001 以计算 mAP + iou_thres = self.hyp.get('iou', 0.6) + preds_nms = non_max_suppression(preds, conf_thres=0.001, iou_thres=iou_thres) + return preds_nms + + def update_metrics(self, preds, batch): + """对比预测边界框与真实标签，统计各个阈值下的 True Positives""" + batch_idx = batch["batch_idx"].reshape(-1).asnumpy() + targets_bboxes = batch["bboxes"].asnumpy() + targets_cls = batch["cls"].reshape(-1).asnumpy() + + for si, pred in enumerate(preds): + # 1. 提取当前图像的 Ground Truth + mask = (batch_idx == si) + t_cls = targets_cls[mask] + t_box = targets_bboxes[mask] + + # 2. 构建格式化标签矩阵 [N, 5] -> [class, x1, y1, x2, y2] + if len(t_cls) > 0: + imgsz = getattr(self.args, 'imgsz', 640) + t_box_xyxy = xywh2xyxy_np(t_box) * imgsz + labels = np.concatenate((t_cls[:, None], t_box_xyxy), axis=1) + else: + labels = np.zeros((0, 5)) + + # 3. 提取当前图像的预测结果 + p_det = pred.asnumpy() if len(pred) > 0 else np.zeros((0, 6)) + + + # 4. 极端情况处理：无预测框 + if len(p_det) == 0: + if len(labels): + self.metrics.update_stats(np.zeros((0, 10), dtype=bool), np.zeros(0), np.zeros(0), labels[:, 0]) + continue + + # 5. 执行 IoU 匹配，计算 TP 矩阵 + tp = process_batch(p_det, labels, self.iou_v) + self.metrics.update_stats(tp, p_det[:, 4], p_det[:, 5], labels[:, 0]) + + def finalize_metrics(self): + """计算最终聚合的 mAP 指标""" + self.stats = self.metrics.process() + + def get_stats(self): + """整合信息与验证指标并返回""" + stats = super().get_stats() + stats.update(self.stats) + + # 兼容 Trainer 生命周期中的模型保存判定逻辑 + stats['fitness'] = stats.get('metrics/mAP50-95(B)', 0.0) + self.results_dict = stats + + return stats \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/__init__.py b/src/mindnlp/models/ultralytics/models/yolo/pose/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/predict.py b/src/mindnlp/models/ultralytics/models/yolo/pose/predict.py new file mode 100644 index 000000000..6ee1717db --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/pose/predict.py @@ -0,0 +1,108 @@ +import os +import numpy as np +import cv2 +import mindspore as ms +from engine.predictor import BasePredictor +from utils.ops import non_max_suppression + +# 预测结果封装与渲染类 +class Results: + """ + 姿态估计结果封装类 + 负责存储原始图像、边界框检测结果、关键点坐标，并提供可视化渲染与保存功能 + """ + def __init__(self, orig_img, det, kpts, names): + self.orig_img = orig_img + self.det = det # 边界框与类别信息: [N, 6] -> x1, y1, x2, y2, conf, cls + self.kpts = kpts # 关键点信息: [N, 17, 3] -> x, y, conf + self.names = names + + # COCO 数据集标准的 17 关键点物理连接拓扑结构 (骨架) + self.skeleton = [ + [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], + [5, 6], [5, 7], [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], + [1, 3], [2, 4], [3, 5], [4, 6] + ] + + def save(self, save_dir=".", file_name="result.jpg"): + """将边界框、有效关键点及其骨架连线渲染至原图并保存""" + res_img = self.orig_img.copy() + + for i, (*xyxy, conf, cls) in enumerate(self.det): + # 1. 渲染检测边界框与置信度 + bx1, by1, bx2, by2 = map(int, xyxy) + cv2.rectangle(res_img, (bx1, by1), (bx2, by2), (0, 0, 255), 2) + cv2.putText(res_img, f"{self.names.get(int(cls), 'Pose')} {conf:.2f}", + (bx1, max(by1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) + + # 2. 渲染关键点与骨架拓扑 + if i < len(self.kpts): + kpts = self.kpts[i] + + # 绘制有效关键点 (置信度阈值 > 0.5) + for kx, ky, kconf in kpts: + if kconf > 0.5: + cv2.circle(res_img, (int(kx), int(ky)), 5, (0, 255, 0), -1) + + # 绘制骨架连线 (仅当两端关键点均有效时连线) + for sk in self.skeleton: + p1, p2 = kpts[sk[0]], kpts[sk[1]] + if p1[2] > 0.5 and p2[2] > 0.5: + cv2.line(res_img, (int(p1[0]), int(p1[1])), (int(p2[0]), int(p2[1])), + (255, 255, 0), 2, cv2.LINE_AA) + + os.makedirs(save_dir, exist_ok=True) + save_path = os.path.join(save_dir, file_name) + cv2.imwrite(save_path, res_img) + print(f"[Info] 渲染结果已保存至: {save_path}") + +# 姿态预测器核心逻辑 +class PosePredictor(BasePredictor): + """ + YOLO11 姿态估计预测器子类 + 继承自 BasePredictor，扩展了关键点维度的后处理与坐标系逆映射逻辑 + """ + def setup_model(self, model, ckpt_path=None): + """初始化模型参数及拓扑配置""" + super().setup_model(model, ckpt_path) + self.nkpt = getattr(self.model.config, 'kpt_shape', [17, 3])[0] + self.names = getattr(self.model, 'names', {0: 'person'}) + + def postprocess(self, preds, orig_img, prep_info): + """ + 推理后处理：执行 NMS，并将坐标逆映射回原始图像尺度 + """ + if isinstance(preds, (list, tuple)): + preds = preds[0] + + # 1. 执行非极大值抑制 (NMS) + preds_nms = non_max_suppression( + preds, + conf_thres=self.conf_thres, + iou_thres=self.iou_thres, + nc=len(self.names) + ) + + det = preds_nms[0].asnumpy() if len(preds_nms) > 0 else np.zeros((0, 6 + self.nkpt * 3)) + if len(det) == 0: + return Results(orig_img, np.zeros((0, 6)), [], self.names) + + # 2. 获取预处理时的缩放比例与填充信息 + ratio, pad_w, pad_h = prep_info + + # 3. 逆仿射变换：还原边界框坐标 + det[:, :4] = (det[:, :4] - [pad_w, pad_h, pad_w, pad_h]) / ratio + + # 4. 逆仿射变换：还原关键点坐标并重塑维度 + kpts = det[:, 6:].reshape(-1, self.nkpt, 3) + kpts[..., 0] = (kpts[..., 0] - pad_w) / ratio + kpts[..., 1] = (kpts[..., 1] - pad_h) / ratio + + # 5. 越界裁剪：确保所有坐标不超出原始图像边界 + h0, w0 = orig_img.shape[:2] + det[:, [0, 2]] = np.clip(det[:, [0, 2]], 0, w0) + det[:, [1, 3]] = np.clip(det[:, [1, 3]], 0, h0) + kpts[..., 0] = np.clip(kpts[..., 0], 0, w0) + kpts[..., 1] = np.clip(kpts[..., 1], 0, h0) + + return Results(orig_img, det[:, :6], kpts, self.names) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/train.py b/src/mindnlp/models/ultralytics/models/yolo/pose/train.py new file mode 100644 index 000000000..db6fafd3f --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/pose/train.py @@ -0,0 +1,118 @@ +import os +import yaml +import mindspore as ms +from mindspore import nn, ops + +from engine.trainer import BaseTrainer +from modeling_yolo import YOLO11ForPose, YOLOConfig +from utils.loss import v8PoseLoss +from utils.optimizer import build_optimizer, get_lr +from data.loaders import create_dataloader + +class PoseTrainer(BaseTrainer): + """ + 姿态估计(Pose)任务专属 Trainer + 继承自 BaseTrainer，封装了数据加载、模型初始化、损失构建及优化器配置等底层训练循环 + """ + def __init__(self, args): + # 解析数据集 YAML 获取类别数和关键点形状 + with open(args.data, 'r', encoding='utf-8') as f: + self.data = yaml.safe_load(f) + + # 解析超参数 YAML (hyp.yaml) + self.hyp = {} + if hasattr(args, 'hyp') and args.hyp and os.path.exists(args.hyp): + with open(args.hyp, 'r', encoding='utf-8') as f: + self.hyp = yaml.safe_load(f) + + super().__init__(args) + + # 姿态估计专属的损失函数名称集 (包含边界框回归、分类、分布焦点、关键点坐标及置信度) + self.loss_names = ["box_loss", "cls_loss", "dfl_loss", "pose_loss", "kobj_loss"] + + def get_dataloader(self, is_training): + """实例化并返回姿态估计专属的 DataLoader""" + split_key = 'train' if is_training else 'val' + dataset_path = os.path.join(self.data.get('path', ''), self.data[split_key]) + + return create_dataloader( + dataset_path, + imgsz=self.args.imgsz, + batch_size=getattr(self.args, 'batch_size', self.args.batch), + task='pose', + is_training=is_training, + num_workers=getattr(self.args, 'workers', 8) + ) + + def get_model(self, cfg=None, weights=None, verbose=True): + """模型结构构建与预训练权重加载验证""" + config = YOLOConfig(yaml_path=self.args.model_cfg, scale=self.args.scale, task='pose') + + # 动态配置 Pose 专属属性 + config.nc = self.data.get('nc', 1) # 姿态估计通常默认为 1 个类 (person) + config.kpt_shape = self.data.get('kpt_shape', [17, 3]) + model = YOLO11ForPose(config) + + # 补充 Loss 计算依赖的核心属性，防止前向传播期间出现 AttributeError + model.nc = config.nc + model.kpt_shape = config.kpt_shape + model.reg_max = getattr(config, 'reg_max', 16) + model.stride = ms.Tensor([8, 16, 32], dtype=ms.float32) + + if weights and os.path.exists(weights): + if verbose: + print(f"[Info] 加载预训练权重进行微调: {weights}") + param_dict = ms.load_checkpoint(weights) + + # 开启 strict_load=False 以允许部分权重不匹配 (适用于网络结构微调) + param_not_load, ckpt_not_load = ms.load_param_into_net(model, param_dict, strict_load=False) + + if verbose: + print("-" * 50) + print("[Info] 预训练权重加载完整性校验：") + print(f" - 模型总参数量: {len(model.parameters_dict())}") + print(f" - 未能加载的参数数量: {len(param_not_load)}") + if len(param_not_load) > 0: + print(f"[Warning] 存在 {len(param_not_load)} 个参数未成功匹配。") + print(f" - 未加载参数示例: {param_not_load[:5]}") + print("-" * 50) + + return model + + def build_optimizer(self, model): + """基于 hyp.yaml 或默认配置构建学习率策略与优化器""" + steps_per_epoch = self.train_loader.get_dataset_size() + lr_list = get_lr(self.args, self.hyp, steps_per_epoch) + return build_optimizer(model, lr_list, self.hyp) + + def get_loss_fn(self): + """实例化姿态估计专用损失函数""" + criterion = v8PoseLoss(self.model) + criterion.imgsz = self.args.imgsz + return criterion + + def get_validator(self): + """实例化姿态验证器，用于训练周期间的性能评估""" + from models.yolo.pose.val import PoseValidator + + self.args.task = 'pose' + self.args.mode = 'val' + + names = self.data.get('names', {i: f'class_{i}' for i in range(self.data.get('nc', 1))}) + kpt_shape = self.data.get('kpt_shape', [17, 3]) + + return PoseValidator( + dataloader=self.val_loader, + save_dir=self.save_dir, + args=self.args, + names=names, + kpt_shape=kpt_shape + ) + + def label_loss_items(self, loss_items=None, prefix="train"): + """格式化输出损失字典，便于日志系统记录""" + keys = [f"{prefix}/{x}" for x in self.loss_names] + if loss_items is not None: + loss_items = [round(float(x), 5) for x in loss_items] + return dict(zip(keys, loss_items)) + return keys \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/val.py b/src/mindnlp/models/ultralytics/models/yolo/pose/val.py new file mode 100644 index 000000000..7ec6d091a --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/pose/val.py @@ -0,0 +1,219 @@ +import numpy as np +import mindspore as ms +from mindspore import ops + +from engine.validator import BaseValidator +from utils.ops import non_max_suppression, xywh2xyxy_np +from utils.metrics import PoseMetrics, kpt_iou + +class PoseValidator(BaseValidator): + """ + YOLO11 姿态估计专属验证器 + 继承自 BaseValidator，处理数据加载、前向推理、NMS后处理以及 OKS 指标计算逻辑 + """ + def __init__(self, dataloader=None, save_dir=None, args=None, names=None, kpt_shape=None): + super().__init__(dataloader, save_dir, args) + self.names = names if names else {0: 'person'} + self.nc = len(self.names) + + # 姿态专属属性 + self.kpt_shape = kpt_shape if kpt_shape else [17, 3] + self.nkpt = self.kpt_shape[0] + + # COCO 格式 OKS 计算所需的各个关键点容差常数 (Sigma) + self.sigma = np.array( + [.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89] + ) / 10.0 + + self.metrics = None + self.stats = {} + + def get_dataloader(self, dataset_path, batch_size=16): + """ + 验证集专属数据加载逻辑 + 传递 task='pose' 标识以触发数据后端生成 OKS 计算所需的完整关键点字段 + """ + return create_dataloader( + path=dataset_path, + imgsz=getattr(self.args, 'imgsz', 640), + batch_size=batch_size, + task='pose', + is_training=False, + num_workers=getattr(self.args, 'workers', 8) + ) + + def init_metrics(self, model): + """初始化性能评估指标容器""" + self.metrics = PoseMetrics(names=self.names) + + def preprocess(self, batch): + """ + 数据预处理阶段 + """ + for k in ["img", "image"]: + if k in batch: + input_tensor = batch[k] + batch[k] = ops.cast(input_tensor, ms.float32) / 255.0 + return batch + + def postprocess(self, preds): + """ + 后处理逻辑：执行维度对齐与非极大值抑制 (NMS)。 + """ + p = preds[0] if isinstance(preds, (tuple, list)) else preds + + # 维度翻转适配：当特征维度小于锚框维度时，转置为 NMS 期望输入的 [B, anchors, features] + if len(p.shape) == 3 and p.shape[1] < p.shape[2]: + p = p.swapaxes(1, 2) + + conf_thres = getattr(self.args, 'conf', 0.001) + iou_thres = getattr(self.args, 'iou', 0.6) + + out = non_max_suppression( + p, + conf_thres=conf_thres, + iou_thres=iou_thres, + nc=getattr(self, 'nc', 1) + ) + return out + + def update_metrics(self, preds, batch): + """基于模型预测输出与真实标签，计算并统计当前 Batch 的 IoU 与 OKS 匹配结果""" + batch_idx = batch["batch_idx"].view(-1).asnumpy() + gt_cls = batch["cls"].view(-1).asnumpy() + gt_bboxes = batch["bboxes"].asnumpy().copy() + gt_keypoints = batch["keypoints"].asnumpy().copy() + + imgsz = getattr(self.args, 'imgsz', 640) + + # --- 数据格式统一与坐标系对齐 --- + # 尺度还原：若检测到真实边界框为 [0, 1] 归一化坐标，则放大至绝对像素尺度 + valid_boxes = gt_bboxes[gt_bboxes > 0] + if len(valid_boxes) > 0 and valid_boxes.max() <= 1.01: + gt_bboxes[:, 0:4] *= imgsz + gt_keypoints[..., 0:2] *= imgsz + + # 格式转换：统一将目标框坐标转换为绝对像素级的 [x1, y1, x2, y2] 格式 + labels_box_xyxy = xywh2xyxy_np(gt_bboxes) + + # 边界约束：执行物理越界裁剪 + labels_box_xyxy[:, [0, 2]] = np.clip(labels_box_xyxy[:, [0, 2]], 0, imgsz) + labels_box_xyxy[:, [1, 3]] = np.clip(labels_box_xyxy[:, [1, 3]], 0, imgsz) + + # 逐图执行指标评估统计 + for i, pred in enumerate(preds): + idx_mask = (batch_idx == i) + labels_cls = gt_cls[idx_mask] + labels_box = labels_box_xyxy[idx_mask] + labels_kpt = gt_keypoints[idx_mask] + + pred_np = pred.asnumpy() if len(pred) > 0 else np.zeros((0, 6 + self.nkpt * 3)) + + # 漏检处理 + if len(pred_np) == 0: + if len(labels_cls) > 0: + self.metrics.update_stats( + tp_b=np.zeros((0, 10), dtype=bool), + tp_p=np.zeros((0, 10), dtype=bool), + conf=np.zeros(0), pred_cls=np.zeros(0), target_cls=labels_cls + ) + continue + + p_boxes = pred_np[:, :4] + p_conf = pred_np[:, 4] + p_cls = pred_np[:, 5] + p_kpts = pred_np[:, 6:].reshape(-1, self.nkpt, 3) + + tp_b = np.zeros((len(pred_np), 10), dtype=bool) + tp_p = np.zeros((len(pred_np), 10), dtype=bool) + + if len(labels_box) > 0: + # [A] 边界框 IoU 匹配计算 + iou_b = self.box_iou_np(p_boxes, labels_box) + tp_b = self.match_predictions(p_cls, labels_cls, iou_b) + + # [B] 关键点 OKS 匹配计算 + # 依据目标框计算近似包围面积 + w = labels_box[:, 2] - labels_box[:, 0] + h = labels_box[:, 3] - labels_box[:, 1] + area = (w * h) * 0.53 + + if len(p_kpts) > 0 and len(labels_kpt) > 0: + iou_p = kpt_iou( + ms.Tensor(p_kpts), + ms.Tensor(labels_kpt), + ms.Tensor(area), + self.sigma + ).asnumpy() + tp_p = self.match_predictions(p_cls, labels_cls, iou_p) + + # 更新当前迭代批次的评估状态 + self.metrics.update_stats(tp_b, tp_p, p_conf, p_cls, labels_cls) + + def finalize_metrics(self): + """验证循环终止阶段：计算评估指标（如 mAP）的最终结果""" + if len(self.metrics.stats) == 0: + print("[Warning] 验证过程未收集到有效目标数据，mAP 返回默认值 0.0。") + self.stats = {"metrics/mAP50-95(B)": 0.0, "metrics/mAP50-95(P)": 0.0} + return + + try: + self.stats = self.metrics.process() + except Exception as e: + print(f"[Error] 评估指标计算过程发生异常: {e}") + self.stats = {"metrics/mAP50-95(B)": 0.0, "metrics/mAP50-95(P)": 0.0} + + def get_stats(self): + """向 Trainer 反馈格式化后的各项验证统计数据""" + stats = super().get_stats() + stats.update(self.stats) + + if self.metrics is not None: + stats['fitness'] = self.metrics.fitness + else: + stats['fitness'] = 0.0 + self.results_dict = stats + return stats + + def box_iou_np(self, box1, box2, eps=1e-7): + """利用 NumPy 高效计算预测框与真实框之间的交并比 (IoU) 矩阵""" + b1_x1, b1_y1, b1_x2, b1_y2 = np.split(box1, 4, axis=1) + b2_x1, b2_y1, b2_x2, b2_y2 = np.split(box2, 4, axis=1) + + inter_area = (np.minimum(b1_x2, b2_x2.T) - np.maximum(b1_x1, b2_x1.T)).clip(0) * \ + (np.minimum(b1_y2, b2_y2.T) - np.maximum(b1_y1, b2_y1.T)).clip(0) + + area1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1) + area2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + + return inter_area / (area1 + area2.T - inter_area + eps) + + def match_predictions(self, pred_cls, target_cls, iou, iou_thresholds=None): + """采用贪心匹配算法，求解预测框与真实目标类别间的最佳一对一映射""" + if iou_thresholds is None: + iou_thresholds = np.linspace(0.5, 0.95, 10) + + tp = np.zeros((pred_cls.shape[0], iou_thresholds.shape[0]), dtype=bool) + correct_class = pred_cls[:, None] == target_cls[None, :] + iou = iou * correct_class + + for j, thr in enumerate(iou_thresholds): + matches = np.argwhere(iou >= thr) + if matches.shape[0] == 0: + continue + + match_scores = iou[matches[:, 0], matches[:, 1]] + match_data = np.concatenate((matches, match_scores[:, None]), axis=1) + + # 按匹配分数降序排序并执行去重操作 + match_data = match_data[match_data[:, 2].argsort()[::-1]] + _, unique_preds = np.unique(match_data[:, 0], return_index=True) + match_data = match_data[unique_preds] + + match_data = match_data[match_data[:, 2].argsort()[::-1]] + _, unique_gts = np.unique(match_data[:, 1], return_index=True) + match_data = match_data[unique_gts] + + tp[match_data[:, 0].astype(int), j] = True + + return tp \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py b/src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/predict.py b/src/mindnlp/models/ultralytics/models/yolo/segment/predict.py new file mode 100644 index 000000000..5cffc6cce --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/segment/predict.py @@ -0,0 +1,152 @@ +import os +import numpy as np +import cv2 +import mindspore as ms + +from engine.predictor import BasePredictor +from utils.ops import non_max_suppression + +# 实例分割结果实体与可视化类 +class Results: + """ + 分割结果封装类 + 提供在原图上叠印目标边界框及像素级掩码 (Mask) 的持久化保存方法 + """ + def __init__(self, orig_img, det, masks, names): + self.orig_img = orig_img + self.det = det # [N, 6] 格式：x1, y1, x2, y2, conf, cls + self.masks = masks # 包含 N 个掩码矩阵的列表，尺寸与 orig_img 一致 + self.names = names + + def save(self, save_dir=".", file_name="result.jpg"): + """将预测边界框与掩码着色后保存为图像文件""" + res_img = self.orig_img.copy() + + for i, (*xyxy, conf, cls) in enumerate(self.det): + bx1, by1, bx2, by2 = map(int, xyxy) + cls_id = int(cls) + label_text = self.names.get(cls_id, f"ID:{cls_id}") + + # 掩码渲染 (橙色通透叠加处理) + if i < len(self.masks): + mask_orig = self.masks[i] + color = np.array([255, 120, 0], dtype=np.uint8) # 设定掩码基础色 (BGR) + weight_mask = (np.clip(mask_orig, 0, 1) ** 2.0)[:, :, None] + overlay = (weight_mask * color).astype(np.uint8) + res_img = cv2.addWeighted(res_img, 1.0, overlay, 0.45, 0) + + # 计算并绘制掩码轮廓的平滑描边 + binary_mask = (mask_orig > 0.4).astype(np.uint8) + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1) + cv2.drawContours(res_img, contours, -1, (255, 255, 255), 1, lineType=cv2.LINE_AA) + + # 边界框与文本标签渲染 + cv2.rectangle(res_img, (bx1, by1), (bx2, by2), (255, 255, 255), 1) + text = f"{label_text} {conf:.2f}" + (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) + cv2.rectangle(res_img, (bx1, by1 - 20), (bx1 + tw, by1), (0, 255, 0), -1) + cv2.putText(res_img, text, (bx1, by1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) + + os.makedirs(save_dir, exist_ok=True) + save_path = os.path.join(save_dir, file_name) + cv2.imwrite(save_path, res_img) + + +# 实例分割推理算子 +class SegmentationPredictor(BasePredictor): + """ + YOLO11 实例分割任务预测器 + 扩展基类能力以解析掩码系数，并在特征映射后执行 Mask 合成 + """ + def setup_model(self, model, ckpt_path=None): + super().setup_model(model, ckpt_path) + if hasattr(self.model, 'names'): + self.names = self.model.names + else: + nc = getattr(self.model.config, 'nc', 80) if hasattr(self.model, 'config') else 80 + self.names = {i: f"class_{i}" for i in range(nc)} + + def postprocess(self, preds, orig_img, prep_info): + """ + 后处理解析：分离预测边界框与掩码系数 -> NMS -> 合成实例掩码 -> 坐标空间映射 + """ + # 拆解网络多头输出：(预测矩阵, 原型掩码特征图) + prediction, proto = preds + + # 应用非极大值抑制。内部会提取附带的 32 维 Mask 系数 + preds_nms = non_max_suppression( + prediction, + conf_thres=self.conf_thres, + iou_thres=self.iou_thres + ) + + det = preds_nms[0].asnumpy() if len(preds_nms) > 0 else np.zeros((0, 38)) + if len(det) == 0: + return Results(orig_img, np.zeros((0, 6)), [], self.names) + + # 剥离空间边界框数据与对应的特征系数 + boxes_640 = det[:, :4].copy() + mask_coeffs = det[:, 6:] + + # 在网络输入尺度 (通常为 640x640) 下合成掩码实例 + proto_data = proto.asnumpy()[0] + masks_640 = self._process_mask_ultimate_precision( + proto_data, mask_coeffs, boxes_640, (self.imgsz, self.imgsz) + ) + + # 执行边界框的逆向尺度还原 + ratio, pad_w, pad_h = prep_info + det[:, 0] = (det[:, 0] - pad_w) / ratio + det[:, 1] = (det[:, 1] - pad_h) / ratio + det[:, 2] = (det[:, 2] - pad_w) / ratio + det[:, 3] = (det[:, 3] - pad_h) / ratio + + h0, w0 = orig_img.shape[:2] + det[:, [0, 2]] = np.clip(det[:, [0, 2]], 0, w0) + det[:, [1, 3]] = np.clip(det[:, [1, 3]], 0, h0) + + # 掩码阵列反向映射：剥离 Padding 后缩放至原始图像空间 + final_masks_orig = [] + top, bottom = int(round(pad_h - 0.1)), int(round(self.imgsz - pad_h + 0.1)) + left, right = int(round(pad_w - 0.1)), int(round(self.imgsz - pad_w + 0.1)) + + for m_img in masks_640: + m_cropped = m_img[top:bottom, left:right] + mask_orig = cv2.resize(m_cropped, (w0, h0), interpolation=cv2.INTER_LINEAR) + final_masks_orig.append(mask_orig) + + return Results(orig_img, det[:, :6], final_masks_orig, self.names) + + def _process_mask_ultimate_precision(self, protos, masks_in, bboxes, shape): + """ + 利用矩阵乘法结合 Sigmoid 激活还原特征图概率，并应用边界框进行严格裁剪 + """ + c, mh, mw = protos.shape + ih, iw = shape + + masks = (masks_in @ protos.reshape(c, -1)).reshape(-1, mh, mw) + width_ratio, height_ratio = mw / iw, mh / ih + + scaled_bboxes = bboxes.copy() + scaled_bboxes[:, [0, 2]] *= width_ratio + scaled_bboxes[:, [1, 3]] *= height_ratio + + full_masks = [] + for i in range(len(masks)): + m_raw = masks[i] + m_min, m_max = m_raw.min(), m_raw.max() + m_norm = (m_raw - m_min) / (m_max - m_min + 1e-6) + m_prob = 1 / (1 + np.exp(-((m_norm - 0.5) * 10))) + + x1, y1, x2, y2 = scaled_bboxes[i] + r = np.arange(mw, dtype=np.float32)[None, :] + col = np.arange(mh, dtype=np.float32)[:, None] + + # 使用逻辑张量截断框外冗余特征 + m_prob = m_prob * ((r >= x1) * (r < x2) * (col >= y1) * (col < y2)) + + m_img = cv2.resize(m_prob, (iw, ih), interpolation=cv2.INTER_CUBIC) + full_masks.append(m_img) + + return full_masks \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/train.py b/src/mindnlp/models/ultralytics/models/yolo/segment/train.py new file mode 100644 index 000000000..8a99b8f49 --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/segment/train.py @@ -0,0 +1,103 @@ +import os +import yaml +import logging +import mindspore as ms + +# --- 核心架构组件 --- +from engine.trainer import BaseTrainer +from modeling_yolo import YOLO11ForSegmentation, YOLOConfig +from utils.loss import v8SegmentationLoss +from utils.optimizer import build_optimizer, get_lr +from data.loaders import create_dataloader + +LOGGER = logging.getLogger(__name__) + +class SegmentationTrainer(BaseTrainer): + """ + 实例分割任务专属 Trainer + 继承 BaseTrainer 以复用标准的训练循环 + """ + def __init__(self, args): + with open(args.data, 'r', encoding='utf-8') as f: + self.data = yaml.safe_load(f) + + super().__init__(args) + # 定义分割任务独有的 4 项损失函数 + self.loss_names = ["box_loss", "cls_loss", "dfl_loss", "seg_loss"] + + def get_dataloader(self, is_training): + """挂载分割任务特有的 DataLoader""" + split_key = 'train' if is_training else 'val' + dataset_path = os.path.join(self.data.get('path', ''), self.data[split_key]) + + return create_dataloader( + dataset_path, + imgsz=self.args.imgsz, + batch_size=self.args.batch_size, + task='segment', + is_training=is_training, + num_workers=getattr(self.args, 'workers', 8), + hyp=self.hyp # 传递超参数字典以驱动数据增强 + ) + + def get_model(self, cfg=None, weights=None, verbose=True): + """构建分割模型拓扑并加载检查点""" + config = YOLOConfig(yaml_path=self.args.model_cfg, scale=self.args.scale, task='segment') + config.nc = self.data.get('nc', 80) + model = YOLO11ForSegmentation(config) + + # 注入辅助特征解码的核心张量属性 + model.nc = config.nc + model.reg_max = getattr(config, 'reg_max', 16) + model.stride = ms.Tensor([8, 16, 32], dtype=ms.float32) + + if weights and os.path.exists(weights): + LOGGER.info(f"正在载入初始检查点权重: {weights}") + param_dict = ms.load_checkpoint(weights) + + # 使用 strict_load=False 容忍迁移学习阶段的分类头维度差异 + param_not_load, _ = ms.load_param_into_net(model, param_dict, strict_load=False) + + LOGGER.info("-" * 40) + LOGGER.info("[权重匹配核查报告]") + LOGGER.info(f"待加载参数总量: {len(param_dict)}") + LOGGER.info(f"网络未被初始化的参数数目: {len(param_not_load)}") + + # 若缺失参数过多 (通常由于配置失误或网络拓扑更改导致)，可依据具体环境抛出警告 + if len(param_not_load) > 0: + LOGGER.warning("部分网络参数未能在 Checkpoint 中寻址匹配，样本如下:") + for p in param_not_load[:5]: + LOGGER.warning(f" -> {p}") + LOGGER.info("-" * 40) + + return model + + def build_optimizer(self, model): + """基于 hyp.yaml 或默认配置构建学习率策略与优化器""" + steps_per_epoch = self.train_loader.get_dataset_size() + lr_list = get_lr(self.args, self.hyp, steps_per_epoch) + return build_optimizer(model, lr_list, self.hyp) + + def get_loss_fn(self): + """配置支持 Prototype 掩码评估的损失计算图""" + criterion = v8SegmentationLoss(self.model) + criterion.imgsz = self.args.imgsz + return criterion + + def get_validator(self): + """装载针对 Mask mAP 计算定制的实例分割验证器""" + from models.yolo.segment.val import SegmentationValidator + + args_dict = vars(self.args) + args_dict['task'] = 'segment' + args_dict['mode'] = 'val' + + names = self.data.get('names', {i: f'class_{i}' for i in range(self.data.get('nc', 80))}) + return SegmentationValidator(dataloader=self.val_loader, save_dir=self.save_dir, args=self.args, names=names) + + def label_loss_items(self, loss_items=None, prefix="train"): + keys = [f"{prefix}/{x}" for x in self.loss_names] + if loss_items is not None: + loss_items = [round(float(x), 5) for x in loss_items] + return dict(zip(keys, loss_items)) + return keys \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/val.py b/src/mindnlp/models/ultralytics/models/yolo/segment/val.py new file mode 100644 index 000000000..79b1cbf40 --- /dev/null +++ b/src/mindnlp/models/ultralytics/models/yolo/segment/val.py @@ -0,0 +1,205 @@ +import logging +import numpy as np +import cv2 +import mindspore as ms +from mindspore import Tensor, ops + +from engine.validator import BaseValidator +from utils.ops import non_max_suppression, process_mask, xywh2xyxy_np +from utils.metrics import SegmentMetrics + +LOGGER = logging.getLogger(__name__) + +class SegmentationValidator(BaseValidator): + """ + 实例分割任务评估验证器 + 重写 BaseValidator 的前处理及后处理节点，以计算 BBox mAP 与 Mask mAP 双重验证指标 + """ + def __init__(self, dataloader=None, save_dir=None, args=None, names=None): + super().__init__(dataloader, save_dir, args) + self.names = names + self.metrics = None + self.stats = {} + self.niou = 10 + self.iou_v = np.linspace(0.5, 0.95, self.niou) + + def get_dataloader(self, dataset_path, batch_size=16): + """激活具备 Mask 特征字段打包逻辑的数据加载器""" + return create_dataloader( + path=dataset_path, + imgsz=getattr(self.args, 'imgsz', 640), + batch_size=batch_size, + task='segment', + is_training=False, + num_workers=getattr(self.args, 'workers', 8) + ) + + def init_metrics(self, model): + """建立分工独立的 BBox 与 Mask 精度统计容器""" + if not self.names: + self.names = getattr(model, 'names', {i: f'class_{i}' for i in range(getattr(self.args, 'nc', 80))}) + self.metrics = SegmentMetrics(names=self.names) + + def preprocess(self, batch): + """ + 预处理阶段：确保图像输入浮点化域为 [0, 1] + 统一目标真值边框 (GT BBoxes) 的数据格式为物理像素下的绝对坐标 (xyxy) + """ + batch["image"] = ops.cast(batch["image"], ms.float32) / 255.0 + + _, _, h, w = batch["image"].shape + bboxes_np = batch["bboxes"].asnumpy() + + if len(bboxes_np) > 0: + bboxes_xyxy = xywh2xyxy_np(bboxes_np) + + # 若原始数据停留在归一化相对空间，则实施放大操作 + if bboxes_xyxy.max() <= 1.01: + bboxes_xyxy[:, [0, 2]] *= w + bboxes_xyxy[:, [1, 3]] *= h + + batch["bboxes"] = ms.Tensor(bboxes_xyxy, dtype=ms.float32) + + return batch + + def postprocess(self, preds): + """ + 后处理逻辑链：从网络提取原生预测簇，利用 NMS 算子进行空间重叠度消解 + 验证阶段要求捕获极大限度的召回率边界，故应用预设的超低置信度阈值 + """ + preds_box, proto = preds[0], preds[1] + + # 为了捕获完整的 Precision-Recall 曲线全貌，验证期的 conf_thres 被极小化约束 + iou_thres = self.hyp.get('iou', 0.6) + preds_nms = non_max_suppression( + preds_box, + conf_thres=0.001, + iou_thres=iou_thres, + max_det=300 + ) + return preds_nms, proto + + def update_metrics(self, preds, batch): + """计算图像批次内的真实阳性实例 (TP)，记录针对 Mask 与 Box 的双模评估指标""" + preds_nms, proto = preds + + batch_idx = batch["batch_idx"].view(-1).asnumpy() + all_gt_cls = batch["cls"].view(-1).asnumpy() + all_gt_bboxes = batch["bboxes"].asnumpy() + all_gt_masks = batch["masks"] + + imgsz = getattr(self.args, 'imgsz', 640) + + for i, pred in enumerate(preds_nms): + mask_gt = (batch_idx == i) + gt_cls = all_gt_cls[mask_gt] + gt_bboxes = all_gt_bboxes[mask_gt] + gt_masks = all_gt_masks[mask_gt] + + pred_np = pred.asnumpy() if len(pred) > 0 else np.zeros((0, 6)) + + if len(pred) == 0: + if len(gt_cls) > 0: + tp_empty = np.zeros((0, self.niou), dtype=bool) + self.metrics.update_stats( + tp_b=tp_empty, tp_m=tp_empty, conf=np.zeros(0), + pred_cls=np.zeros(0), target_cls=gt_cls + ) + continue + + # 约束预测框溢出原图边际 + pred_np[:, :4] = np.clip(pred_np[:, :4], 0, imgsz) + pred = ms.Tensor(pred_np, dtype=ms.float32) + + # 调用高度优化的底层算子，在浮点计算域还原实例级 Mask 蒙版 + masks_pred_bin = process_mask(proto[i], pred[:, 6:], pred[:, :4], (imgsz, imgsz), upsample=False) + + # 生成各个 IoU 阈值刻度对应的真阳性判断矩阵 + tp_b = self.calculate_tp(pred_np[:, :4], pred_np[:, 5], gt_bboxes, gt_cls, is_mask=False) + tp_m = self.calculate_tp(masks_pred_bin, pred_np[:, 5], gt_masks, gt_cls, is_mask=True) + + self.metrics.update_stats( + tp_b=tp_b, tp_m=tp_m, conf=pred_np[:, 4], + pred_cls=pred_np[:, 5], target_cls=gt_cls + ) + + def finalize_metrics(self): + """执行最终度量计算""" + if len(self.metrics.stats) == 0: + LOGGER.warning("[WARNING] 当前评估轮次无有效度量样本输入。") + self.stats = {"metrics/mAP50-95(B)": 0.0, "metrics/mAP50-95(M)": 0.0} + return + + try: + self.stats = self.metrics.process() + except Exception as e: + LOGGER.error(f"[ERROR] 指标聚合运算发生内部异常终止: {e}") + self.stats = {"metrics/mAP50-95(B)": 0.0, "metrics/mAP50-95(M)": 0.0} + + def get_stats(self): + """交还评估快照""" + stats = super().get_stats() + stats.update(self.stats) + + if self.metrics is not None: + stats['fitness'] = self.metrics.fitness + else: + stats['fitness'] = 0.0 + + self.results_dict = stats + return stats + + def calculate_tp(self, preds, pred_cls, gts, gt_cls, is_mask=False): + """ + 全任务域通用匹配校验核心逻辑：解析 IoU 并构建布尔记录矩阵 + 内置了预防因批尺寸特征不对齐引发的矩阵相乘 (MatMul) 异常保护 + """ + if len(gts) == 0: + return np.zeros((preds.shape[0], self.niou), dtype=bool) + + p_tensor = ms.Tensor(preds) if not isinstance(preds, ms.Tensor) else preds + g_tensor = ms.Tensor(gts) if not isinstance(gts, ms.Tensor) else gts + pc_tensor = ms.Tensor(pred_cls) if not isinstance(pred_cls, ms.Tensor) else pred_cls + gc_tensor = ms.Tensor(gt_cls) if not isinstance(gt_cls, ms.Tensor) else gt_cls + + if is_mask: + n, h, w = p_tensor.shape + m, gh, gw = g_tensor.shape + + # 若空间解析度与特征原型图失配，则强制将真实标签蒙版执行重采样收缩 + if (gh, gw) != (h, w): + g_resized = ops.ResizeNearestNeighbor((h, w))(g_tensor.expand_dims(1))[:, 0] + g = g_resized.view(m, -1).astype(ms.float32) + else: + g = g_tensor.view(m, -1).astype(ms.float32) + + p = p_tensor.view(n, -1).astype(ms.float32) + + inter = ops.matmul(p, g.T) + area_p = p.sum(1, keepdims=True) + area_g = g.sum(1, keepdims=True).T + iou = inter / (area_p + area_g - inter + 1e-7) + else: + g_xyxy = ms.Tensor(gts) if not isinstance(gts, ms.Tensor) else gts + + lt = ops.maximum(p_tensor[:, None, :2], g_xyxy[:, :2]) + rb = ops.minimum(p_tensor[:, None, 2:], g_xyxy[:, 2:]) + + wh = ops.maximum(rb - lt, 0.0) + inter = wh[:, :, 0] * wh[:, :, 1] + + area_p = (p_tensor[:, 2] - p_tensor[:, 0]) * (p_tensor[:, 3] - p_tensor[:, 1]) + area_g = (g_xyxy[:, 2] - g_xyxy[:, 0]) * (g_xyxy[:, 3] - g_xyxy[:, 1]) + iou = inter / (area_p[:, None] + area_g - inter + 1e-7) + + correct_class = (pc_tensor.view(-1, 1) == gc_tensor.view(1, -1)) + + tp = np.zeros((preds.shape[0], self.niou), dtype=bool) + iou_np = iou.asnumpy() + cc_np = correct_class.asnumpy() + + for i, threshold in enumerate(self.iou_v): + match_matrix = (iou_np >= threshold) & cc_np + tp[:, i] = match_matrix.any(axis=1) + + return tp \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/modules.py b/src/mindnlp/models/ultralytics/modules.py new file mode 100644 index 000000000..cce68b49b --- /dev/null +++ b/src/mindnlp/models/ultralytics/modules.py @@ -0,0 +1,469 @@ +import math +import numpy as np +import mindspore as ms +import mindspore.numpy as mnp +from mindspore import nn, ops, Tensor, Parameter + +from utils.ops import dist2bbox +from utils.tal import make_anchors + + +# 基础工具函数与算子 + +def autopad(k, p=None, d=1): + """自动计算填充 (Padding) 以保持输出特征图尺寸一致""" + if d > 1: + k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] + if p is None: + p = k // 2 if isinstance(k, int) else [x // 2 for x in k] + return p + +class Identity(nn.Cell): + """恒等映射层""" + def construct(self, x): + return x + +class Concat(nn.Cell): + """张量拼接层""" + def __init__(self, axis=1): + super(Concat, self).__init__() + self.axis = axis + def construct(self, x): + return ops.concat(x, self.axis) + +class Upsample(nn.Cell): + """ + 上采样模块 + """ + def __init__(self, size=None, scale_factor=None, mode='nearest'): + super().__init__() + self.size = size + self.scale_factor = scale_factor + self.mode = mode + + def construct(self, x): + if self.scale_factor is not None: + shape = x.shape + new_size = (int(shape[-2] * self.scale_factor), int(shape[-1] * self.scale_factor)) + return ops.interpolate(x, size=new_size, mode=self.mode) + + return ops.interpolate(x, size=self.size, mode=self.mode) + + +# 核心卷积模块 + +class ConvNormAct(nn.Cell): + """标准卷积块：包含 Conv2d + BatchNorm2d + 激活函数""" + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True, momentum=0.97, eps=1e-3, sync_bn=False): + super(ConvNormAct, self).__init__() + self.conv = nn.Conv2d( + c1, c2, k, s, pad_mode="pad", padding=autopad(k, p, d), group=g, dilation=d, has_bias=False + ) + self.bn = nn.BatchNorm2d(c2, momentum=momentum, eps=eps).to_float(ms.float32) + self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Cell) else Identity()) + + def construct(self, x): + return self.act(self.bn(self.conv(x))) + +class DWConv(ConvNormAct): + """深度可分离卷积 (Depthwise Convolution)""" + def __init__(self, c1, c2, k=1, s=1, d=1, act=True, sync_bn=False): + # 使用 math.gcd 确保 group 数等于输入输出的最大公约数 + super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act, sync_bn=sync_bn) + +class Bottleneck(nn.Cell): + """标准瓶颈块 (Bottleneck)""" + def __init__(self, c1, c2, shortcut=True, k=(3, 3), g=(1, 1), e=0.5, act=True, momentum=0.97, eps=1e-3, sync_bn=False): + super().__init__() + c_ = int(c2 * e) + # 根据官方设计，此处强制设置 group=1 + self.conv1 = ConvNormAct(c1, c_, k[0], 1, g=1, act=act, momentum=momentum, eps=eps, sync_bn=sync_bn) + self.conv2 = ConvNormAct(c_, c2, k[1], 1, g=1, act=act, momentum=momentum, eps=eps, sync_bn=sync_bn) + self.add = shortcut and c1 == c2 + + def construct(self, x): + return x + self.conv2(self.conv1(x)) if self.add else self.conv2(self.conv1(x)) + +class C2f(nn.Cell): + """C2f 模块 (CSP Bottleneck with 2 convolutions)""" + def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, momentum=0.97, eps=1e-3, sync_bn=False): + super().__init__() + self.c = int(c2 * e) + self.cv1 = ConvNormAct(c1, 2 * self.c, 1, 1, momentum=momentum, eps=eps, sync_bn=sync_bn) + self.cv2 = ConvNormAct((2 + n) * self.c, c2, 1, momentum=momentum, eps=eps, sync_bn=sync_bn) + self.m = nn.CellList([Bottleneck(self.c, self.c, shortcut, k=(3, 3), g=(1, g), e=1.0, + momentum=momentum, eps=eps, sync_bn=sync_bn) for _ in range(n)]) + def construct(self, x): + x = self.cv1(x) + y = list(ops.split(x, axis=1, split_size_or_sections=self.c)) + for m in self.m: + y.append(m(y[-1])) + return self.cv2(ops.concat(y, axis=1)) + +class C3k(nn.Cell): + """C3k 模块""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3, sync_bn=False): + super().__init__() + c_ = int(c2 * e) + self.cv1 = ConvNormAct(c1, c_, 1, 1, sync_bn=sync_bn) + self.cv2 = ConvNormAct(c1, c_, 1, 1, sync_bn=sync_bn) + self.cv3 = ConvNormAct(2 * c_, c2, 1, sync_bn=sync_bn) + self.m = nn.SequentialCell([Bottleneck(c_, c_, shortcut, k=(k, k), g=1, e=1.0, sync_bn=sync_bn) for _ in range(n)]) + + def construct(self, x): + return self.cv3(ops.concat((self.m(self.cv1(x)), self.cv2(x)), axis=1)) + +class C3k2(C2f): + """YOLO11 核心模块 C3k2 (基于 C2f 优化的 C3k 变体)""" + def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True, sync_bn=False): + super().__init__(c1, c2, n, shortcut, g, e, sync_bn=sync_bn) + self.m = nn.CellList([ + C3k(self.c, self.c, 2, shortcut, g, sync_bn=sync_bn) if c3k else + Bottleneck(self.c, self.c, shortcut, k=(3, 3), g=(1, g), sync_bn=sync_bn) for _ in range(n) + ]) + +class SPPF(nn.Cell): + """空间金字塔池化模块 (Spatial Pyramid Pooling - Fast)""" + def __init__(self, c1, c2, k=5, momentum=0.97, eps=1e-3): + super(SPPF, self).__init__() + c_ = c1 // 2 + self.cv1 = ConvNormAct(c1, c_, 1, 1, momentum=momentum, eps=eps) + self.cv2 = ConvNormAct(c_ * 4, c2, 1, 1, momentum=momentum, eps=eps) + self.m = nn.MaxPool2d(kernel_size=k, stride=1, pad_mode="same") + + def construct(self, x): + x = self.cv1(x) + y1 = self.m(x) + y2 = self.m(y1) + y3 = self.m(y2) + return self.cv2(ops.concat((x, y1, y2, y3), 1)) + + +# 注意力机制模块 + +class Attention(nn.Cell): + """多头自注意力模块 (Multi-Head Self Attention)""" + def __init__(self, dim, num_heads=8, attn_ratio=0.5): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.key_dim = int(self.head_dim * attn_ratio) + self.scale = self.key_dim ** -0.5 + nh_kd = self.key_dim * num_heads + h = dim + nh_kd * 2 + self.qkv = ConvNormAct(c1=dim, c2=h, k=1, act=False) + self.proj = ConvNormAct(c1=dim, c2=dim, k=1, act=False) + self.pe = ConvNormAct(c1=dim, c2=dim, k=3, s=1, g=dim, act=False) + + def construct(self, x): + B, C, H, W = x.shape + N = H * W + qkv = self.qkv(x) + q, k, v = qkv.view(B, self.num_heads, self.key_dim*2 + self.head_dim, N).split([self.key_dim, self.key_dim, self.head_dim], axis=2) + attn = (ops.transpose(q, (0, 1, 3, 2)) @ k) * self.scale + attn = ops.softmax(attn, axis=-1) + x = (v @ ops.transpose(attn, (0, 1, 3, 2))).view(B, C, H, W) + self.pe(v.reshape(B, C, H, W)) + return self.proj(x) + +class PSABlock(nn.Cell): + """PSA (Polarized Self-Attention) 模块""" + def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True, sync_bn=False): + super().__init__() + self.attn = Attention(c, attn_ratio=attn_ratio, num_heads=num_heads) + self.ffn = nn.SequentialCell([ + ConvNormAct(c, c * 2, 1, sync_bn=sync_bn), + ConvNormAct(c * 2, c, 1, act=False, sync_bn=sync_bn) + ]) + self.add = shortcut + def construct(self, x): + x = x + self.attn(x) if self.add else self.attn(x) + x = x + self.ffn(x) if self.add else self.ffn(x) + return x + +class C2PSA(nn.Cell): + """融合 PSA 注意力机制的 C2 模块""" + def __init__(self, c1, c2, n=1, e=0.5, sync_bn=False): + super().__init__() + assert c1 == c2 + self.c = int(c1 * e) + self.cv1 = ConvNormAct(c1, 2 * self.c, 1, 1, sync_bn=sync_bn) + self.cv2 = ConvNormAct(2 * self.c, c1, 1, sync_bn=sync_bn) + self.m = nn.SequentialCell([PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64, sync_bn=sync_bn) for _ in range(n)]) + def construct(self, x): + x = self.cv1(x) + a, b = ops.split(x, axis=1, split_size_or_sections=self.c) + return self.cv2(ops.concat((a, self.m(b)), 1)) + + +# 网络头部模块 (Heads) + +class Classify(nn.Cell): + """YOLO11 分类任务头部""" + def __init__(self, c1, c2, *args, **kwargs): + super().__init__() + c_ = 1280 # efficientnet_b0 默认输出尺寸 + + # 强制设置 k=1, s=1 确保拓扑结构一致 + self.conv = ConvNormAct(c1, c_, k=1, s=1) + self.pool = nn.AdaptiveAvgPool2d(1) + self.drop = nn.Dropout(p=0.0) + self.linear = nn.Dense(c_, c2) + + def construct(self, x): + if isinstance(x, (list, tuple)): + x = ops.concat(x, 1) + + x = self.conv(x) + x = self.pool(x) + x = ops.flatten(x, start_dim=1) + x = self.drop(x) + x = self.linear(x) + + if self.training: + return x + y = ops.softmax(x, axis=1) + return y, x + +class DFL(nn.Cell): + """分布焦点损失 (Distribution Focal Loss) 积分计算模块""" + def __init__(self, c1=16): + super().__init__() + self.conv = nn.Conv2d(c1, 1, 1, has_bias=False) + self.conv.weight.requires_grad = False # 积分序列权重应保持固定 + self.c1 = c1 + self.softmax = ops.Softmax(axis=1) + self.initialize_conv_weight() + + def initialize_conv_weight(self): + """初始化离散积分权重""" + self.conv.weight.set_data( + Tensor(np.arange(self.c1).reshape((1, self.c1, 1, 1)), dtype=ms.float32) + ) + + def construct(self, x): + s = x.shape + b, c, a = s[0], s[1], s[-1] + x = self.softmax(x.view(b, 4, self.c1, a).swapaxes(2, 1)) + x = self.conv(x) + return x.view(b, 4, a) + +class YOLO11DetectHead(nn.Cell): + """YOLO11 目标检测头部""" + def __init__(self, nc=80, reg_max=16, stride=(), ch=(), sync_bn=False): + super().__init__() + self.nc = nc + self.nl = len(ch) + self.reg_max = reg_max + self.no = nc + reg_max * 4 + + self.stride = Parameter(ms.Tensor(stride, ms.int32), requires_grad=False) + + c2 = max((16, ch[0] // 4, reg_max * 4)) + c3 = max(ch[0], min(nc, 100)) + + self.cv2 = nn.CellList([ + nn.SequentialCell([ + ConvNormAct(x, c2, 3), + ConvNormAct(c2, c2, 3), + nn.Conv2d(c2, 4 * reg_max, 1, has_bias=True) + ]) for x in ch + ]) + + self.cv3 = nn.CellList([ + nn.SequentialCell([ + nn.SequentialCell([DWConv(x, x, 3), ConvNormAct(x, c3, 1)]), + nn.SequentialCell([DWConv(c3, c3, 3), ConvNormAct(c3, c3, 1)]), + nn.Conv2d(c3, nc, 1, has_bias=True, pad_mode="pad", padding=0) + ]) for x in ch + ]) + self.dfl = DFL(reg_max) + + def construct(self, x): + """ + YOLO11 检测头前向传播 + 训练模式：返回 (拼接后的展平Tensor, 原始特征图List) + 推理模式：返回 (最终预测结果, 原始特征图List) + """ + # 1. 基础卷积处理 + res = [] + for i in range(self.nl): + # 保持 [B, 64+nc, H, W] 结构 + res.append(ops.concat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)) + + # 2. 统一展平逻辑 (训练和推理都需要这个展平后的 x_cat 来算 Loss) + y = [] + for xi in res: + bs, ch, h, w = xi.shape + y.append(xi.view(bs, ch, -1)) + + # 拼接所有尺度：[B, 144, 19200] + x_cat = ops.concat(y, 2) + + if self.training: + return x_cat, res + + # 3. 推理解码逻辑 (仅在 self.training=False 时执行) + box_dist, cls_logits = ops.split(x_cat, (self.reg_max * 4, self.nc), axis=1) + + # DFL 积分回归 + box_decoded = self.dfl(box_dist) # 形状 [B, 4, 19200] + + # 获取锚点和步长 + # 这里的 x 必须是原始特征图列表 + anchors, strides = self.make_anchors(x, self.stride, 0.5) + + # 解码坐标：直接传入原始 anchors，利用 ops.py 里的自适应逻辑 + dbox = dist2bbox(box_decoded, anchors, strides=strides, xywh=True, axis=1) + + # 最终拼接：[B, 4+nc, 19200] + final_pred = ops.concat((dbox, ops.sigmoid(cls_logits)), 1) + + return final_pred, res + + +class YOLO11Pose(YOLO11DetectHead): + """YOLO11 姿态估计头部""" + def __init__(self, nc=1, kpt_shape=(17, 3), reg_max=16, stride=(), ch=(), sync_bn=False): + nc = 1 + super().__init__(nc, reg_max, stride, ch, sync_bn) + self.kpt_shape = kpt_shape + self.nkpt = kpt_shape[0] * kpt_shape[1] + c4 = max(ch[0] // 4, self.nkpt) + + self.cv4 = nn.CellList([ + nn.SequentialCell([ + ConvNormAct(x, c4, 3), + ConvNormAct(c4, c4, 3), + nn.Conv2d(c4, self.nkpt, 1, has_bias=True) + ]) for x in ch + ]) + + def kpts_decode(self, kpt_flat, anc, strides): + """解析并还原关键点坐标及其可见度""" + B = kpt_flat.shape[0] + + # 统一提升至 float32 精度，防止混合精度计算时发生溢出 + kpt_flat = kpt_flat.astype(ms.float32) + anc = anc.astype(ms.float32) + strides = strides.astype(ms.float32) + + kpt = kpt_flat.transpose(0, 2, 1).view(B, -1, self.kpt_shape[0], self.kpt_shape[1]) + + anc_safe = anc.view(1, -1, 1, 2) + strides_safe = strides.view(1, -1, 1, 1) + + xy = kpt[..., :2] + xy_decoded = (xy * 2.0 + (anc_safe - 0.5)) * strides_safe + vis = ops.sigmoid(kpt[..., 2:3]) + + kpt_decoded = ops.concat((xy_decoded, vis), axis=-1).view(B, -1, self.nkpt) + return kpt_decoded.transpose(0, 2, 1) + + def construct(self, x): + bs = x[0].shape[0] + box_outs, cls_outs, kpt_outs = [], [], [] + for i in range(self.nl): + box_outs.append(self.cv2[i](x[i])) + cls_outs.append(self.cv3[i](x[i])) + kpt_outs.append(self.cv4[i](x[i])) + + if self.training: + out = tuple(ops.concat((box_outs[i], cls_outs[i]), 1) for i in range(self.nl)) + return out, kpt_outs + + box_flat = ops.concat([b.view(bs, self.reg_max * 4, -1) for b in box_outs], axis=2) + cls_flat = ops.concat([c.view(bs, self.nc, -1) for c in cls_outs], axis=2) + kpt_flat = ops.concat([k.view(bs, self.nkpt, -1) for k in kpt_outs], axis=2) + + # 统一提升精度以保证解码阶段数值稳定性 + box_flat = box_flat.astype(ms.float32) + cls_flat = cls_flat.astype(ms.float32) + kpt_flat = kpt_flat.astype(ms.float32) + + cls_prob = ops.sigmoid(cls_flat) + + self.anchors, self.strides = make_anchors(x, self.stride, 0.5) + anc = self.anchors.astype(ms.float32).view(1, -1, 2).transpose(0, 2, 1) + strd = self.strides.astype(ms.float32).view(1, -1, 1).transpose(0, 2, 1) + + box_dist = self.dfl(box_flat) + lt = box_dist[:, :2, :] + rb = box_dist[:, 2:, :] + c_xy = anc + (rb - lt) / 2.0 + wh = rb + lt + dbox = ops.concat((c_xy, wh), axis=1) * strd + + pred_kpt = self.kpts_decode(kpt_flat, self.anchors, self.strides) + + # 按照 (Box, Cls, Pose) 顺序拼接 + y = ops.concat((dbox, cls_prob, pred_kpt), axis=1) + return y.transpose(0, 2, 1) + +class ProtoCell(nn.Cell): + """实例分割任务中的原型 (Prototype) 生成模块""" + def __init__(self, c1, c_=256, c2=32, sync_bn=False): + super().__init__() + self.cv1 = ConvNormAct(c1, c_, 3, sync_bn=sync_bn, momentum=0.9, eps=1e-3) + self.upsample = nn.Conv2dTranspose(c_, c_, 2, 2, has_bias=True) + self.cv2 = ConvNormAct(c_, c_, 3, sync_bn=sync_bn, momentum=0.9, eps=1e-3) + self.cv3 = ConvNormAct(c_, c2, k=1, sync_bn=sync_bn, act=False) + + def construct(self, x): + x = self.cv1(x) + x = self.upsample(x) + x = self.cv2(x) + x = self.cv3(x) + + # 注意：此处限制数值范围旨在避免混合精度下的异常溢出 + # 规范做法建议排查底层 float16 乘加运算，或依赖全局 FP32 转换 + x = ops.clip_by_value(x, -10.0, 10.0) + return x + +class YOLO11Segment(YOLO11DetectHead): + """YOLO11 实例分割头部""" + def __init__(self, nc=80, reg_max=16, nm=32, npr=256, stride=(), ch=(), sync_bn=False): + super().__init__(nc=nc, reg_max=reg_max, stride=stride, ch=ch, sync_bn=sync_bn) + self.nm = nm + c4 = max(ch[0] // 4, nm) + _npr = max(c4, min(ch[0], npr)) + self.proto = ProtoCell(ch[0], _npr, nm, sync_bn=sync_bn) + + self.cv4 = nn.CellList([ + nn.SequentialCell([ + ConvNormAct(x, c4, 3, sync_bn=sync_bn), + ConvNormAct(c4, c4, 3, sync_bn=sync_bn), + nn.Conv2d(c4, nm, 1, has_bias=True) + ]) for x in ch + ]) + + def construct(self, x): + p_out = self.proto(x[0]) + + bs = x[0].shape[0] + out, mc = (), () + for i in range(self.nl): + xi = x[i] + cv2_feat = self.cv2[i](xi) + cv3_feat = self.cv3[i](xi) + + out += (ops.concat((cv2_feat, cv3_feat), 1),) + mc += (self.cv4[i](xi).view(bs, self.nm, -1),) + + if self.training: + return out, ops.concat(mc, 2), p_out + + self.anchors, self.strides = self.make_anchors(out, self.stride, 0.5) + x_all = ops.concat([xi.view(bs, self.no, -1) for xi in out], 2) + + box, cls = ops.split(x_all, (self.reg_max * 4, self.nc), 1) + + box_decoded = self.dfl(box) + box_decoded = ops.relu(box_decoded) + dbox = dist2bbox(box_decoded.transpose(0, 2, 1), + ops.expand_dims(self.anchors, 0), + self.strides, xywh=True) + + mc_all = ops.concat(mc, 2) + final_pred = ops.concat((dbox, ops.sigmoid(cls), mc_all), 1) + + return final_pred, p_out \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/readme.md b/src/mindnlp/models/ultralytics/readme.md new file mode 100644 index 000000000..ec79f30b1 --- /dev/null +++ b/src/mindnlp/models/ultralytics/readme.md @@ -0,0 +1,79 @@ +# mindnlp兼容ultralutics库项目运行指南 + +本项目基于 MindSpore 框架实现了 YOLO11 的四大核心任务：图像分类、目标检测、实例分割和姿态估计。支持从头训练、加载预训练权重微调、模型验证与推理。 + +1. 数据集准备 + +请参考 `./ultralytics/cfg/datasets` 目录下各个数据集的配置文件（`.yaml`）获取下载路径。 +下载完成后，请将数据集解压并放置到以下目录： +`./ultralytics/datasets/` + +2. 预训练权重转换 + +我们提供了将 PyTorch 的 `.pt` 权重转化为 MindSpore 的 `.ckpt` 权重的转换脚本。请在项目根目录下运行以下命令： + +```bash +# 转换分类任务权重 +python tools/convert.py --task classify +# 转换检测任务权重 +python tools/convert.py --task detect +# 转换分割任务权重 +python tools/convert.py --task segment +# 转换姿态估计任务权重 +python tools/convert.py --task pose + +3. 任务运行指令 +以下命令均需在项目根目录下执行。 + +（1）图像分类任务 (Classify) + +# 加载权重微调 (Fine-tune) +python examples/yolo/classify/run_train.py --weights ./yolo11n-cls.ckpt --save_dir ./runs/cls/train_finetune + +# 从头开始训练 (Train from scratch) +python examples/yolo/classify/run_train.py --save_dir ./runs/cls/train_scratch + +# 模型验证 (Validate) +python examples/yolo/classify/run_val.py + +# 模型推理 (Inference) +python examples/yolo/classify/inference.py +（2）目标检测任务 (Detect) + +# 加载权重微调 (Fine-tune) +python examples/yolo/detect/run_train.py --weights ./yolo11n.ckpt --save_dir ./runs/detect/train_finetune + +# 从头开始训练 (Train from scratch) +python examples/yolo/detect/run_train.py --save_dir ./runs/detect/train_scratch + +# 模型验证 (Validate) +python examples/yolo/detect/run_val.py + +# 模型推理 (Inference) +python examples/yolo/detect/inference.py +（3）实例分割任务 (Segment) + +# 加载权重微调 (Fine-tune) +python examples/yolo/segment/run_train.py --weights ./yolo11n-seg.ckpt --save_dir ./runs/segment/train_finetune + +# 从头开始训练 (Train from scratch) +python examples/yolo/segment/run_train.py --save_dir ./runs/segment/train_scratch + +# 模型验证 (Validate) +python examples/yolo/segment/run_val.py + +# 模型推理 (Inference) +python examples/yolo/segment/inference.py +（4）姿态估计任务 (Pose) + +# 加载权重微调 (Fine-tune) +python examples/yolo/pose/run_train.py --weights ./yolo11n-pose.ckpt --save_dir ./runs/pose/train_finetune + +# 从头开始训练 (Train from scratch) +python examples/yolo/pose/run_train.py --save_dir ./runs/pose/train_scratch + +# 模型验证 (Validate) +python examples/yolo/pose/run_val.py + +# 模型推理 (Inference) +python examples/yolo/pose/inference.py \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/tools/convert.py b/src/mindnlp/models/ultralytics/tools/convert.py new file mode 100644 index 000000000..471ac13c5 --- /dev/null +++ b/src/mindnlp/models/ultralytics/tools/convert.py @@ -0,0 +1,158 @@ +import os +import argparse +import numpy as np +import mindspore as ms +from mindspore import Tensor +from ultralytics import YOLO + +import modeling_yolo +from configuration_yolo import YOLOConfig + +def translate_ms_to_pt(ms_name, task): + """ + 参数映射路由表：将 MindSpore 架构下的参数名称转换为 PyTorch 架构对应的参数名称 + """ + # 基础规范化转换：对齐 BatchNorm 以及权重、偏置的命名差异 + pt_name = ms_name.replace(".moving_mean", ".running_mean") \ + .replace(".moving_variance", ".running_var") \ + .replace(".gamma", ".weight") \ + .replace(".beta", ".bias") + + # 结构命名转换：对齐 C2f/Bottleneck 等内部核心模块的命名规范 + pt_name = pt_name.replace(".conv1.", ".cv1.") + pt_name = pt_name.replace(".conv2.", ".cv2.") + pt_name = pt_name.replace(".conv3.", ".cv3.") + + return pt_name + +def universal_convert(task="segment", scale="n"): + """ + 通用权重转换流水线。 + 支持将 Ultralytics 官方 PyTorch (.pt) 权重转换为符合 MindNLP 规范的 MindSpore (.ckpt) 权重 + """ + # 1. 初始化任务及配置映射字典 + model_map = { + "classify": modeling_yolo.YOLO11ForClassification, + "detect": modeling_yolo.YOLO11ForObjectDetection, + "segment": modeling_yolo.YOLO11ForSegmentation, + "pose": modeling_yolo.YOLO11ForPose + } + pt_name_map = { + "classify": f"yolo11{scale}-cls.pt", + "detect": f"yolo11{scale}.pt", + "segment": f"yolo11{scale}-seg.pt", + "pose": f"yolo11{scale}-pose.pt" + } + yaml_map = { + "classify": "cfg/models/11/yolo11-cls.yaml", + "detect": "cfg/models/11/yolo11.yaml", + "segment": "cfg/models/11/yolo11-seg.yaml", + "pose": "cfg/models/11/yolo11-pose.yaml" + } + + nc_map = {"classify": 1000, "detect": 80, "segment": 80, "pose": 1} + + print(f"[INFO] 启动 YOLO11-{task.upper()} 权重转换流程...") + + current_yaml = yaml_map[task] + + # 2. 动态构建网络配置与 MindSpore 模型实例 + if task == "pose": + cfg = YOLOConfig(yaml_path=current_yaml, scale=scale, nc=nc_map[task], kpt_shape=[17, 3]) + else: + cfg = YOLOConfig(yaml_path=current_yaml, scale=scale, nc=nc_map[task]) + + ms_model = model_map[task](cfg) + ms_model.set_train(False) + + # 3. 加载 PyTorch 原生权重 + print(f"[INFO] 正在解析 PyTorch 权重文件: {pt_name_map[task]}") + pt_yolo = YOLO(pt_name_map[task]) + pt_dict = pt_yolo.model.state_dict() + + print("[DEBUG] 官方 PyTorch 权重键名抽样 (末尾 5 项):") + pt_keys = list(pt_dict.keys()) + for k in pt_keys[-5:]: + print(f" {k:50} | 尺寸: {list(pt_dict[k].shape)}") + print("-" * 80) + + new_ms_ckpt = [] + matched_count = 0 + ms_params = list(ms_model.parameters_and_names()) + + print("[INFO] 开始执行参数映射与维度校验...") + + # 4. 执行转换主循环 + for ms_name, ms_param in ms_params: + ms_shape = tuple(ms_param.shape) + ms_size = ms_param.size + + # 获取期望映射的 PyTorch 键名 + expected_pt_name = translate_ms_to_pt(ms_name, task) + + if expected_pt_name in pt_dict: + pt_v = pt_dict[expected_pt_name] + + # 维度一致性校验 + if pt_v.numel() == ms_size: + val_np = pt_v.cpu().numpy().reshape(ms_shape) + + # 安全策略：限制 BatchNorm 历史方差下限，防止半精度推理时出现除零溢出 + if "moving_variance" in ms_name: + val_np = np.maximum(val_np, 1e-5) + + ms_param.set_data(Tensor(val_np, ms.float32)) + new_ms_ckpt.append({'name': ms_name, 'data': ms_param.data}) + + print(f"[对齐成功] {ms_name:<55} <- {expected_pt_name}") + matched_count += 1 + + # 内存优化：匹配成功后从字典中移除该键，以便最后审计遗留项 + del pt_dict[expected_pt_name] + else: + print(f"[形状冲突] 参数 {ms_name} 期望尺寸 {ms_shape}，实际载入尺寸 {tuple(pt_v.shape)}") + else: + # 处理常量和不可训练张量 (如 DFL 积分权重、步长锚点) + # 由于网络初始化时已经通过 Config 构建了正确的值，此处直接保留并存入 Checkpoint 即可 + if "stride" in ms_name or "dfl.conv.weight" in ms_name: + new_ms_ckpt.append({'name': ms_name, 'data': ms_param.data}) + print(f"[保留原值] {ms_name:<55} (框架内置固定张量)") + matched_count += 1 + else: + # 若非常量且未能在 PT 字典中找到映射，留交审计模块处理 + pass + + # 5. 输出转换审计报告 + print("-" * 80) + print("[INFO] 权重转换审计清单") + matched_names = [x['name'] for x in new_ms_ckpt] + failed_names = [n for n, _ in ms_params if n not in matched_names] + + if not failed_names: + print(" 校验通过：所有模型参数均已成功映射。") + else: + print(f" 校验警告：存在 {len(failed_names)} 个未匹配参数，请检查拓扑结构：") + for fn in failed_names: + print(f" - {fn}") + + # 6. 持久化存储 + base_name = os.path.splitext(pt_name_map[task])[0] + save_ckpt_path = f"{base_name}.ckpt" + + ms.save_checkpoint(new_ms_ckpt, save_ckpt_path) + print(f"[INFO] 转换结束。参数对齐率: {matched_count}/{len(ms_params)}。已序列化至: {save_ckpt_path}") + + return save_ckpt_path + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="YOLO11 参数转换工具 (PyTorch to MindSpore Checkpoint)") + parser.add_argument("--task", "-t", type=str, default="pose", + choices=["classify", "detect", "segment", "pose"], + help="目标模型的基础任务类型") + parser.add_argument("--scale", "-s", type=str, default="n", + choices=["n", "s", "m", "l", "x"], + help="指定模型的规模缩放因子") + args = parser.parse_args() + + universal_convert(args.task, args.scale) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/__init__.py b/src/mindnlp/models/ultralytics/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/models/ultralytics/utils/ema.py b/src/mindnlp/models/ultralytics/utils/ema.py new file mode 100644 index 000000000..df8da3d2e --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/ema.py @@ -0,0 +1,46 @@ +import math +import copy +import mindspore as ms +from mindspore import ops + +class ModelEMA: + """ + MindSpore版 ModelEMA + """ + def __init__(self, model, decay=0.9999, tau=2000, updates=0): + # 深度拷贝模型实例 + self.ema_model = copy.deepcopy(model) + self.ema_model.set_train(False) + self.updates = updates + + # 动态衰减率策略 + self.decay = lambda x: decay * (1 - math.exp(-x / tau)) + + # 提取参数列表 (转为字典以防顺序错乱) + self.ema_params = {p.name: p for p in self.ema_model.get_parameters()} + self.model_params = {p.name: p for p in model.get_parameters()} + + # 锁定梯度 + for param in self.ema_params.values(): + param.requires_grad = False + + def update(self, model): + """ 执行单步 EMA 更新 """ + self.updates += 1 + d = self.decay(self.updates) + + # 提取当前最新模型参数 + current_model_params = {p.name: p for p in model.get_parameters()} + + for name, ema_p in self.ema_params.items(): + model_p = current_model_params.get(name) + if model_p is None: + continue + + if model_p.dtype in [ms.float16, ms.float32]: + if "moving_mean" in name or "moving_variance" in name: + ema_p.set_data(model_p.value()) + else: + # 常规权重执行 EMA 平滑 + new_val = d * ema_p.value() + (1.0 - d) * model_p.value() + ema_p.set_data(new_val) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/loss.py b/src/mindnlp/models/ultralytics/utils/loss.py new file mode 100644 index 000000000..ce80fa9eb --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/loss.py @@ -0,0 +1,535 @@ +import numpy as np +import mindspore as ms +import mindspore.ops as ops +from mindspore import nn, Tensor + +from utils.tal import TaskAlignedAssigner, make_anchors +from utils.ops import bbox_iou, xywh2xyxy, dist2bbox, bbox2dist + +# COCO 数据集 17 个关键点的标准差（用于计算 OKS 相似度） +OKS_SIGMA = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0 + + +# 分类任务损失函数 +class YOLOClassificationLoss(nn.Cell): + """YOLO11 图像分类任务损失函数 (基于交叉熵)""" + def __init__(self): + super(YOLOClassificationLoss, self).__init__() + self.loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') + + def construct(self, preds, targets): + if isinstance(targets, dict): + label_tensor = targets.get('label', targets.get('cls')) + else: + label_tensor = targets + + return self.loss_fn(preds, label_tensor) + +class v8ClassificationLoss(nn.Cell): + """YOLO 多标签分类逻辑损失函数 (基于 BCE)""" + def __init__(self): + super().__init__() + self.loss_fn = nn.BCEWithLogitsLoss(reduction='mean') + + def construct(self, preds, targets): + num_classes = preds.shape[-1] + dtype = preds.dtype + on_value = Tensor(1.0, dtype) + off_value = Tensor(0.0, dtype) + + # 修复：展平 targets 防止多出维度 + targets_flat = targets.view(-1) + targets_one_hot = ops.one_hot(targets_flat, num_classes, on_value, off_value) + + return self.loss_fn(preds, targets_one_hot) + +# 目标检测基础损失组件 + +class DFLoss(nn.Cell): + """分布焦点损失 (Distribution Focal Loss)""" + def __init__(self, reg_max=16): + super().__init__() + self.reg_max = reg_max + + def construct(self, pred_dist, target): + # 将输入分布重塑为交叉熵格式 [N * 4, reg_max] + pred_dist = pred_dist.view(-1, self.reg_max) + + # 拆分目标边界偏移量以计算积分权重 + target_left = ops.cast(target, ms.int32) + target_right = target_left + 1 + + weight_left = target_right.astype(pred_dist.dtype) - target + weight_right = target - target_left.astype(pred_dist.dtype) + + tl = target_left.view(-1) + tr = target_right.view(-1) + + loss_left = ops.cross_entropy(pred_dist, tl, reduction="none").view(target.shape) * weight_left + loss_right = ops.cross_entropy(pred_dist, tr, reduction="none").view(target.shape) * weight_right + + return (loss_left + loss_right).mean() + + +class BboxLoss(nn.Cell): + """边界框回归损失 (包含 CIoU/GIoU 损失与 DFL 损失)""" + def __init__(self, reg_max=16): + super().__init__() + self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None + + def construct(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask): + fg_indices = ops.nonzero(fg_mask) + loss_iou = ops.zeros((), ms.float32) + loss_dfl = ops.zeros((), ms.float32) + + if fg_indices.shape[0] > 0: + p_boxes = ops.gather_nd(pred_bboxes, fg_indices).astype(ms.float32) + t_boxes = ops.gather_nd(target_bboxes, fg_indices).astype(ms.float32) + weight = ops.gather_nd(target_scores.sum(-1), fg_indices).expand_dims(-1).astype(ms.float32) + + # 计算 IoU 损失并限制范围以确保数值稳定性 + iou = bbox_iou(p_boxes, t_boxes, xywh=False, CIoU=False).clip(0, 1) + loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum + + # 计算 DFL 损失 + if self.dfl_loss: + selected_p_dist = ops.gather_nd(pred_dist, fg_indices) + target_ltrb = bbox2dist(anchor_points, target_bboxes, self.dfl_loss.reg_max - 1) + selected_t_ltrb = ops.gather_nd(target_ltrb, fg_indices) + loss_dfl = self.dfl_loss(selected_p_dist, selected_t_ltrb) + + return loss_iou, loss_dfl + + +# YOLO11 主干任务损失类 + +class v8DetectionLoss(nn.Cell): + """YOLO11 目标检测任务损失函数计算类""" + def __init__(self, model): + super().__init__() + self.bce = nn.BCEWithLogitsLoss(reduction="none") + self.nc = model.nc + self.imgsz = getattr(model, 'imgsz', 640) + self.reg_max = model.reg_max + self.bbox_loss = BboxLoss(self.reg_max) + self.stride = model.stride + + # 动态分配器 + self.assigner = TaskAlignedAssigner(topk=10, num_classes=self.nc, alpha=0.5, beta=6.0) + self.proj = ops.arange(self.reg_max).astype(ms.float32) + + # 将默认超参数转化为类属性，避免硬编码 + self.hyp_box = getattr(model, 'hyp_box', 7.5) + self.hyp_cls = getattr(model, 'hyp_cls', 0.5) + self.hyp_dfl = getattr(model, 'hyp_dfl', 1.5) + + def decode_predictions(self, feats): + """ + 解析特征图，分离分类得分与边界框分布 + + Args: + feats (Tensor): [Batch, 4*reg_max + nc, Anchors] 形状的展平张量 + + Returns: + pred_distri (Tensor): 边界框分布信息 + pred_scores (Tensor): 分类得分 + """ + # 如果传入的是未拼接的原始特征图列表，执行自动拼接 (兼容性逻辑) + if isinstance(feats, (list, tuple)): + x_list = [] + for xi in feats: + b, c, h, w = xi.shape + x_list.append(xi.view(b, c, -1)) + x = ops.concat(x_list, axis=2) + else: + x = feats + + # 按通道分割分类与回归信息 + pred_distri, pred_scores = ops.split(x, (self.reg_max * 4, self.nc), axis=1) + + # 交换轴以适配后续计算：[B, C, N] -> [B, N, C] + return pred_distri.swapaxes(1, 2), pred_scores.swapaxes(1, 2) + + def bbox_decode(self, anchor_points, pred_dist): + """ + 通过积分机制将预测分布解码为物理边界框坐标 + + Args: + anchor_points (Tensor): 锚点中心坐标 + pred_dist (Tensor): 预测的 ltrb 分布 + """ + if self.reg_max > 1: + batch, anchors, channels = pred_dist.shape + + pred_dist = pred_dist.view(batch, anchors, 4, self.reg_max) + pred_dist = ops.softmax(pred_dist, axis=-1) + pred_dist = ops.matmul(pred_dist, self.proj) + + return dist2bbox(pred_dist, anchor_points[:pred_dist.shape[1]], xywh=False) + + def preprocess_targets(self, batch, feats): + """ + 将散装标签 (Collapsed Labels) 重组为标准 Batch 格式 + 解决 [Batch_Idx, Class, x, y, w, h] 格式的维度对齐问题 + """ + input_tensor = feats[0] if isinstance(feats, (list, tuple)) else feats + bs = input_tensor.shape[0] + + bboxes = batch["bboxes"] + cls = batch["cls"] + batch_idx = batch["batch_idx"].view(-1) + + if not isinstance(bboxes, ms.Tensor): + bboxes = ms.Tensor(bboxes, ms.float32) + if not isinstance(cls, ms.Tensor): + cls = ms.Tensor(cls, ms.float32) + + # 统计每张图的目标数并确定本批次最大目标量，用于 Tensor 填充 + counts = [] + for i in range(bs): + counts.append(int(ops.sum(batch_idx == i).asnumpy())) + max_obj = max(max(counts), 1) + + # 构建标准的 [Batch, Max_Obj, 5] 容器 + targets = ops.zeros((bs, max_obj, 5), bboxes.dtype) + + for i in range(bs): + mask = (batch_idx == i) + num = counts[i] + if num > 0: + actual_num = min(num, max_obj) + # 填充 Class ID 与 Bbox 坐标 + targets[i, :actual_num, 0] = cls[mask][:actual_num].view(-1) + targets[i, :actual_num, 1:5] = bboxes[mask][:actual_num] + + # 坐标映射：将归一化坐标转换为像素坐标 + if hasattr(self, 'imgsz') and self.imgsz: + targets[..., 1:5] *= self.imgsz + + return targets + + def construct(self, preds, batch): + """ + Loss 计算图主入口 + """ + # 解包模型输出：x_cat 为展平预测值，feats_list 为多尺度特征列表 + if isinstance(preds, (tuple, list)): + x_cat = preds[0] + feats_list = preds[1] + else: + x_cat = preds + feats_list = [preds] + + # 解析预测分布与分类得分 + pred_distri, pred_scores = self.decode_predictions(x_cat) + + # 生成锚点中心与步长向量 (基于特征图原始分辨率) + anchor_points, stride_tensor = make_anchors(feats_list, self.stride, 0.5) + + # 标签重组与拆分 + targets = self.preprocess_targets(batch, feats_list) + gt_labels, gt_bboxes = targets.split((1, 4), 2) # [B, N, 1], [B, N, 4] + mask_gt = gt_bboxes.sum(2, keepdim=True) > 0 # 有效目标掩码 + + # 坐标解码：生成解码后的预测框并转换真实框格式 + pred_bboxes = self.bbox_decode(anchor_points, pred_distri) + current_gt_bboxes = xywh2xyxy(gt_bboxes) + + # 正负样本分配 (Assigner) + # 注意：此处输入均已还原至像素空间 + target_bboxes, target_scores, fg_mask = self.assigner( + pred_scores.sigmoid(), + (pred_bboxes * stride_tensor), + anchor_points * stride_tensor, + gt_labels, + current_gt_bboxes, + mask_gt + )[1:4] + + # 损失计算逻辑 + target_scores_sum = ops.maximum(target_scores.sum(), 1.0) + + # 分类损失 (BCE) + loss_cls = self.bce(pred_scores, target_scores).sum() / target_scores_sum + + # 回归损失 (IoU + DFL) + loss_box, loss_dfl = ops.zeros((), ms.float32), ops.zeros((), ms.float32) + fg_mask_bool = ops.cast(fg_mask, ms.bool_) + + if fg_mask_bool.any(): + loss_box, loss_dfl = self.bbox_loss( + pred_distri, pred_bboxes, anchor_points, + target_bboxes / stride_tensor, target_scores, target_scores_sum, fg_mask_bool + ) + + # 记录 Loss Items 用于日志展示 (stop_gradient 保证记录不干扰反传) + self.loss_items = ops.stop_gradient(ops.stack([ + loss_box * self.hyp_box, + loss_cls * self.hyp_cls, + loss_dfl * self.hyp_dfl + ])) + + return (loss_box * self.hyp_box), (loss_cls * self.hyp_cls), (loss_dfl * self.hyp_dfl) + + +class v8SegmentationLoss(v8DetectionLoss): + """YOLO11 实例分割任务损失函数计算类""" + def __init__(self, model): + super().__init__(model) + self.overlap = True + self.hyp_seg = getattr(model, 'hyp_seg', 2.5) + + def preprocess_masks(self, batch, mask_h, mask_w, bs): + """将扁平化的分割掩码打包为批量张量""" + masks = batch["masks"].astype(ms.float32) + batch_idx = batch["batch_idx"].view(-1) + + if masks.shape[0] == 0: + return ops.zeros((bs, 1, mask_h, mask_w), ms.float32) + + if masks.shape[-2:] != (mask_h, mask_w): + masks = ops.ResizeNearestNeighbor((mask_h, mask_w))(masks.expand_dims(1))[:, 0] + + counts = [int((batch_idx == i).sum()) for i in range(bs)] + max_obj = max(max(counts) if counts else 0, 1) + + target_masks = ops.zeros((bs, max_obj, mask_h, mask_w), ms.float32) + for i in range(bs): + mask_flag = (batch_idx == i) + num = counts[i] + if num > 0: + actual_num = min(num, max_obj) + target_masks[i, :actual_num] = masks[mask_flag][:actual_num] + + return target_masks + + def calculate_segmentation_loss(self, fg_mask, gt_masks, target_gt_idx, target_bboxes, proto, pred_masks): + """计算掩码交叉熵损失""" + bs, nm, ph, pw = proto.shape + loss_mask = ops.zeros((), ms.float32) + + for i in range(bs): + mask_indices = ops.nonzero(fg_mask[i]).view(-1) + if mask_indices.shape[0] == 0: + continue + + coeffs = pred_masks[i][mask_indices] + p = proto[i].view(nm, -1) + pred_mask_full = ops.matmul(coeffs, p).view(-1, ph, pw) + + current_target_idx = target_gt_idx[i].view(-1) + obj_idx = ops.cast(current_target_idx[mask_indices], ms.int32) + target_mask = ops.cast(gt_masks[i][obj_idx] > 0, ms.float32) + + t_box = target_bboxes[i][mask_indices] + ratio_w = pw / getattr(self, 'imgsz', 640) + ratio_h = ph / getattr(self, 'imgsz', 640) + + x1 = (t_box[:, 0] * ratio_w).view(-1, 1, 1) + y1 = (t_box[:, 1] * ratio_h).view(-1, 1, 1) + x2 = (t_box[:, 2] * ratio_w).view(-1, 1, 1) + y2 = (t_box[:, 3] * ratio_h).view(-1, 1, 1) + + r_x = ops.arange(pw).view(1, 1, pw).astype(ms.float32) + r_y = ops.arange(ph).view(1, ph, 1).astype(ms.float32) + + crop_mask = ops.cast((r_x >= x1) & (r_x < x2) & (r_y >= y1) & (r_y < y2), ms.float32) + target_mask = target_mask * crop_mask + + loss_m = ops.binary_cross_entropy_with_logits(pred_mask_full, target_mask, reduction="none") + loss_mask += loss_m.mean(axis=(1, 2)).sum() + + return loss_mask + + def construct(self, preds, batch): + feats, pred_masks, proto = preds + bs, nm, mask_h, mask_w = proto.shape + + pred_distri, pred_scores = self.decode_predictions(feats) + pred_masks = pred_masks.transpose(0, 2, 1) + anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5) + + targets = self.preprocess_targets(batch, feats) + gt_labels, gt_bboxes = targets.split((1, 4), 2) + + if gt_bboxes.max() <= 1.0: + gt_bboxes = gt_bboxes * getattr(self, 'imgsz', 640) + mask_gt = gt_bboxes.sum(2, keepdim=True) > 0 + + pred_bboxes = self.bbox_decode(anchor_points, pred_distri) + current_gt_bboxes = xywh2xyxy(gt_bboxes) + + pred_scores_sg = ops.stop_gradient(pred_scores.sigmoid()) + pred_bboxes_sg = ops.stop_gradient(pred_bboxes * stride_tensor) + + _, target_bboxes, target_scores, fg_mask, target_gt_idx = self.assigner( + pred_scores_sg, pred_bboxes_sg, anchor_points * stride_tensor, + gt_labels, current_gt_bboxes, mask_gt + ) + + target_scores_sum = ops.maximum(target_scores.sum(), 1.0) + loss_cls = self.bce(pred_scores, target_scores).sum() / target_scores_sum + + loss_box, loss_dfl, loss_mask = ops.zeros((), ms.float32), ops.zeros((), ms.float32), ops.zeros((), ms.float32) + + if fg_mask.any(): + loss_box, loss_dfl = self.bbox_loss( + pred_distri, pred_bboxes, anchor_points, + target_bboxes / stride_tensor, target_scores, target_scores_sum, fg_mask + ) + gt_masks = self.preprocess_masks(batch, mask_h, mask_w, bs) + loss_mask = self.calculate_segmentation_loss( + fg_mask, gt_masks, target_gt_idx, target_bboxes, proto, pred_masks + ) + loss_mask = loss_mask / target_scores_sum + + return (loss_box * self.hyp_box), (loss_cls * self.hyp_cls), (loss_dfl * self.hyp_dfl), (loss_mask * self.hyp_seg) + + +class KeypointLoss(nn.Cell): + """关键点估计损失计算 (基于目标关键点相似度 OKS)""" + def __init__(self, sigmas): + super().__init__() + self.sigmas = ms.Tensor(sigmas, ms.float32) + + def construct(self, pred_kpts, gt_kpts, kpt_mask, area): + pred_kpts = ops.cast(pred_kpts, ms.float32) + gt_kpts = ops.cast(gt_kpts, ms.float32) + kpt_mask_f = ops.cast(kpt_mask, ms.float32) + area = ops.cast(area, ms.float32) + + d = ops.pow(pred_kpts[..., 0] - gt_kpts[..., 0], 2) + ops.pow(pred_kpts[..., 1] - gt_kpts[..., 1], 2) + factor = pred_kpts.shape[1] / (ops.sum(kpt_mask_f, dim=1) + 1e-9) + sigmas_sq = ops.pow(2 * self.sigmas, 2) + + e = d / (sigmas_sq * (area + 1e-9) * 2) + loss = factor.expand_dims(-1) * ((1.0 - ops.exp(-e)) * kpt_mask_f) + + return loss.mean(axis=1).sum() + + +class v8PoseLoss(v8DetectionLoss): + """YOLO11 姿态估计任务专用损失函数计算类""" + def __init__(self, model): + super().__init__(model) + self.kpt_shape = getattr(model.model[-1], 'kpt_shape', [17, 3]) + self.nkpt = self.kpt_shape[0] + self.ndim = self.kpt_shape[1] + self.bce_pose = nn.BCEWithLogitsLoss(reduction="mean") + + sigmas = OKS_SIGMA if self.kpt_shape == [17, 3] else (np.ones(self.nkpt) / self.nkpt) + self.keypoint_loss = KeypointLoss(sigmas=sigmas) + + self.hyp_pose = getattr(model, 'hyp_pose', 12.0) + self.hyp_kobj = getattr(model, 'hyp_kobj', 1.0) + + def preprocess_keypoints(self, batch, bs): + """格式化数据集目标关键点""" + keypoints = batch["keypoints"].astype(ms.float32) + batch_idx = batch["batch_idx"].view(-1) + counts = [int((batch_idx == i).sum()) for i in range(bs)] + max_obj = max(max(counts) if counts else 0, 1) + + targets_kpt = ops.zeros((bs, max_obj, self.nkpt, self.ndim), keypoints.dtype) + for i in range(bs): + num = counts[i] + if num > 0: + actual_num = min(num, max_obj) + targets_kpt[i, :actual_num] = keypoints[batch_idx == i][:actual_num] + + if hasattr(self, 'imgsz') and self.imgsz: + targets_kpt[..., 0] *= self.imgsz + targets_kpt[..., 1] *= self.imgsz + + return targets_kpt + + def kpts_decode(self, anchor_points, pred_kpts, stride_tensor): + """还原预测关键点的物理坐标""" + anc = anchor_points.view(1, anchor_points.shape[0], 1, 2) + strides = stride_tensor.view(1, anchor_points.shape[0], 1, 1) + xy = (pred_kpts[..., :2] * 2.0 + (anc - 0.5)) * strides + vis = pred_kpts[..., 2:3] + return ops.concat((xy, vis), axis=-1) + + def calculate_keypoints_loss(self, fg_mask, target_gt_idx, targets_kpt, stride_tensor, target_bboxes, pred_kpts_decoded, pred_kpts_raw): + bs = fg_mask.shape[0] + loss_pose, loss_kobj = ops.zeros((), ms.float32), ops.zeros((), ms.float32) + + for i in range(bs): + mask_indices = ops.nonzero(fg_mask[i]).view(-1) + if mask_indices.shape[0] == 0: + continue + + p_kpt_dec = pred_kpts_decoded[i][mask_indices] + p_kpt_raw = pred_kpts_raw[i][mask_indices] + current_target_idx = target_gt_idx[i].reshape(-1) + obj_idx = current_target_idx[mask_indices] + gt_kpt = targets_kpt[i][obj_idx] + + gt_kpt_pixel = gt_kpt[..., :2] + t_box = target_bboxes[i][mask_indices] + + width_f32 = ops.cast(t_box[:, 2] - t_box[:, 0], ms.float32) + height_f32 = ops.cast(t_box[:, 3] - t_box[:, 1], ms.float32) + area_pixel = (width_f32 * height_f32).view(-1, 1) + + kpt_mask = gt_kpt[..., 2] > 0 + loss_pose += self.keypoint_loss(p_kpt_dec[..., :2], gt_kpt_pixel, kpt_mask, area_pixel) + + if self.ndim == 3: + kpt_mask_f = ops.cast(kpt_mask, ms.float32) + loss_kobj += self.bce_pose(p_kpt_raw[..., 2], kpt_mask_f) + + return loss_pose, loss_kobj + + def construct(self, preds, batch): + feats, pred_kpts = preds[0], preds[1] + pred_distri, pred_scores = self.decode_predictions(feats) + bs = pred_scores.shape[0] + + kpt_list = [pk.view(pk.shape[0], pk.shape[1], -1) for pk in pred_kpts] + pred_kpts_cat = ops.concat(kpt_list, axis=2) + pred_kpts = pred_kpts_cat.swapaxes(1, 2).view(bs, -1, self.nkpt, self.ndim) + + anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5) + + targets = self.preprocess_targets(batch, feats) + gt_labels, gt_bboxes = targets.split((1, 4), 2) + + if gt_bboxes.max() <= 1.0: + gt_bboxes = gt_bboxes * getattr(self, 'imgsz', 640) + + mask_gt = gt_bboxes.sum(2, keepdim=True) > 0 + targets_kpt = self.preprocess_keypoints(batch, bs) + + pred_bboxes = self.bbox_decode(anchor_points, pred_distri) + current_gt_bboxes = xywh2xyxy(gt_bboxes) + + pred_scores_sg = ops.stop_gradient(pred_scores.sigmoid()) + pred_bboxes_sg = ops.stop_gradient(pred_bboxes * stride_tensor) + + _, target_bboxes, target_scores, fg_mask, target_gt_idx = self.assigner( + pred_scores_sg, pred_bboxes_sg, anchor_points * stride_tensor, + gt_labels, current_gt_bboxes, mask_gt + ) + + fg_mask_bool = ops.cast(fg_mask, ms.bool_) + pred_kpts_decoded = self.kpts_decode(anchor_points, pred_kpts, stride_tensor) + + target_scores_sum = ops.maximum(target_scores.sum(), 1.0) + loss_cls = self.bce(pred_scores, target_scores).sum() / target_scores_sum + + loss_box, loss_dfl, loss_pose, loss_kobj = ops.zeros((), ms.float32), ops.zeros((), ms.float32), ops.zeros((), ms.float32), ops.zeros((), ms.float32) + + if fg_mask_bool.any(): + loss_box, loss_dfl = self.bbox_loss( + pred_distri, pred_bboxes, anchor_points, + target_bboxes / stride_tensor, target_scores, target_scores_sum, fg_mask_bool + ) + loss_pose, loss_kobj = self.calculate_keypoints_loss( + fg_mask_bool, target_gt_idx, targets_kpt, + stride_tensor, target_bboxes, pred_kpts_decoded, pred_kpts + ) + loss_pose /= target_scores_sum + + return (loss_box * self.hyp_box), (loss_cls * self.hyp_cls), (loss_dfl * self.hyp_dfl), (loss_pose * self.hyp_pose), (loss_kobj * self.hyp_kobj) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/metrics.py b/src/mindnlp/models/ultralytics/utils/metrics.py new file mode 100644 index 000000000..bf6d1c286 --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/metrics.py @@ -0,0 +1,315 @@ +import numpy as np +import mindspore as ms +from mindspore import ops + +def box_iou(box1, box2, eps=1e-7): + """ + 计算两组边界框的交并比 (IoU) + + Args: + box1: (N, 4) 形状的张量，格式为 [x1, y1, x2, y2] + box2: (M, 4) 形状的张量，格式为 [x1, y1, x2, y2] + + Returns: + (N, M) 形状的 IoU 矩阵 + """ + be1 = ops.expand_dims(box1, 1) + be2 = ops.expand_dims(box2, 0) + + lt = ops.maximum(be1[..., :2], be2[..., :2]) + rb = ops.minimum(be1[..., 2:], be2[..., 2:]) + + wh = ops.clamp(rb - lt, min=0) + inter = wh[..., 0] * wh[..., 1] + + area1 = (box1[..., 2] - box1[..., 0]) * (box1[..., 3] - box1[..., 1]) + area2 = (box2[..., 2] - box2[..., 0]) * (box2[..., 3] - box2[..., 1]) + union = ops.expand_dims(area1, 1) + ops.expand_dims(area2, 0) - inter + eps + + return inter / union + + +def ap_per_class(tp, conf, pred_cls, target_cls, eps=1e-16): + """ + 计算各分类的 Average Precision (AP) + 通过对 Precision-Recall 曲线进行数值积分实现 + """ + i = np.argsort(-conf) + tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] + + unique_classes, nt = np.unique(target_cls, return_counts=True) + nc = unique_classes.shape[0] + + ap = np.zeros((nc, tp.shape[1])) + for ci, c in enumerate(unique_classes): + i = (pred_cls == c) + n_l = nt[ci] + n_p = i.sum() + + if n_p == 0 or n_l == 0: + continue + + fpc = (1 - tp[i]).cumsum(0) + tpc = tp[i].cumsum(0) + + recall = tpc / (n_l + eps) + precision = tpc / (tpc + fpc) + + for j in range(tp.shape[1]): + ap[ci, j] = compute_ap(recall[:, j], precision[:, j]) + + return ap, unique_classes + + +def compute_ap(recall, precision, method='interp'): + """ + 基于 Precision-Recall 序列计算平均精度 (AP) + 支持 101 点插值法 (COCO 默认) 或连续积分法 + """ + mrec = np.concatenate(([0.0], recall, [1.0])) + mpre = np.concatenate(([1.0], precision, [0.0])) + + # 计算 Precision 包络线，确保单调递减特性 + mpre = np.maximum.accumulate(mpre[::-1])[::-1] + + if method == 'interp': + x = np.linspace(0, 1, 101) + ap = np.trapz(np.interp(x, mrec, mpre), x) + else: + i = np.where(mrec[1:] != mrec[:-1])[0] + ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) + + return ap + + +def process_batch(detections, labels, iou_thresholds): + """ + 基于不同的 IoU 阈值，判定检测结果的真阳性 (TP) 状态 + + Args: + detections: (N, 6) 格式为 [x1, y1, x2, y2, conf, cls] + labels: (M, 5) 格式为 [cls, x1, y1, x2, y2] + iou_thresholds: 评估用的 IoU 阈值数组 + + Returns: + tp: (N, len(iou_thresholds)) 形状的布尔矩阵 + """ + tp = np.zeros((detections.shape[0], iou_thresholds.shape[0]), dtype=bool) + + if labels.shape[0] == 0: + return tp + + ious = box_iou(ms.Tensor(detections[:, :4]), ms.Tensor(labels[:, 1:])).asnumpy() + + detected_cls = detections[:, 5] + target_cls = labels[:, 0] + + for j, iou_thr in enumerate(iou_thresholds): + matched_gt = np.zeros(labels.shape[0], dtype=bool) + + for i, det in enumerate(detections): + cls = det[5] + matches = np.where((target_cls == cls) & (matched_gt == 0) & (ious[i] >= iou_thr))[0] + + if matches.shape[0] > 0: + best_match = matches[ious[i, matches].argmax()] + tp[i, j] = True + matched_gt[best_match] = True + + return tp + + +def mask_iou(mask1, mask2, eps=1e-7): + """计算像素级掩码的交并比 (Mask IoU)""" + m1 = mask1.view(mask1.shape[0], -1).astype(ms.float32) + m2 = mask2.view(mask2.shape[0], -1).astype(ms.float32) + + intersection = ops.matmul(m1, m2.T) + area1 = m1.sum(1).expand_dims(1) + area2 = m2.sum(1).expand_dims(0) + union = area1 + area2 - intersection + eps + + return intersection / union + + +# 各类任务评估计算 + +class ClassifyMetrics: + """图像分类任务评估类，计算 Top-1 和 Top-5 准确率""" + def __init__(self): + self.top1 = 0.0 + self.top5 = 0.0 + self.task = "classify" + + def process(self, targets, preds): + if not targets or not preds: + return + + preds = np.concatenate(preds, axis=0) + targets = np.concatenate(targets, axis=0) + + correct = (targets[:, None] == preds) + + self.top1 = float(correct[:, 0].mean()) + self.top5 = float(correct.any(axis=1).mean()) + + @property + def fitness(self): + """分类任务适应度得分 (Top-1 与 Top-5 的均值)""" + return (self.top1 + self.top5) / 2.0 + + @property + def results_dict(self): + return { + "metrics/accuracy_top1": self.top1, + "metrics/accuracy_top5": self.top5, + "fitness": self.fitness + } + + +class DetMetrics: + """目标检测任务评估类，维护预测与标签状态，计算全类别 mAP""" + def __init__(self, names): + self.names = names + self.nc = len(names) + self.stats = [] + + def update_stats(self, tp, conf, pred_cls, target_cls): + self.stats.append({ + "tp": tp, "conf": conf, "pred_cls": pred_cls, "target_cls": target_cls + }) + + def process(self): + if not self.stats: + return {"metrics/mAP50(B)": 0, "metrics/mAP50-95(B)": 0} + + tp = np.concatenate([x["tp"] for x in self.stats], 0) + conf = np.concatenate([x["conf"] for x in self.stats], 0) + pred_cls = np.concatenate([x["pred_cls"] for x in self.stats], 0) + target_cls = np.concatenate([x["target_cls"] for x in self.stats], 0) + + ap, _ = ap_per_class(tp, conf, pred_cls, target_cls) + + return { + "metrics/mAP50(B)": ap[:, 0].mean(), + "metrics/mAP50-95(B)": ap.mean() + } + + @property + def fitness(self): + """检测任务适应度得分：0.1 * mAP@50 + 0.9 * mAP@50-95""" + stats = self.process() + return stats.get("metrics/mAP50(B)", 0) * 0.1 + stats.get("metrics/mAP50-95(B)", 0) * 0.9 + + +class SegmentMetrics(DetMetrics): + """实例分割任务评估类，计算边界框及掩码的双重 mAP""" + def __init__(self, names): + super().__init__(names) + self.stats_m = [] + + def update_stats(self, tp_b, tp_m, conf, pred_cls, target_cls): + super().update_stats(tp_b, conf, pred_cls, target_cls) + self.stats_m.append(tp_m) + + def process(self): + results = super().process() + if not self.stats_m: + return results + + tp_m = np.concatenate(self.stats_m, 0) + ap_m, _ = ap_per_class( + tp_m, + np.concatenate([x["conf"] for x in self.stats], 0), + np.concatenate([x["pred_cls"] for x in self.stats], 0), + np.concatenate([x["target_cls"] for x in self.stats], 0) + ) + + results.update({ + "metrics/mAP50(M)": ap_m[:, 0].mean(), + "metrics/mAP50-95(M)": ap_m.mean() + }) + return results + + @property + def fitness(self): + stats = self.process() + # 计算目标检测 (Box) 的分数 + box_fitness = stats.get("metrics/mAP50(B)", 0) * 0.1 + stats.get("metrics/mAP50-95(B)", 0) * 0.9 + # 计算实例分割 (Mask) 的分数 + mask_fitness = stats.get("metrics/mAP50(M)", 0) * 0.1 + stats.get("metrics/mAP50-95(M)", 0) * 0.9 + # 综合两者的表现 + return box_fitness + mask_fitness + + +OKS_SIGMA = np.array( + [0.26, 0.25, 0.25, 0.35, 0.35, 0.79, 0.79, 0.72, 0.72, 0.62, 0.62, 1.07, 1.07, 0.87, 0.87, 0.89, 0.89], + dtype=np.float32, +) / 10.0 + + +def kpt_iou(kpt1, kpt2, area, sigma, eps=1e-7): + """计算关键点目标相似度 (Object Keypoint Similarity, OKS)""" + p_xy = ops.expand_dims(kpt1[..., :2], 1) + g_xy = ops.expand_dims(kpt2[..., :2], 0) + + d2 = ops.pow(p_xy - g_xy, 2).sum(-1) + + g_mask = kpt2[..., 2] > 0 + g_mask_exp = ops.expand_dims(g_mask, 0) + + sigma = ms.Tensor(sigma, dtype=kpt1.dtype) + sig2 = ops.pow(sigma * 2, 2) + area_exp = area.view(1, -1, 1) + + denominator = area_exp * sig2 * 2 + + exponent = -d2 / (denominator + eps) + exponent = ops.clamp(exponent, -100, 0) + oks_all = ops.exp(exponent) + + g_mask_float = g_mask_exp.astype(kpt1.dtype) + num = (oks_all * g_mask_float).sum(-1) + den = g_mask_float.sum(-1) + eps + + return num / den + + +class PoseMetrics(DetMetrics): + """姿态估计任务评估类，计算边界框及关键点 OKS 的双重 mAP""" + def __init__(self, names): + super().__init__(names) + self.stats_p = [] + + def update_stats(self, tp_b, tp_p, conf, pred_cls, target_cls): + super().update_stats(tp_b, conf, pred_cls, target_cls) + self.stats_p.append(tp_p) + + def process(self): + results = super().process() + if not self.stats_p: + return results + + tp_p = np.concatenate(self.stats_p, 0) + ap_p, _ = ap_per_class( + tp_p, + np.concatenate([x["conf"] for x in self.stats], 0), + np.concatenate([x["pred_cls"] for x in self.stats], 0), + np.concatenate([x["target_cls"] for x in self.stats], 0) + ) + + results.update({ + "metrics/mAP50(P)": ap_p[:, 0].mean() if len(ap_p) else 0.0, + "metrics/mAP50-95(P)": ap_p.mean() if len(ap_p) else 0.0 + }) + return results + + @property + def fitness(self): + stats = self.process() + # 计算目标检测 (Box) 的分数 + box_fitness = stats.get("metrics/mAP50(B)", 0) * 0.1 + stats.get("metrics/mAP50-95(B)", 0) * 0.9 + # 计算姿态估计 (Pose/Keypoint) 的分数 + pose_fitness = stats.get("metrics/mAP50(P)", 0) * 0.1 + stats.get("metrics/mAP50-95(P)", 0) * 0.9 + # 综合两者的表现 + return box_fitness + pose_fitness \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/ops.py b/src/mindnlp/models/ultralytics/utils/ops.py new file mode 100644 index 000000000..1f17d8ff7 --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/ops.py @@ -0,0 +1,401 @@ +import math +import numpy as np +import mindspore as ms +from mindspore import ops + +def xyxy2xywh(x): + """边界框格式转换：(x1, y1, x2, y2) -> (xc, yc, w, h)""" + y = ops.deepcopy(x) if isinstance(x, ms.Tensor) else np.copy(x) + y[..., 0] = (x[..., 0] + x[..., 2]) / 2 + y[..., 1] = (x[..., 1] + x[..., 3]) / 2 + y[..., 2] = x[..., 2] - x[..., 0] + y[..., 3] = x[..., 3] - x[..., 1] + return y + + +def xywh2xyxy(x): + """边界框格式转换：(xc, yc, w, h) -> (x1, y1, x2, y2)""" + y = ops.deepcopy(x) if isinstance(x, ms.Tensor) else np.copy(x) + + y[..., 0] = x[..., 0] - x[..., 2] / 2 # x1 + y[..., 1] = x[..., 1] - x[..., 3] / 2 # y1 + y[..., 2] = x[..., 0] + x[..., 2] / 2 # x2 + y[..., 3] = x[..., 1] + x[..., 3] / 2 # y2 + return y + + +def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0): + """归一化相对坐标转为物理像素坐标""" + y = ops.deepcopy(x) if isinstance(x, ms.Tensor) else np.copy(x) + y[..., 0] = w * (x[..., 0] - x[..., 2] / 2) + padw + y[..., 1] = h * (x[..., 1] - x[..., 3] / 2) + padh + y[..., 2] = w * (x[..., 0] + x[..., 2] / 2) + padw + y[..., 3] = h * (x[..., 1] + x[..., 3] / 2) + padh + return y + + +def clip_boxes(boxes, shape): + """将边界框坐标截断至图像物理边界内""" + h, w = shape[:2] + if isinstance(boxes, ms.Tensor): + boxes[..., 0] = ops.clamp(boxes[..., 0], 0, w) + boxes[..., 1] = ops.clamp(boxes[..., 1], 0, h) + boxes[..., 2] = ops.clamp(boxes[..., 2], 0, w) + boxes[..., 3] = ops.clamp(boxes[..., 3], 0, h) + else: + boxes[..., [0, 2]] = np.clip(boxes[..., [0, 2]], 0, w) + boxes[..., [1, 3]] = np.clip(boxes[..., [1, 3]], 0, h) + return boxes + + +def segment2box(segment, width=640, height=640): + """由不规则多边形轮廓 (Segments) 推导其最小外接矩形 (Bbox)""" + if isinstance(segment, ms.Tensor): + x = segment[..., 0] + y = segment[..., 1] + else: + x, y = segment.T + return np.array([x.min(), y.min(), x.max(), y.max()]) + + +def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): + """ + 计算目标边界框惩罚度量 (MindSpore 算子适配版) + 支持 IoU 及 GIoU, DIoU, CIoU 等变体计算 + """ + if xywh: + x1, y1, w1, h1 = ops.split(box1, 1, axis=-1) + x2, y2, w2, h2 = ops.split(box2, 1, axis=-1) + + w1_half, h1_half = w1 / 2, h1 / 2 + w2_half, h2_half = w2 / 2, h2 / 2 + b1_x1, b1_x2 = x1 - w1_half, x1 + w1_half + b1_y1, b1_y2 = y1 - h1_half, y1 + h1_half + b2_x1, b2_x2 = x2 - w2_half, x2 + w2_half + b2_y1, b2_y2 = y2 - h2_half, y2 + h2_half + else: + # 分离坐标点并确保各个 Tensor 独立 + b1_x1, b1_y1, b1_x2, b1_y2 = ops.split(box1, 1, axis=-1) + b2_x1, b2_y1, b2_x2, b2_y2 = ops.split(box2, 1, axis=-1) + + inter = (ops.minimum(b1_x2, b2_x2) - ops.maximum(b1_x1, b2_x1)).clip(0) * \ + (ops.minimum(b1_y2, b2_y2) - ops.maximum(b1_y1, b2_y1)).clip(0) + + w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps + w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps + union = w1 * h1 + w2 * h2 - inter + eps + + iou = inter / union + + if CIoU or DIoU or GIoU: + cw = ops.maximum(b1_x2, b2_x2) - ops.minimum(b1_x1, b2_x1) + ch = ops.maximum(b1_y2, b2_y2) - ops.minimum(b1_y1, b2_y1) + if CIoU or DIoU: + c2 = cw ** 2 + ch ** 2 + eps + rho2 = ((b1_x1 + b1_x2 - b2_x1 - b2_x2) ** 2 + (b1_y1 + b1_y2 - b2_y1 - b2_y2) ** 2) / 4 + if CIoU: + v = (4 / (math.pi ** 2)) * ops.pow(ops.atan(w2 / h2) - ops.atan(w1 / h1), 2) + v = ops.stop_gradient(v) + alpha = v / ((1 - iou) + v + eps) + return iou - (rho2 / c2 + v * alpha) + return iou - rho2 / c2 + + c_area = cw * ch + eps + return iou - (c_area - union) / c_area + + return iou + + +def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None): + """基于缩放比例系数，将边界框坐标还原回原始图像的空间尺度""" + if ratio_pad is None: + gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 + else: + gain = ratio_pad[0][0] + pad = ratio_pad[1] + + boxes[..., [0, 2]] -= pad[0] + boxes[..., [1, 3]] -= pad[1] + boxes[..., :4] /= gain + return clip_boxes(boxes, img0_shape) + + +def dist2bbox(distance, anchor_points, strides=None, xywh=True, axis=-1): + # 1. 拆分偏移量 + lt, rb = ops.split(distance, split_size_or_sections=2, axis=axis) + + # 2. 锚点维度自适应 + ap = anchor_points + if axis == 1 or axis == -2: # Detect + if ap.ndim == 2: ap = ops.expand_dims(ap, 0).transpose(0, 2, 1) + elif ap.ndim == 3 and ap.shape[-1] == 2: ap = ap.transpose(0, 2, 1) + else: # Segment/Pose + if ap.ndim == 2: ap = ops.expand_dims(ap, 0) + elif ap.ndim == 3 and ap.shape[1] == 2: ap = ap.transpose(0, 2, 1) + + # 3. 计算坐标 + x1y1 = ap - lt + x2y2 = ap + rb + + # 4. 步长缩放 + if strides is not None: + s = strides + if axis == 1 or axis == -2: + if s.ndim == 1: s = s.view(1, 1, -1) + elif s.ndim == 2: s = ops.expand_dims(s, 1) + elif s.ndim == 3 and s.shape[-1] == 1: s = s.transpose(0, 2, 1) + else: + if s.ndim == 1: s = s.view(1, -1, 1) + elif s.ndim == 2: s = ops.expand_dims(s, -1) + elif s.ndim == 3 and s.shape[1] == 1: s = s.transpose(0, 2, 1) + x1y1 = x1y1 * s + x2y2 = x2y2 * s + + # 5. 格式输出 + if xywh: + c_xy = (x1y1 + x2y2) / 2 + wh = x2y2 - x1y1 + return ops.concat((c_xy, wh), axis=axis) + return ops.concat((x1y1, x2y2), axis=axis) + +def bbox2dist(anchor_points, bbox, reg_max): + """ + 将边界框(x1y1x2y2)转换为相对于锚点的偏移量(ltrb) + """ + x1y1, x2y2 = ops.split(bbox, split_size_or_sections=2, axis=-1) + lt = anchor_points - x1y1 + rb = x2y2 - anchor_points + dist = ops.concat((lt, rb), axis=-1) + return dist.clip(0, reg_max - 0.01) # 限制范围防止溢出 + + +def make_divisible(x, divisor): + """调整特征通道数，使其严格对齐硬件算子的内存对齐规范 (通常为 8 的倍数)""" + if isinstance(divisor, ms.Tensor): + divisor = int(divisor.max().asnumpy()) + return math.ceil(x / divisor) * divisor + + +def batch_iou(batch_box1, batch_box2, eps=1e-7): + """ + 计算批处理数据的交叉 IoU 矩阵 + 内置 Batch 截断策略，以兼容 DataLoader 处理尾部不完整批次时的隐式 Padding 行为 + """ + bs1 = batch_box1.shape[0] + bs2 = batch_box2.shape[0] + + # 应对静态图引擎 Padding 造成的批量不一致情况 + if bs1 != bs2: + actual_bs = min(bs1, bs2) + batch_box1 = batch_box1[:actual_bs] + batch_box2 = batch_box2[:actual_bs] + + b1_x1, b1_y1, b1_x2, b1_y2 = ops.split(batch_box1.expand_dims(2), 1, axis=-1) + b2_x1, b2_y1, b2_x2, b2_y2 = ops.split(batch_box2.expand_dims(1), 1, axis=-1) + + inter_w = (ops.minimum(b1_x2, b2_x2) - ops.maximum(b1_x1, b2_x1)).clip(0) + inter_h = (ops.minimum(b1_y2, b2_y2) - ops.maximum(b1_y1, b2_y1)).clip(0) + inter_area = inter_w * inter_h + + area1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1) + area2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + + union_area = area1 + area2 - inter_area + eps + + return inter_area / union_area + + +def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, max_det=300, nc=80): + """ + 通用非极大值抑制 (Non-Maximum Suppression, NMS) + 自适应目标检测、实例分割以及姿态估计的特征图切片解析 + + Args: + prediction: 网络输出的预测张量 + conf_thres: 用于初筛目标的置信度阈值 + iou_thres: 用于框去重的重叠度阈值 + max_det: 单张图像最大保留的目标数 + nc: 任务的基础类别数，用于动态推断后置的额外特征维度 + """ + if isinstance(prediction, (list, tuple)): + prediction = prediction[0] + + if prediction.shape[1] < prediction.shape[2]: + prediction = prediction.transpose(0, 2, 1) + + bs = prediction.shape[0] + + # 根据额外的通道特征自动判断是否附带了分割掩码系数或姿态关键点 + extra_dim = prediction.shape[2] - 4 - nc + + output = [ops.zeros((0, 6 + extra_dim), prediction.dtype)] * bs + + for xi, x in enumerate(prediction): + x = x.astype(ms.float32) + + box = x[:, :4] + cls = x[:, 4: (4 + nc)] + + conf, cls_id = ops.max(cls, axis=-1) + mask = conf > conf_thres + + if not mask.any(): + continue + + x_filtered = x[mask] + curr_box = xywh2xyxy(x_filtered[:, :4]) + curr_conf = conf[mask].expand_dims(-1).astype(ms.float32) + curr_clsid = cls_id[mask].expand_dims(-1).astype(ms.float32) + + if extra_dim > 0: + curr_extras = x_filtered[:, (4 + nc): (4 + nc + extra_dim)].astype(ms.float32) + x_combined = ops.concat((curr_box, curr_conf, curr_clsid, curr_extras), -1) + else: + x_combined = ops.concat((curr_box, curr_conf, curr_clsid), -1) + + x_np = x_combined.asnumpy() + + indices = np.argsort(-x_np[:, 4]) + x_np = x_np[indices] + + keep = [] + while x_np.shape[0] > 0: + current_box = x_np[0] + keep.append(current_box) + if x_np.shape[0] == 1: + break + + ious = calculate_nms_iou(current_box[:4], x_np[1:, :4]) + x_np = x_np[1:][ious < iou_thres] + + if len(keep) > 0: + res = np.stack(keep) + out_tensor = ms.Tensor(res[:max_det], prediction.dtype) + output[xi] = out_tensor + + return output + + +def calculate_nms_iou(box1, boxes): + """提供纯 NumPy 环境支持的 IoU 计算，辅助完成 NMS""" + x1 = np.maximum(box1[0], boxes[:, 0]) + y1 = np.maximum(box1[1], boxes[:, 1]) + x2 = np.minimum(box1[2], boxes[:, 2]) + y2 = np.minimum(box1[3], boxes[:, 3]) + + inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1) + area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) + area2 = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + + return inter / (area1 + area2 - inter + 1e-7) + + +def xywh2xyxy_np(x): + """基于 NumPy 的边界框格式解析函数，主要用于脱离计算图的后置评估流""" + y = np.copy(x) + y[..., 0] = x[..., 0] - x[..., 2] / 2 + y[..., 1] = x[..., 1] - x[..., 3] / 2 + y[..., 2] = x[..., 0] + x[..., 2] / 2 + y[..., 3] = x[..., 1] + x[..., 3] / 2 + return y + + +def crop_mask(masks, boxes): + """ + 掩码边界裁剪操作 + """ + n, h, w = masks.shape + + x1, y1, x2, y2 = ops.split(boxes, 1, axis=-1) + + x1 = x1.expand_dims(-1) + y1 = y1.expand_dims(-1) + x2 = x2.expand_dims(-1) + y2 = y2.expand_dims(-1) + + r_x = ops.arange(w, dtype=ms.float32).view(1, 1, w) + r_y = ops.arange(h, dtype=ms.float32).view(1, h, 1) + + mask_x = ops.logical_and(r_x >= x1, r_x < x2).astype(ms.float32) + mask_y = ops.logical_and(r_y >= y1, r_y < y2).astype(ms.float32) + + crop_mask_val = mask_x * mask_y + + return masks * crop_mask_val + + +def process_mask(protos, masks_in, bboxes, shape, upsample=False): + """ + 通过模型产生的原型映射层 (Protos) 与掩码系数生成特定目标实例分割图 + 支持可选的双线性插值上采样操作 + """ + c, mh, mw = protos.shape + + masks_in = masks_in.astype(ms.float32) + protos_flat = protos.view(c, -1).astype(ms.float32) + + masks = ops.matmul(masks_in, protos_flat) + masks = ops.sigmoid(masks).view(-1, mh, mw) + + width_ratio = mw / shape[1] + height_ratio = mh / shape[0] + + scale_tensor = ms.Tensor([width_ratio, height_ratio, width_ratio, height_ratio], dtype=ms.float32) + scaled_bboxes = bboxes.astype(ms.float32) * scale_tensor + + masks = crop_mask(masks, scaled_bboxes) + + if upsample: + masks = ops.interpolate(masks.expand_dims(1), size=shape, mode='bilinear', align_corners=False).squeeze(1) + + return ops.cast(masks > 0.5, ms.float32) + + +def process_mask_native(protos, masks_in, bboxes, shape): + """原图尺寸下采样后方执行裁剪校验""" + c, mh, mw = protos.shape + masks = ops.matmul(masks_in, protos.reshape(c, -1)).reshape(-1, mh, mw) + + masks = ops.interpolate(masks.expand_dims(1), size=shape, mode='bilinear', align_corners=False).squeeze(1) + masks = crop_mask(masks, bboxes) + + return masks > 0 + + +def clip_coords(coords, shape): + """对点坐标集实行裁剪截断边界约束""" + h, w = shape[:2] + if isinstance(coords, ms.Tensor): + coords[..., 0] = ops.clamp(coords[..., 0], 0, w) + coords[..., 1] = ops.clamp(coords[..., 1], 0, h) + else: + coords[..., 0] = np.clip(coords[..., 0], 0, w) + coords[..., 1] = np.clip(coords[..., 1], 0, h) + return coords + + +def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): + """ + 基于预定义的空间比例，将目标点列 (如姿态点集) 从计算域尺度等距缩放回原图尺度域 + """ + if ratio_pad is None: + gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 + else: + gain = ratio_pad[0][0] + pad = ratio_pad[1] + + if isinstance(coords, ms.Tensor): + coords = coords.astype(ms.float32) + else: + coords = np.asarray(coords).astype(np.float32) + + coords[..., 0] -= pad[0] + coords[..., 1] -= pad[1] + coords[..., 0] /= gain + coords[..., 1] /= gain + + coords = clip_coords(coords, img0_shape) + + return coords \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/optimizer.py b/src/mindnlp/models/ultralytics/utils/optimizer.py new file mode 100644 index 000000000..aba3e3d6d --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/optimizer.py @@ -0,0 +1,89 @@ +import math +import mindspore.nn as nn + +def get_lr(args, hyp, steps_per_epoch): + """ + 生成带有线性预热 (Linear Warmup) 与余弦退火 (Cosine Annealing) 的学习率调度序列 + + Args: + args: 包含全局训练轮数 (epochs) 等宏观配置的参数对象 + hyp (dict): 包含学习率、预热轮数等微观调优配置的超参数字典 + steps_per_epoch (int): 每个 Epoch 包含的迭代步数 + + Returns: + list[float]: 按步长展开的学习率列表 + """ + total_steps = args.epochs * steps_per_epoch + + # 从超参数字典中安全提取参数，提供默认值作为工程兜底 + warmup_epochs = hyp.get('warmup_epochs', 3.0) + warmup_steps = int(warmup_epochs * steps_per_epoch) + + base_lr = hyp.get('lr0', 0.01) + lrf = hyp.get('lrf', 0.01) # 最终学习率比例 + min_lr = base_lr * lrf + + lr_each_step = [] + for i in range(total_steps): + if i < warmup_steps: + # 预热阶段：从极小值线性增长至初始学习率 base_lr + lr = base_lr * (i + 1) / warmup_steps + else: + # 余弦退火阶段 + cur_step = i - warmup_steps + total_decay_steps = total_steps - warmup_steps + + # 使用标准的余弦退火公式计算当前步的学习率 + decay_ratio = 0.5 * (1 + math.cos(math.pi * cur_step / total_decay_steps)) + lr = min_lr + (base_lr - min_lr) * decay_ratio + + lr_each_step.append(lr) + + return lr_each_step + +def build_optimizer(model, lr_list, hyp): + """ + 基于参数分组与配置字典构建优化器 + + 对一维张量（如 BatchNorm 权重）和偏置项（Bias）取消权重衰减（Weight Decay）， + 以防限制模型的拟合能力或导致梯度异常 + + Args: + model (nn.Cell): 待优化的网络模型 + lr_list (list[float]): 预先计算好的学习率步长序列 + hyp (dict): 超参数字典，需包含 'optimizer', 'weight_decay', 'momentum' 等键 + + Returns: + nn.Optimizer: 实例化后的 MindSpore 优化器 + """ + decay_params = [] + no_decay_params = [] + + # 执行参数分组：过滤不需要 L2 正则化的参数 + for param in model.trainable_params(): + if len(param.shape) == 1 or param.name.endswith(".bias"): + no_decay_params.append(param) + else: + decay_params.append(param) + + weight_decay = hyp.get('weight_decay', 0.0005) + + group_params = [ + {'params': decay_params, 'weight_decay': weight_decay}, + {'params': no_decay_params, 'weight_decay': 0.0} + ] + + # 动态解析并构建指定的优化器 + opt_name = hyp.get('optimizer', 'SGD').upper() + momentum = hyp.get('momentum', 0.937) + + if opt_name == 'SGD': + # YOLO 官方推荐默认使用带 Nesterov 动量的 SGD + optimizer = nn.SGD(group_params, learning_rate=lr_list, momentum=momentum, nesterov=True) + elif opt_name in ['ADAM', 'ADAMW']: + # 微调或特定任务下使用的 AdamW + optimizer = nn.AdamWeightDecay(group_params, learning_rate=lr_list, beta1=momentum, beta2=0.999) + else: + raise ValueError(f"[ERROR] 不支持的优化器类型: {opt_name}。请在 hyp.yaml 中指定 SGD 或 AdamW。") + + return optimizer \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/utils/tal.py b/src/mindnlp/models/ultralytics/utils/tal.py new file mode 100644 index 000000000..8cd109ea7 --- /dev/null +++ b/src/mindnlp/models/ultralytics/utils/tal.py @@ -0,0 +1,140 @@ +import numpy as np +import mindspore as ms +from mindspore import nn, ops + +# 规范引入算子 +from utils.ops import batch_iou + +def make_anchors(feats, strides, grid_cell_offset=0.5): + """ + 根据特征图的尺度生成对应的锚点中心坐标与步长张量 + + Args: + feats (list[Tensor]): 多尺度特征图列表 + strides (list[int]): 对应特征图的下采样步长 + grid_cell_offset (float): 网格中心点偏移量，默认 0.5 + + Returns: + tuple[Tensor, Tensor]: 展平后的锚点坐标张量与对应的步长张量 + """ + anchor_points, stride_tensor = [], [] + dtype = feats[0].dtype + + for i, stride in enumerate(strides): + if i >= len(feats): + break + + shape = feats[i].shape + # 解析特征图空间维度 (支持 [B, C, H, W] 或 [B, H, W] 格式) + if len(shape) == 4: + _, _, h, w = shape + elif len(shape) == 3: + _, h, w = shape + else: + raise ValueError(f"[ERROR] 特征图维度解析失败，期望维度为 3 或 4，实际为 {len(shape)}。") + + sx = ops.arange(w).astype(dtype) + grid_cell_offset + sy = ops.arange(h).astype(dtype) + grid_cell_offset + + # 生成网格坐标 (基于 ij 索引) + grid_y, grid_x = ops.meshgrid(sy, sx, indexing='ij') + + anchor_points.append(ops.stack((grid_x, grid_y), -1).view(-1, 2)) + stride_tensor.append(ops.full((h * w, 1), stride, dtype=dtype)) + + return ops.concat(anchor_points, 0), ops.concat(stride_tensor, 0) + + +class TaskAlignedAssigner(nn.Cell): + """ + 任务对齐样本分配器 (Task-Aligned Assigner) + 依据分类得分与边界框 IoU 的加权指标，为真实目标动态分配正样本锚点 + """ + def __init__(self, topk=10, num_classes=80, alpha=1.0, beta=6.0, eps=1e-9): + super().__init__() + self.topk = topk + self.num_classes = num_classes + self.bg_idx = num_classes + self.alpha = alpha + self.beta = beta + self.eps = eps + + def construct(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt): + """ + 任务对齐分配器：修正冲突消解与软标签逻辑。 + """ + bs = pd_bboxes.shape[0] + num_anchors = pd_bboxes.shape[1] + n_max_boxes = gt_bboxes.shape[1] + + # 1. 计算 IoU [bs, n_max_boxes, num_anchors] + ious = batch_iou(gt_bboxes, pd_bboxes) + if ious.ndim == 4: ious = ious.squeeze(-1) + + # 2. 提取对应 GT 类别的预测得分 + t_labels = gt_labels.squeeze(-1).astype(ms.int32) + # 构造索引矩阵 [bs, n_max_boxes, num_anchors] + scores = ops.gather_elements( + pd_scores.transpose(0, 2, 1), + 1, + ops.broadcast_to(t_labels.expand_dims(-1), (bs, n_max_boxes, num_anchors)) + ) + + # 3. 计算对齐度 Alignment Metric + align_metric = ops.pow(scores, self.alpha) * ops.pow(ious, self.beta) + + # 4. 选取 Top-k 候选 + topk_values, topk_indices = ops.topk(align_metric, self.topk, dim=-1) + is_in_topk = ops.tensor_scatter_elements( + ops.zeros_like(align_metric), topk_indices, ops.ones_like(topk_values), axis=2 + ) + + # 5. 中心点过滤 (确保锚点在 GT 框内) + mask_in_gts = self.select_candidates_in_gts(anc_points, gt_bboxes) + mask_pos = is_in_topk * mask_in_gts.astype(ms.float32) * mask_gt.astype(ms.float32) + + # 6. 每个锚点只允许对应一个最优 GT + mask_pos_sum = mask_pos.sum(-2) # [bs, num_anchors] + if (mask_pos_sum > 1).any(): + # 找到每个锚点 metric 最大的 GT 索引 + max_idx = align_metric.argmax(1) # [bs, num_anchors] + one_hot_mask = ops.one_hot(max_idx, n_max_boxes, 1.0, 0.0).transpose(0, 2, 1) + mask_pos *= one_hot_mask + + # 7. 生成最终标签 + fg_mask = mask_pos.sum(1) > 0 # 前景掩码 + best_gt_idx = mask_pos.argmax(1) # [bs, num_anchors] + + # 提取 target_bboxes + target_bboxes = ops.gather_elements( + ops.broadcast_to(gt_bboxes.expand_dims(1), (bs, num_anchors, n_max_boxes, 4)), + 2, + best_gt_idx.view(bs, num_anchors, 1, 1).expand_as(ops.zeros((bs, num_anchors, 1, 4), ms.int32)) + ).squeeze(2) + + # 提取 target_labels 并生成 Soft Targets (分类分 * IoU) + target_labels = ops.gather_elements(t_labels, 1, best_gt_idx) + target_scores = ops.one_hot(target_labels, self.num_classes, 1.0, 0.0) + + # 结合 IoU 调整置信度 + max_iou = ops.gather_elements(ious.transpose(0, 2, 1), 2, best_gt_idx.expand_dims(-1)).squeeze(-1) + target_scores *= max_iou.expand_dims(-1) + target_scores *= fg_mask.expand_dims(-1).astype(ms.float32) + + return target_labels, target_bboxes, target_scores, fg_mask.astype(ms.bool_), best_gt_idx + + def select_candidates_in_gts(self, anc_points, gt_bboxes): + """ + 中心点空间过滤策略 + 判定网格中心锚点是否物理落入对应的真实边界框范围内 + """ + anc_points = anc_points.expand_dims(0).expand_dims(0) + + lt = gt_bboxes[..., :2].expand_dims(2) + rb = gt_bboxes[..., 2:].expand_dims(2) + + # 计算从中心点到边界框四个边的几何增量 + bbox_deltas = ops.concat((anc_points - lt, rb - anc_points), axis=-1) + + # 空间范围校验：所有几何增量必须大于极小值 + return bbox_deltas.amin(axis=-1) > self.eps From 772b35b9c8a4a851924d3787e2f78b5f46b84ad7 Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Thu, 19 Mar 2026 14:56:50 +0800 Subject: [PATCH 2/6] fix: add datasets config files --- .../ultralytics/cfg/datasets/coco128-seg.yaml | 101 ++++++++++++++++++ .../ultralytics/cfg/datasets/coco128.yaml | 101 ++++++++++++++++++ .../ultralytics/cfg/datasets/coco8-pose.yaml | 47 ++++++++ .../cfg/datasets/imagenette2-160.yaml | 28 +++++ 4 files changed, 277 insertions(+) create mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml create mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml new file mode 100644 index 000000000..a496fdb1e --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml @@ -0,0 +1,101 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO128-seg dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/segment/coco/ +# Example usage: yolo train data=coco128-seg.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco128-seg ← downloads here (7 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: /root/autodl-tmp/ultralytics/datasets/coco128-seg # dataset root dir +train: /root/autodl-tmp/ultralytics/datasets/coco128-seg/images/train2017 # train images (relative to 'path') 128 images +val: /root/autodl-tmp/ultralytics/datasets/coco128-seg/images/train2017 # val images (relative to 'path') 128 images +test: # test images (optional) + +# Classes +names: + 0: person + 1: bicycle + 2: car + 3: motorcycle + 4: airplane + 5: bus + 6: train + 7: truck + 8: boat + 9: traffic light + 10: fire hydrant + 11: stop sign + 12: parking meter + 13: bench + 14: bird + 15: cat + 16: dog + 17: horse + 18: sheep + 19: cow + 20: elephant + 21: bear + 22: zebra + 23: giraffe + 24: backpack + 25: umbrella + 26: handbag + 27: tie + 28: suitcase + 29: frisbee + 30: skis + 31: snowboard + 32: sports ball + 33: kite + 34: baseball bat + 35: baseball glove + 36: skateboard + 37: surfboard + 38: tennis racket + 39: bottle + 40: wine glass + 41: cup + 42: fork + 43: knife + 44: spoon + 45: bowl + 46: banana + 47: apple + 48: sandwich + 49: orange + 50: broccoli + 51: carrot + 52: hot dog + 53: pizza + 54: donut + 55: cake + 56: chair + 57: couch + 58: potted plant + 59: bed + 60: dining table + 61: toilet + 62: tv + 63: laptop + 64: mouse + 65: remote + 66: keyboard + 67: cell phone + 68: microwave + 69: oven + 70: toaster + 71: sink + 72: refrigerator + 73: book + 74: clock + 75: vase + 76: scissors + 77: teddy bear + 78: hair drier + 79: toothbrush + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128-seg.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml new file mode 100644 index 000000000..3db40eff3 --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml @@ -0,0 +1,101 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO128 dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/detect/coco/ +# Example usage: yolo train data=coco128.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco128 ← downloads here (7 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: /root/autodl-tmp/ultralytics/datasets/coco128/coco128 # dataset root dir +train: images/train2017 # train images (relative to 'path') 128 images +val: images/train2017 # val images (relative to 'path') 128 images +test: # test images (optional) + +# Classes +names: + 0: person + 1: bicycle + 2: car + 3: motorcycle + 4: airplane + 5: bus + 6: train + 7: truck + 8: boat + 9: traffic light + 10: fire hydrant + 11: stop sign + 12: parking meter + 13: bench + 14: bird + 15: cat + 16: dog + 17: horse + 18: sheep + 19: cow + 20: elephant + 21: bear + 22: zebra + 23: giraffe + 24: backpack + 25: umbrella + 26: handbag + 27: tie + 28: suitcase + 29: frisbee + 30: skis + 31: snowboard + 32: sports ball + 33: kite + 34: baseball bat + 35: baseball glove + 36: skateboard + 37: surfboard + 38: tennis racket + 39: bottle + 40: wine glass + 41: cup + 42: fork + 43: knife + 44: spoon + 45: bowl + 46: banana + 47: apple + 48: sandwich + 49: orange + 50: broccoli + 51: carrot + 52: hot dog + 53: pizza + 54: donut + 55: cake + 56: chair + 57: couch + 58: potted plant + 59: bed + 60: dining table + 61: toilet + 62: tv + 63: laptop + 64: mouse + 65: remote + 66: keyboard + 67: cell phone + 68: microwave + 69: oven + 70: toaster + 71: sink + 72: refrigerator + 73: book + 74: clock + 75: vase + 76: scissors + 77: teddy bear + 78: hair drier + 79: toothbrush + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml new file mode 100644 index 000000000..185e366ad --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml @@ -0,0 +1,47 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/pose/coco8-pose/ +# Example usage: yolo train data=coco8-pose.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco8-pose ← downloads here (1 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/coco8-pose # dataset root dir +train: images/train # train images (relative to 'path') 4 images +val: images/val # val images (relative to 'path') 4 images +test: # test images (optional) + +# Keypoints +kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) +flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] + +# Classes +names: + 0: person + +# Keypoint names per class +kpt_names: + 0: + - nose + - left_eye + - right_eye + - left_ear + - right_ear + - left_shoulder + - right_shoulder + - left_elbow + - right_elbow + - left_wrist + - right_wrist + - left_hip + - right_hip + - left_knee + - right_knee + - left_ankle + - right_ankle + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-pose.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml new file mode 100644 index 000000000..3d699147a --- /dev/null +++ b/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml @@ -0,0 +1,28 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license +# Imagenette2-160 Dataset by fast.ai +# Example usage: python train.py --data imagenette2-160.yaml + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/imagenette2-160 # 你的数据集根目录 (根据你的截图和运行路径调整) +train: train # 训练集相对路径 (会与 path 拼接) +val: val # 验证集相对路径 (会与 path 拼接) +test: # 测试集 (Imagenette 没有单独的 test，可以留空) + +# Classes +nc: 10 # 类别数量 +names: + 0: tench # 丁鱼 + 1: English springer # 英国斯宾格犬 + 2: cassette player # 录音机 + 3: chain saw # 链锯 + 4: church # 教堂 + 5: French horn # 圆号 + 6: garbage truck # 垃圾车 + 7: gas pump # 加油泵 + 8: golf ball # 高尔夫球 + 9: parachute # 降落伞 + +# Download script/URL (optional) +download: | + wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz + tar -xzvf imagenette2-160.tgz \ No newline at end of file From 86369c3e46bb5ddc6a0550dd98cf4f083d31808e Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Mon, 30 Mar 2026 22:39:15 +0800 Subject: [PATCH 3/6] feat: migrate ultralytics module to mindnlp core directory and add usage example --- .../ultralytics/cfg/datasets/coco128-seg.yaml | 101 --------- .../ultralytics/cfg/datasets/coco128.yaml | 101 --------- .../ultralytics/cfg/datasets/coco8-pose.yaml | 47 ---- .../cfg/datasets/imagenette2-160.yaml | 28 --- .../models/yolo/segment/__init__.py | 0 .../models/ultralytics/utils/__init__.py | 0 src/mindnlp/ultralytics/__init__.py | 3 + .../{models => }/ultralytics/cfg/hyp.yaml | 0 .../ultralytics/cfg/models/11/yolo11-cls.yaml | 0 .../cfg/models/11/yolo11-pose.yaml | 0 .../ultralytics/cfg/models/11/yolo11-seg.yaml | 0 .../ultralytics/cfg/models/11/yolo11.yaml | 0 .../ultralytics/configuration_yolo.py | 30 +-- .../engine}/__init__.py | 0 .../ultralytics/engine/predictor.py | 55 +++-- .../ultralytics/engine/trainer.py | 3 +- .../ultralytics/engine/validator.py | 2 +- .../examples/yolo/classify/inference.py | 15 +- .../examples/yolo/classify/run_train.py | 13 +- .../examples/yolo/classify/run_val.py | 16 +- .../examples/yolo/detect/inference.py | 16 +- .../examples/yolo/detect/run_train.py | 11 +- .../examples/yolo/detect/run_val.py | 17 +- .../examples/yolo/pose/inference.py | 14 +- .../examples/yolo/pose/run_train.py | 13 +- .../ultralytics/examples/yolo/pose/run_val.py | 16 +- .../examples/yolo/segment/inference.py | 14 +- .../examples/yolo/segment/run_train.py | 11 +- .../examples/yolo/segment/run_val.py | 17 +- .../{models => }/ultralytics/modeling_yolo.py | 26 +-- src/mindnlp/ultralytics/models/__init__.py | 1 + src/mindnlp/ultralytics/models/model.py | 201 ++++++++++++++++++ .../models/yolo/classify}/__init__.py | 0 .../models/yolo/classify/predict.py | 2 +- .../ultralytics/models/yolo/classify/train.py | 14 +- .../ultralytics/models/yolo/classify/val.py | 6 +- .../models/yolo/detect}/__init__.py | 0 .../ultralytics/models/yolo/detect/predict.py | 4 +- .../ultralytics/models/yolo/detect/train.py | 12 +- .../ultralytics/models/yolo/detect/val.py | 8 +- .../models/yolo/pose}/__init__.py | 0 .../ultralytics/models/yolo/pose/predict.py | 5 +- .../ultralytics/models/yolo/pose/train.py | 10 +- .../ultralytics/models/yolo/pose/val.py | 6 +- .../models/yolo/segment}/__init__.py | 0 .../models/yolo/segment/predict.py | 4 +- .../ultralytics/models/yolo/segment/train.py | 10 +- .../ultralytics/models/yolo/segment/val.py | 6 +- .../{models => }/ultralytics/modules.py | 79 ++++--- .../{models => }/ultralytics/readme.md | 57 +++-- src/mindnlp/ultralytics/requirements.txt | 26 +++ src/mindnlp/ultralytics/standalone.py | 64 ++++++ .../{models => }/ultralytics/tools/convert.py | 119 ++++++----- .../pose => ultralytics/utils}/__init__.py | 0 .../{models => }/ultralytics/utils/ema.py | 0 .../{models => }/ultralytics/utils/loss.py | 4 +- .../{models => }/ultralytics/utils/metrics.py | 0 .../{models => }/ultralytics/utils/ops.py | 0 .../ultralytics/utils/optimizer.py | 0 .../{models => }/ultralytics/utils/tal.py | 2 +- 60 files changed, 685 insertions(+), 524 deletions(-) delete mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml delete mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml delete mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml delete mode 100644 src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml delete mode 100644 src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py delete mode 100644 src/mindnlp/models/ultralytics/utils/__init__.py create mode 100644 src/mindnlp/ultralytics/__init__.py rename src/mindnlp/{models => }/ultralytics/cfg/hyp.yaml (100%) rename src/mindnlp/{models => }/ultralytics/cfg/models/11/yolo11-cls.yaml (100%) rename src/mindnlp/{models => }/ultralytics/cfg/models/11/yolo11-pose.yaml (100%) rename src/mindnlp/{models => }/ultralytics/cfg/models/11/yolo11-seg.yaml (100%) rename src/mindnlp/{models => }/ultralytics/cfg/models/11/yolo11.yaml (100%) rename src/mindnlp/{models => }/ultralytics/configuration_yolo.py (77%) rename src/mindnlp/{models/ultralytics => ultralytics/engine}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/engine/predictor.py (71%) rename src/mindnlp/{models => }/ultralytics/engine/trainer.py (99%) rename src/mindnlp/{models => }/ultralytics/engine/validator.py (98%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/classify/inference.py (88%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/classify/run_train.py (73%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/classify/run_val.py (87%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/detect/inference.py (81%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/detect/run_train.py (86%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/detect/run_val.py (85%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/pose/inference.py (87%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/pose/run_train.py (79%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/pose/run_val.py (85%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/segment/inference.py (84%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/segment/run_train.py (85%) rename src/mindnlp/{models => }/ultralytics/examples/yolo/segment/run_val.py (87%) rename src/mindnlp/{models => }/ultralytics/modeling_yolo.py (90%) create mode 100644 src/mindnlp/ultralytics/models/__init__.py create mode 100644 src/mindnlp/ultralytics/models/model.py rename src/mindnlp/{models/ultralytics/engine => ultralytics/models/yolo/classify}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/models/yolo/classify/predict.py (98%) rename src/mindnlp/{models => }/ultralytics/models/yolo/classify/train.py (84%) rename src/mindnlp/{models => }/ultralytics/models/yolo/classify/val.py (94%) rename src/mindnlp/{models/ultralytics/models => ultralytics/models/yolo/detect}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/models/yolo/detect/predict.py (96%) rename src/mindnlp/{models => }/ultralytics/models/yolo/detect/train.py (92%) rename src/mindnlp/{models => }/ultralytics/models/yolo/detect/val.py (94%) rename src/mindnlp/{models/ultralytics/models/yolo/classify => ultralytics/models/yolo/pose}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/models/yolo/pose/predict.py (97%) rename src/mindnlp/{models => }/ultralytics/models/yolo/pose/train.py (94%) rename src/mindnlp/{models => }/ultralytics/models/yolo/pose/val.py (97%) rename src/mindnlp/{models/ultralytics/models/yolo/detect => ultralytics/models/yolo/segment}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/models/yolo/segment/predict.py (98%) rename src/mindnlp/{models => }/ultralytics/models/yolo/segment/train.py (91%) rename src/mindnlp/{models => }/ultralytics/models/yolo/segment/val.py (97%) rename src/mindnlp/{models => }/ultralytics/modules.py (89%) rename src/mindnlp/{models => }/ultralytics/readme.md (68%) create mode 100644 src/mindnlp/ultralytics/requirements.txt create mode 100644 src/mindnlp/ultralytics/standalone.py rename src/mindnlp/{models => }/ultralytics/tools/convert.py (57%) rename src/mindnlp/{models/ultralytics/models/yolo/pose => ultralytics/utils}/__init__.py (100%) rename src/mindnlp/{models => }/ultralytics/utils/ema.py (100%) rename src/mindnlp/{models => }/ultralytics/utils/loss.py (99%) rename src/mindnlp/{models => }/ultralytics/utils/metrics.py (100%) rename src/mindnlp/{models => }/ultralytics/utils/ops.py (100%) rename src/mindnlp/{models => }/ultralytics/utils/optimizer.py (100%) rename src/mindnlp/{models => }/ultralytics/utils/tal.py (99%) diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml deleted file mode 100644 index a496fdb1e..000000000 --- a/src/mindnlp/models/ultralytics/cfg/datasets/coco128-seg.yaml +++ /dev/null @@ -1,101 +0,0 @@ -# Ultralytics AGPL-3.0 License - https://ultralytics.com/license - -# COCO128-seg dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics -# Documentation: https://docs.ultralytics.com/datasets/segment/coco/ -# Example usage: yolo train data=coco128-seg.yaml -# parent -# ├── ultralytics -# └── datasets -# └── coco128-seg ← downloads here (7 MB) - -# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] -path: /root/autodl-tmp/ultralytics/datasets/coco128-seg # dataset root dir -train: /root/autodl-tmp/ultralytics/datasets/coco128-seg/images/train2017 # train images (relative to 'path') 128 images -val: /root/autodl-tmp/ultralytics/datasets/coco128-seg/images/train2017 # val images (relative to 'path') 128 images -test: # test images (optional) - -# Classes -names: - 0: person - 1: bicycle - 2: car - 3: motorcycle - 4: airplane - 5: bus - 6: train - 7: truck - 8: boat - 9: traffic light - 10: fire hydrant - 11: stop sign - 12: parking meter - 13: bench - 14: bird - 15: cat - 16: dog - 17: horse - 18: sheep - 19: cow - 20: elephant - 21: bear - 22: zebra - 23: giraffe - 24: backpack - 25: umbrella - 26: handbag - 27: tie - 28: suitcase - 29: frisbee - 30: skis - 31: snowboard - 32: sports ball - 33: kite - 34: baseball bat - 35: baseball glove - 36: skateboard - 37: surfboard - 38: tennis racket - 39: bottle - 40: wine glass - 41: cup - 42: fork - 43: knife - 44: spoon - 45: bowl - 46: banana - 47: apple - 48: sandwich - 49: orange - 50: broccoli - 51: carrot - 52: hot dog - 53: pizza - 54: donut - 55: cake - 56: chair - 57: couch - 58: potted plant - 59: bed - 60: dining table - 61: toilet - 62: tv - 63: laptop - 64: mouse - 65: remote - 66: keyboard - 67: cell phone - 68: microwave - 69: oven - 70: toaster - 71: sink - 72: refrigerator - 73: book - 74: clock - 75: vase - 76: scissors - 77: teddy bear - 78: hair drier - 79: toothbrush - -# Download script/URL (optional) -download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128-seg.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml deleted file mode 100644 index 3db40eff3..000000000 --- a/src/mindnlp/models/ultralytics/cfg/datasets/coco128.yaml +++ /dev/null @@ -1,101 +0,0 @@ -# Ultralytics AGPL-3.0 License - https://ultralytics.com/license - -# COCO128 dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics -# Documentation: https://docs.ultralytics.com/datasets/detect/coco/ -# Example usage: yolo train data=coco128.yaml -# parent -# ├── ultralytics -# └── datasets -# └── coco128 ← downloads here (7 MB) - -# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] -path: /root/autodl-tmp/ultralytics/datasets/coco128/coco128 # dataset root dir -train: images/train2017 # train images (relative to 'path') 128 images -val: images/train2017 # val images (relative to 'path') 128 images -test: # test images (optional) - -# Classes -names: - 0: person - 1: bicycle - 2: car - 3: motorcycle - 4: airplane - 5: bus - 6: train - 7: truck - 8: boat - 9: traffic light - 10: fire hydrant - 11: stop sign - 12: parking meter - 13: bench - 14: bird - 15: cat - 16: dog - 17: horse - 18: sheep - 19: cow - 20: elephant - 21: bear - 22: zebra - 23: giraffe - 24: backpack - 25: umbrella - 26: handbag - 27: tie - 28: suitcase - 29: frisbee - 30: skis - 31: snowboard - 32: sports ball - 33: kite - 34: baseball bat - 35: baseball glove - 36: skateboard - 37: surfboard - 38: tennis racket - 39: bottle - 40: wine glass - 41: cup - 42: fork - 43: knife - 44: spoon - 45: bowl - 46: banana - 47: apple - 48: sandwich - 49: orange - 50: broccoli - 51: carrot - 52: hot dog - 53: pizza - 54: donut - 55: cake - 56: chair - 57: couch - 58: potted plant - 59: bed - 60: dining table - 61: toilet - 62: tv - 63: laptop - 64: mouse - 65: remote - 66: keyboard - 67: cell phone - 68: microwave - 69: oven - 70: toaster - 71: sink - 72: refrigerator - 73: book - 74: clock - 75: vase - 76: scissors - 77: teddy bear - 78: hair drier - 79: toothbrush - -# Download script/URL (optional) -download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml deleted file mode 100644 index 185e366ad..000000000 --- a/src/mindnlp/models/ultralytics/cfg/datasets/coco8-pose.yaml +++ /dev/null @@ -1,47 +0,0 @@ -# Ultralytics AGPL-3.0 License - https://ultralytics.com/license - -# COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics -# Documentation: https://docs.ultralytics.com/datasets/pose/coco8-pose/ -# Example usage: yolo train data=coco8-pose.yaml -# parent -# ├── ultralytics -# └── datasets -# └── coco8-pose ← downloads here (1 MB) - -# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] -path: ./datasets/coco8-pose # dataset root dir -train: images/train # train images (relative to 'path') 4 images -val: images/val # val images (relative to 'path') 4 images -test: # test images (optional) - -# Keypoints -kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) -flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] - -# Classes -names: - 0: person - -# Keypoint names per class -kpt_names: - 0: - - nose - - left_eye - - right_eye - - left_ear - - right_ear - - left_shoulder - - right_shoulder - - left_elbow - - right_elbow - - left_wrist - - right_wrist - - left_hip - - right_hip - - left_knee - - right_knee - - left_ankle - - right_ankle - -# Download script/URL (optional) -download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-pose.zip diff --git a/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml b/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml deleted file mode 100644 index 3d699147a..000000000 --- a/src/mindnlp/models/ultralytics/cfg/datasets/imagenette2-160.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# Ultralytics AGPL-3.0 License - https://ultralytics.com/license -# Imagenette2-160 Dataset by fast.ai -# Example usage: python train.py --data imagenette2-160.yaml - -# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] -path: ./datasets/imagenette2-160 # 你的数据集根目录 (根据你的截图和运行路径调整) -train: train # 训练集相对路径 (会与 path 拼接) -val: val # 验证集相对路径 (会与 path 拼接) -test: # 测试集 (Imagenette 没有单独的 test，可以留空) - -# Classes -nc: 10 # 类别数量 -names: - 0: tench # 丁鱼 - 1: English springer # 英国斯宾格犬 - 2: cassette player # 录音机 - 3: chain saw # 链锯 - 4: church # 教堂 - 5: French horn # 圆号 - 6: garbage truck # 垃圾车 - 7: gas pump # 加油泵 - 8: golf ball # 高尔夫球 - 9: parachute # 降落伞 - -# Download script/URL (optional) -download: | - wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz - tar -xzvf imagenette2-160.tgz \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py b/src/mindnlp/models/ultralytics/models/yolo/segment/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/mindnlp/models/ultralytics/utils/__init__.py b/src/mindnlp/models/ultralytics/utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/mindnlp/ultralytics/__init__.py b/src/mindnlp/ultralytics/__init__.py new file mode 100644 index 000000000..25da3acbe --- /dev/null +++ b/src/mindnlp/ultralytics/__init__.py @@ -0,0 +1,3 @@ +from .models.model import YOLO + +__all__ = ["YOLO"] \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/cfg/hyp.yaml b/src/mindnlp/ultralytics/cfg/hyp.yaml similarity index 100% rename from src/mindnlp/models/ultralytics/cfg/hyp.yaml rename to src/mindnlp/ultralytics/cfg/hyp.yaml diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml b/src/mindnlp/ultralytics/cfg/models/11/yolo11-cls.yaml similarity index 100% rename from src/mindnlp/models/ultralytics/cfg/models/11/yolo11-cls.yaml rename to src/mindnlp/ultralytics/cfg/models/11/yolo11-cls.yaml diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml b/src/mindnlp/ultralytics/cfg/models/11/yolo11-pose.yaml similarity index 100% rename from src/mindnlp/models/ultralytics/cfg/models/11/yolo11-pose.yaml rename to src/mindnlp/ultralytics/cfg/models/11/yolo11-pose.yaml diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml b/src/mindnlp/ultralytics/cfg/models/11/yolo11-seg.yaml similarity index 100% rename from src/mindnlp/models/ultralytics/cfg/models/11/yolo11-seg.yaml rename to src/mindnlp/ultralytics/cfg/models/11/yolo11-seg.yaml diff --git a/src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml b/src/mindnlp/ultralytics/cfg/models/11/yolo11.yaml similarity index 100% rename from src/mindnlp/models/ultralytics/cfg/models/11/yolo11.yaml rename to src/mindnlp/ultralytics/cfg/models/11/yolo11.yaml diff --git a/src/mindnlp/models/ultralytics/configuration_yolo.py b/src/mindnlp/ultralytics/configuration_yolo.py similarity index 77% rename from src/mindnlp/models/ultralytics/configuration_yolo.py rename to src/mindnlp/ultralytics/configuration_yolo.py index f8cf92f29..1168e9da1 100644 --- a/src/mindnlp/models/ultralytics/configuration_yolo.py +++ b/src/mindnlp/ultralytics/configuration_yolo.py @@ -1,26 +1,7 @@ import os import yaml -#from mindnlp.models.utils import PretrainedConfig -# 尝试获取 MindNLP 基类，增强套件兼容性 - -try: - - from mindnlp.models.utils import PretrainedConfig - -except Exception: - - # 兼容性防御：如果环境未完全安装，使用基础 Config 类防止崩溃 - - class PretrainedConfig: - - def __init__(self, **kwargs): - - for k, v in kwargs.items(): - - setattr(self, k, v) - -class YOLOConfig(PretrainedConfig): +class YOLOConfig: """ YOLO11 全任务通用配置类 支持：分类 (Classify)、检测 (Detect)、分割 (Segment)、姿态 (Pose) @@ -30,11 +11,11 @@ class YOLOConfig(PretrainedConfig): def __init__( self, yaml_path=None, - scale='n', - nc=80, + scale='n', + nc=80, kpt_shape=None, reg_max=16, - nm=32, + nm=32, npr=256, **kwargs ): @@ -80,4 +61,5 @@ def __init__( else: print(f"警告: YAML 中未找到 scale '{scale}'，将使用默认值 1.0") - super().__init__(**kwargs) \ No newline at end of file + for k, v in kwargs.items(): + setattr(self, k, v) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/__init__.py b/src/mindnlp/ultralytics/engine/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/__init__.py rename to src/mindnlp/ultralytics/engine/__init__.py diff --git a/src/mindnlp/models/ultralytics/engine/predictor.py b/src/mindnlp/ultralytics/engine/predictor.py similarity index 71% rename from src/mindnlp/models/ultralytics/engine/predictor.py rename to src/mindnlp/ultralytics/engine/predictor.py index 749296837..0d79cf8de 100644 --- a/src/mindnlp/models/ultralytics/engine/predictor.py +++ b/src/mindnlp/ultralytics/engine/predictor.py @@ -17,21 +17,29 @@ class BasePredictor: 负责端到端的推理：数据加载 -> 预处理 -> 前向推理 -> 后处理 """ - def __init__(self, cfg=None): + def __init__(self, cfg=None, **kwargs): """ 初始化预测器配置 """ - self.cfg = cfg or {} + if cfg is not None and not isinstance(cfg, dict): + if hasattr(cfg, '__dict__'): + cfg = vars(cfg) + else: + cfg = dict(cfg) + + self.cfg = cfg or kwargs or {} # 解析超参数配置 (支持从 hyp.yaml 读取 conf, iou 等推理参数) self.hyp = {} - if hasattr(self.cfg, 'hyp') and self.cfg.hyp and os.path.exists(self.cfg.hyp): - with open(self.cfg.hyp, "r", encoding="utf-8") as f: + hyp_path = self.cfg.get('hyp', None) + + if hyp_path and os.path.exists(str(hyp_path)): + with open(str(hyp_path), "r", encoding="utf-8") as f: self.hyp = yaml.safe_load(f) - self.conf_thres = self.hyp.get('conf', getattr(self.cfg, 'conf', 0.25)) - self.iou_thres = self.hyp.get('iou', getattr(self.cfg, 'iou', 0.45)) - self.imgsz = getattr(self.cfg, 'imgsz', 640) + self.conf_thres = self.hyp.get('conf', self.cfg.get('conf', 0.25)) + self.iou_thres = self.hyp.get('iou', self.cfg.get('iou', 0.45)) + self.imgsz = self.cfg.get('imgsz', 640) self.model = None self.dataset = [] @@ -113,7 +121,10 @@ def postprocess(self, preds, orig_img, preprocess_info): raise NotImplementedError("BasePredictor 不执行特定的解析逻辑，请在子类中重写 postprocess 方法。") def __call__(self, source, model=None, ckpt_path=None): - """推理流水线主调度入口""" + """ + 推理流水线主调度入口 + 集成了：环境设置 -> 数据加载 -> 循环推理 -> 自动保存 -> 结果返回 + """ if model is not None: self.setup_model(model, ckpt_path) @@ -123,26 +134,44 @@ def __call__(self, source, model=None, ckpt_path=None): self.setup_source(source) self.results = [] + save_dir = self.cfg.get('save_dir', 'runs/detect/predict') + should_save = self.cfg.get('save', True) + + if should_save: + os.makedirs(save_dir, exist_ok=True) + LOGGER.info(f" 推理结果将保存至: {os.path.abspath(save_dir)}") + LOGGER.info(f"推理引擎启动，共探测到 {len(self.dataset)} 份输入样本。") for img_path_or_arr in self.dataset: - # 1. 预处理 + # 预处理 t1 = time.time() im_tensor, orig_img, prep_info = self.preprocess(img_path_or_arr) - # 2. 推理 + # 推理 t2 = time.time() preds = self.inference(im_tensor) - # 3. 后处理 + # 后处理 (这里会返回各任务自定义的 Results 实例) t3 = time.time() result = self.postprocess(preds, orig_img, prep_info) + + if should_save and result is not None: + if isinstance(img_path_or_arr, str): + fname = os.path.basename(img_path_or_arr) + else: + fname = f"result_{int(time.time()*1000)}.jpg" + + if hasattr(result, 'save'): + result.save(save_dir=save_dir, file_name=fname) + result.save_dir = os.path.abspath(save_dir) + self.results.append(result) - # 4. 耗时统计 + # 耗时统计 inf_time = (t3 - t2) * 1000 post_time = (time.time() - t3) * 1000 name = img_path_or_arr if isinstance(img_path_or_arr, str) else "numpy_array" - LOGGER.info(f"处理完成 [{os.path.basename(name)}] | 前向推理: {inf_time:.1f}ms | 后处理: {post_time:.1f}ms") + LOGGER.info(f"处理完成 [{os.path.basename(name) if isinstance(name, str) else 'arr'}] | 前向推理: {inf_time:.1f}ms | 后处理: {post_time:.1f}ms") return self.results \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/engine/trainer.py b/src/mindnlp/ultralytics/engine/trainer.py similarity index 99% rename from src/mindnlp/models/ultralytics/engine/trainer.py rename to src/mindnlp/ultralytics/engine/trainer.py index 64beb6b3c..1cd02ec4d 100644 --- a/src/mindnlp/models/ultralytics/engine/trainer.py +++ b/src/mindnlp/ultralytics/engine/trainer.py @@ -4,7 +4,8 @@ import mindspore as ms import numpy as np from mindspore import nn, ops -from utils.ema import ModelEMA + +from ultralytics.utils.ema import ModelEMA class TrainStepWithClip(nn.TrainOneStepCell): """支持梯度裁剪的单步训练封装类""" diff --git a/src/mindnlp/models/ultralytics/engine/validator.py b/src/mindnlp/ultralytics/engine/validator.py similarity index 98% rename from src/mindnlp/models/ultralytics/engine/validator.py rename to src/mindnlp/ultralytics/engine/validator.py index 5872d767a..1d803116b 100644 --- a/src/mindnlp/models/ultralytics/engine/validator.py +++ b/src/mindnlp/ultralytics/engine/validator.py @@ -7,7 +7,7 @@ from mindspore import ops from tqdm import tqdm -from utils.ops import non_max_suppression +from ultralytics.utils.ops import non_max_suppression # 统一日志配置 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py b/src/mindnlp/ultralytics/examples/yolo/classify/inference.py similarity index 88% rename from src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py rename to src/mindnlp/ultralytics/examples/yolo/classify/inference.py index 84da1c62c..d1d1434e4 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/classify/inference.py +++ b/src/mindnlp/ultralytics/examples/yolo/classify/inference.py @@ -1,11 +1,20 @@ import argparse import os +import sys import logging import mindspore as ms -from models.yolo.classify.predict import ClassificationPredictor -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForClassification +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.classify.predict import ClassificationPredictor +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForClassification # 统一日志配置 logging.basicConfig(level=logging.INFO, format='%(message)s') diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py b/src/mindnlp/ultralytics/examples/yolo/classify/run_train.py similarity index 73% rename from src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py rename to src/mindnlp/ultralytics/examples/yolo/classify/run_train.py index 7614c7c5a..49e417b14 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_train.py +++ b/src/mindnlp/ultralytics/examples/yolo/classify/run_train.py @@ -3,15 +3,22 @@ import sys import os -# 从核心库中导入纯净的 Trainer -from models.yolo.classify.train import ClassificationTrainer +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.classify.train import ClassificationTrainer def main(): parser = argparse.ArgumentParser(description="YOLO11 Classification Training") parser.add_argument('--data', type=str, default='./cfg/datasets/imagenette2-160.yaml') parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-cls.yaml') - parser.add_argument('--weights', type=str, default='./yolo11n-cls.ckpt', help="初始化权重路径，留空则从头训练") + parser.add_argument('--weights', type=str, default='', help="初始化权重路径，留空则从头训练") parser.add_argument('--scale', type=str, default='n') parser.add_argument('--imgsz', type=int, default=224) parser.add_argument('--batch', type=int, default=64) diff --git a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py b/src/mindnlp/ultralytics/examples/yolo/classify/run_val.py similarity index 87% rename from src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py rename to src/mindnlp/ultralytics/examples/yolo/classify/run_val.py index 8222d9b76..d571fec5e 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/classify/run_val.py +++ b/src/mindnlp/ultralytics/examples/yolo/classify/run_val.py @@ -1,9 +1,19 @@ +import os +import sys import argparse import mindspore as ms -from models.yolo.classify.val import ClassificationValidator -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForClassification +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.classify.val import ClassificationValidator +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForClassification def main(): parser = argparse.ArgumentParser(description="YOLO11 Classification Validation Pipeline") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py b/src/mindnlp/ultralytics/examples/yolo/detect/inference.py similarity index 81% rename from src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py rename to src/mindnlp/ultralytics/examples/yolo/detect/inference.py index c3102b93b..4ccedb432 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/detect/inference.py +++ b/src/mindnlp/ultralytics/examples/yolo/detect/inference.py @@ -3,15 +3,23 @@ import sys import os -from models.yolo.detect.predict import DetectionPredictor -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForObjectDetection +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.detect.predict import DetectionPredictor +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForObjectDetection def main(): parser = argparse.ArgumentParser(description="YOLO11 Object Detection Inference Pipeline") # --- 核心推理参数 --- - parser.add_argument('--source', type=str, default='./datasets/coco128/coco128/images/train2017', help='待预测的图像或目录路径') + parser.add_argument('--source', type=str, default='./datasets/coco128/images/train2017', help='待预测的图像或目录路径') parser.add_argument('--weights', type=str, default='./yolo11n.ckpt', help='预训练权重文件路径') parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11.yaml', help='模型架构 YAML 配置文件') parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='模型规模标识') diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py b/src/mindnlp/ultralytics/examples/yolo/detect/run_train.py similarity index 86% rename from src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py rename to src/mindnlp/ultralytics/examples/yolo/detect/run_train.py index e5ed4f6a8..2d0f87323 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_train.py +++ b/src/mindnlp/ultralytics/examples/yolo/detect/run_train.py @@ -1,9 +1,18 @@ import argparse import os +import sys import yaml import mindspore as ms -from models.yolo.detect.train import DetectionTrainer +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.detect.train import DetectionTrainer def main(): parser = argparse.ArgumentParser(description="YOLO11 Object Detection Training Pipeline") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py b/src/mindnlp/ultralytics/examples/yolo/detect/run_val.py similarity index 85% rename from src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py rename to src/mindnlp/ultralytics/examples/yolo/detect/run_val.py index 9a806cc24..480c847b1 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/detect/run_val.py +++ b/src/mindnlp/ultralytics/examples/yolo/detect/run_val.py @@ -1,12 +1,21 @@ import argparse import mindspore as ms import os +import sys import yaml -from models.yolo.detect.val import DetectionValidator -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForObjectDetection -from data.loaders import create_dataloader +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.detect.val import DetectionValidator +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForObjectDetection +from ultralytics.data.loaders import create_dataloader def main(): parser = argparse.ArgumentParser(description="YOLO11 Object Detection Standalone Validation") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py b/src/mindnlp/ultralytics/examples/yolo/pose/inference.py similarity index 87% rename from src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py rename to src/mindnlp/ultralytics/examples/yolo/pose/inference.py index e331501d9..a9738a0fc 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/pose/inference.py +++ b/src/mindnlp/ultralytics/examples/yolo/pose/inference.py @@ -3,9 +3,17 @@ import sys import os -from models.yolo.pose.predict import PosePredictor -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForPose +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.pose.predict import PosePredictor +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForPose def main(): parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Inference Pipeline") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py b/src/mindnlp/ultralytics/examples/yolo/pose/run_train.py similarity index 79% rename from src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py rename to src/mindnlp/ultralytics/examples/yolo/pose/run_train.py index 81961cef4..ab9017488 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_train.py +++ b/src/mindnlp/ultralytics/examples/yolo/pose/run_train.py @@ -1,8 +1,17 @@ import argparse import os +import sys import mindspore as ms -from models.yolo.pose.train import PoseTrainer +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.pose.train import PoseTrainer def main(): parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Training Pipeline") @@ -12,7 +21,7 @@ def main(): parser.add_argument('--model_cfg', type=str, default='./cfg/models/11/yolo11-pose.yaml', help='模型架构 YAML 配置文件') parser.add_argument('--hyp', type=str, default='./cfg/hyp.yaml', help='训练超参数配置文件 (学习率、衰减等)') parser.add_argument('--scale', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'], help='网络规模标识') - parser.add_argument('--weights', type=str, default='./yolo11n-pose.ckpt', help='预训练权重文件路径 (.ckpt)') + parser.add_argument('--weights', type=str, default='', help='预训练权重文件路径 (.ckpt)') # 2. 训练周期与计算超参数 parser.add_argument('--epochs', type=int, default=100, help='总训练迭代轮数') diff --git a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py b/src/mindnlp/ultralytics/examples/yolo/pose/run_val.py similarity index 85% rename from src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py rename to src/mindnlp/ultralytics/examples/yolo/pose/run_val.py index 845b42806..e1ffaaf71 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/pose/run_val.py +++ b/src/mindnlp/ultralytics/examples/yolo/pose/run_val.py @@ -4,10 +4,18 @@ import sys import yaml -from models.yolo.pose.val import PoseValidator -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForPose -from data.loaders import create_dataloader +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.pose.val import PoseValidator +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForPose +from ultralytics.data.loaders import create_dataloader def main(): parser = argparse.ArgumentParser(description="YOLO11 Pose Estimation Standalone Validation") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py b/src/mindnlp/ultralytics/examples/yolo/segment/inference.py similarity index 84% rename from src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py rename to src/mindnlp/ultralytics/examples/yolo/segment/inference.py index d42415f24..59a82aedf 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/segment/inference.py +++ b/src/mindnlp/ultralytics/examples/yolo/segment/inference.py @@ -3,9 +3,17 @@ import sys import os -from models.yolo.segment.predict import SegmentationPredictor -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForSegmentation +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.segment.predict import SegmentationPredictor +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForSegmentation def main(): parser = argparse.ArgumentParser(description="YOLO11 Segmentation Inference Pipeline") diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py b/src/mindnlp/ultralytics/examples/yolo/segment/run_train.py similarity index 85% rename from src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py rename to src/mindnlp/ultralytics/examples/yolo/segment/run_train.py index b306999f4..f095cb54b 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_train.py +++ b/src/mindnlp/ultralytics/examples/yolo/segment/run_train.py @@ -1,9 +1,18 @@ import argparse import os +import sys import logging import mindspore as ms -from models.yolo.segment.train import SegmentationTrainer +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.segment.train import SegmentationTrainer logging.basicConfig(level=logging.INFO, format='%(message)s') LOGGER = logging.getLogger(__name__) diff --git a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py b/src/mindnlp/ultralytics/examples/yolo/segment/run_val.py similarity index 87% rename from src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py rename to src/mindnlp/ultralytics/examples/yolo/segment/run_val.py index 40bb368fc..99e27b1ab 100644 --- a/src/mindnlp/models/ultralytics/examples/yolo/segment/run_val.py +++ b/src/mindnlp/ultralytics/examples/yolo/segment/run_val.py @@ -1,12 +1,21 @@ import argparse import mindspore as ms import os +import sys import yaml -from models.yolo.segment.val import SegmentationValidator -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForSegmentation -from data.loaders import create_dataloader +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from ultralytics.models.yolo.segment.val import SegmentationValidator +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForSegmentation +from ultralytics.data.loaders import create_dataloader def main(): parser = argparse.ArgumentParser(description="YOLO11 Segmentation Standalone Validation") diff --git a/src/mindnlp/models/ultralytics/modeling_yolo.py b/src/mindnlp/ultralytics/modeling_yolo.py similarity index 90% rename from src/mindnlp/models/ultralytics/modeling_yolo.py rename to src/mindnlp/ultralytics/modeling_yolo.py index b5e4f1eb1..b8e549ae0 100644 --- a/src/mindnlp/models/ultralytics/modeling_yolo.py +++ b/src/mindnlp/ultralytics/modeling_yolo.py @@ -3,22 +3,8 @@ import mindspore as ms from mindspore import nn, ops -#from mindnlp.core.nn import PreTrainedModel - -try: - # 尝试导入官方的预训练基类 - from mindnlp.core.nn import PreTrainedModel -except ImportError as e: - print(f"⚠️ 警告: 无法从 mindnlp 导入 PreTrainedModel (可能是 mindtorch 版本冲突)。启动降级模式。错误: {e}") - # 本地降级方案：用一个假的基类糊弄过去，保证 train.py 能正常跑 - class PreTrainedModel(nn.Cell): - config_class = None - def __init__(self, config): - super().__init__() - self.config = config - -from configuration_yolo import YOLOConfig -from modules import ConvNormAct, C3k2, C2PSA, Classify, YOLO11DetectHead, YOLO11Segment, YOLO11Pose, Concat, Identity, SPPF, Upsample +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modules import ConvNormAct, C3k2, C2PSA, Classify, YOLO11DetectHead, YOLO11Segment, YOLO11Pose, Concat, Identity, SPPF, Upsample # 模块映射字典 @@ -36,16 +22,14 @@ def __init__(self, config): 'Classify': Classify } -class YOLO11Base(PreTrainedModel): +class YOLO11Base(nn.Cell): """ YOLO11 通用基类，负责解析 YAML 配置、搭建网络拓扑结构以及权重的初始化。 """ - # 绑定配置类，使得 from_pretrained 能够自动解析对应的 Config - config_class = YOLOConfig - def __init__(self, config): - super().__init__(config) + # nn.Cell 的初始化不需要传 config + super().__init__() self.config = config self.nc = config.nc diff --git a/src/mindnlp/ultralytics/models/__init__.py b/src/mindnlp/ultralytics/models/__init__.py new file mode 100644 index 000000000..5e1d8fee7 --- /dev/null +++ b/src/mindnlp/ultralytics/models/__init__.py @@ -0,0 +1 @@ +from .model import YOLO \ No newline at end of file diff --git a/src/mindnlp/ultralytics/models/model.py b/src/mindnlp/ultralytics/models/model.py new file mode 100644 index 000000000..cb05b7676 --- /dev/null +++ b/src/mindnlp/ultralytics/models/model.py @@ -0,0 +1,201 @@ +import os +import re +from types import SimpleNamespace +import contextlib +import numpy as np +import mindspore as ms +from mindspore.common.initializer import Initializer, Zero, initializer + + +from ultralytics.tools.convert import universal_convert + +# 导入各大任务的 Trainer (训练器) +from ultralytics.models.yolo.classify.train import ClassificationTrainer +from ultralytics.models.yolo.detect.train import DetectionTrainer +from ultralytics.models.yolo.segment.train import SegmentationTrainer +from ultralytics.models.yolo.pose.train import PoseTrainer + +# 导入各大任务的 Validator (验证器) +from ultralytics.models.yolo.classify.val import ClassificationValidator +from ultralytics.models.yolo.detect.val import DetectionValidator +from ultralytics.models.yolo.segment.val import SegmentationValidator +from ultralytics.models.yolo.pose.val import PoseValidator + +# 导入各大任务的 Predictor (推理器) +from ultralytics.models.yolo.classify.predict import ClassificationPredictor +from ultralytics.models.yolo.detect.predict import DetectionPredictor +from ultralytics.models.yolo.segment.predict import SegmentationPredictor +from ultralytics.models.yolo.pose.predict import PosePredictor + +class YOLO: + def __init__(self, model='yolo11n.pt', task=None): + self.model_name = str(model) + self.task = task + self.scale = 'n' + + self.model = None + self._parse_model_info() + + + if self.model_name.endswith('.yaml'): + print(f"[MindNLP YOLO] 检测到传入 YAML 架构文件: {self.model_name}") + print(f"[MindNLP YOLO] 模式: 从头开始随机初始化训练 (跳过权重转换)。") + self.yaml_path = self.model_name + self.ckpt_path = "" # 空路径表示不加载任何预训练权重 + else: + self.yaml_path = None + pt_path = self.model_name + + if self.model_name.endswith('.pt'): + self.ckpt_path = self.model_name.replace('.pt', '.ckpt') + pt_path = self.model_name + elif self.model_name.endswith('.ckpt') and not os.path.exists(self.model_name): + pt_path = self.model_name.replace('.ckpt', '.pt') + self.ckpt_path = self.model_name + else: + self.ckpt_path = self.model_name + + if self.ckpt_path and not os.path.exists(self.ckpt_path): + print(f"[MindNLP YOLO] 未找到本地权重 {self.ckpt_path}，准备获取源文件 {pt_path} 并转换...") + universal_convert(pt_path=pt_path, + ckpt_path=self.ckpt_path, + task=self.task, + scale=self.scale) + elif self.ckpt_path: + print(f"[MindNLP YOLO] 发现已存在权重: {self.ckpt_path}，直接加载。") + + def _parse_model_info(self): + """内部方法：从权重文件名中正则解析出任务类型和模型规模""" + # 解析任务 Task + if self.task is None: + if '-cls' in self.model_name: + self.task = 'classify' + elif '-seg' in self.model_name: + self.task = 'segment' + elif '-pose' in self.model_name: + self.task = 'pose' + else: + self.task = 'detect' # 默认无后缀为目标检测 + + # 解析规模 Scale (正则匹配 yolo11 后面的 n, s, m, l, x) + match = re.search(r'yolo11([nsmlx])', self.model_name.lower()) + if match: + self.scale = match.group(1) + else: + self.scale = 'n' # 若未匹配到，默认使用 nano 规模 + + def train(self, **kwargs): + """统一训练路由方法""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + parent_dir = os.path.dirname(current_dir) + + if getattr(self, 'yaml_path', None): + kwargs['model'] = self.yaml_path + kwargs['weights'] = "" # 无预训练权重 + else: + kwargs['model'] = self.ckpt_path + kwargs['weights'] = self.ckpt_path + + kwargs['scale'] = getattr(self, 'scale', 'n') + + # 3. 智能匹配任务架构文件 + if 'model_cfg' not in kwargs: + cfg_folder = os.path.join(parent_dir, 'cfg', 'models', '11') + task_cfg_map = { + 'detect': 'yolo11.yaml', + 'segment': 'yolo11-seg.yaml', + 'pose': 'yolo11-pose.yaml', + 'classify': 'yolo11-cls.yaml' + } + cfg_file = task_cfg_map.get(self.task, 'yolo11.yaml') + kwargs['model_cfg'] = os.path.join(cfg_folder, cfg_file) + + if not os.path.exists(kwargs['model_cfg']): + print(f"[MindNLP YOLO 错误] 找不到架构文件: {kwargs['model_cfg']}") + + # 4. 定位超参数文件 + if 'hyp' not in kwargs: + kwargs['hyp'] = os.path.join(parent_dir, 'cfg', 'hyp.yaml') + + # 5. 定位数据集配置 + if 'data' in kwargs and not os.path.isabs(kwargs['data']): + data_cfg_path = os.path.join(parent_dir, 'cfg', 'datasets', kwargs['data']) + if os.path.exists(data_cfg_path): + kwargs['data'] = data_cfg_path + print(f"[MindNLP YOLO] 自动定位数据集: {kwargs['data']}") + + # 6. 包装参数并路由 + if 'save_dir' not in kwargs: + kwargs['save_dir'] = os.path.join("runs", self.task, "train") + + args_obj = SimpleNamespace(**kwargs) + + print(f"[MindNLP YOLO] 准备启动 {self.task} 任务的训练...") + + if self.task == 'classify': + trainer = ClassificationTrainer(args=args_obj) + elif self.task == 'detect': + trainer = DetectionTrainer(args=args_obj) + elif self.task == 'segment': + trainer = SegmentationTrainer(args=args_obj) + elif self.task == 'pose': + trainer = PoseTrainer(args=args_obj) + else: + raise ValueError(f"[MindNLP YOLO] 暂不支持的任务类型: {self.task}") + + trainer.train() + if hasattr(trainer, 'model'): + self.model = trainer.model + elif hasattr(trainer, 'ema') and hasattr(trainer.ema, 'ema'): + self.model = trainer.ema.ema # 如果用了指数移动平均 + + return trainer + + def val(self, **kwargs): + """ + 统一验证路由方法。根据实例化的任务类型，拉起对应的验证器。 + """ + if self.model is None: + # 如果内存里没有模型实体，就让验证器去读本地权重文件 + kwargs['model'] = self.ckpt_path + pass_model = self.ckpt_path + print("[MindNLP YOLO] 当前内存无模型实体，将传递权重路径给 Validator 进行初始化。") + else: + # 如果刚刚 train 完，内存里有模型，直接传给验证器 + kwargs['model'] = self.ckpt_path # 依然保留路径作为元数据 + pass_model = self.model + print(f"[MindNLP YOLO] 准备启动 {self.task} 任务的验证...") + + if self.task == 'classify': + validator = ClassificationValidator(args=kwargs) + elif self.task == 'detect': + validator = DetectionValidator(args=kwargs) + elif self.task == 'segment': + validator = SegmentationValidator(args=kwargs) + elif self.task == 'pose': + validator = PoseValidator(args=kwargs) + else: + raise ValueError(f"[MindNLP YOLO] 暂不支持的任务类型: {self.task}") + + return validator(model=self.model) + + def __call__(self, source=None, **kwargs): + """ + 统一推理路由方法 + """ + kwargs['model'] = self.ckpt_path + kwargs['source'] = source + print(f"[MindNLP YOLO] 准备启动 {self.task} 任务的推理...") + + if self.task == 'classify': + predictor = ClassificationPredictor(cfg=kwargs) + elif self.task == 'detect': + predictor = DetectionPredictor(cfg=kwargs) + elif self.task == 'segment': + predictor = SegmentationPredictor(cfg=kwargs) + elif self.task == 'pose': + predictor = PosePredictor(cfg=kwargs) + else: + raise ValueError(f"[MindNLP YOLO] 暂不支持的任务类型: {self.task}") + + return predictor(source=source, model=self.model) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/engine/__init__.py b/src/mindnlp/ultralytics/models/yolo/classify/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/engine/__init__.py rename to src/mindnlp/ultralytics/models/yolo/classify/__init__.py diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/predict.py b/src/mindnlp/ultralytics/models/yolo/classify/predict.py similarity index 98% rename from src/mindnlp/models/ultralytics/models/yolo/classify/predict.py rename to src/mindnlp/ultralytics/models/yolo/classify/predict.py index 8a197cabb..7dc937d7f 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/classify/predict.py +++ b/src/mindnlp/ultralytics/models/yolo/classify/predict.py @@ -5,7 +5,7 @@ import mindspore as ms from mindspore import Tensor, ops -from engine.predictor import BasePredictor +from ultralytics.engine.predictor import BasePredictor # 分类结果实体类 class Results: diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/train.py b/src/mindnlp/ultralytics/models/yolo/classify/train.py similarity index 84% rename from src/mindnlp/models/ultralytics/models/yolo/classify/train.py rename to src/mindnlp/ultralytics/models/yolo/classify/train.py index 02b814538..63862fa7b 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/classify/train.py +++ b/src/mindnlp/ultralytics/models/yolo/classify/train.py @@ -6,13 +6,13 @@ from mindspore import nn # --- 导入底层架构和组件 --- -from engine.trainer import BaseTrainer -from models.yolo.classify.val import ClassificationValidator -from configuration_yolo import YOLOConfig -from modeling_yolo import YOLO11ForClassification -from data.loaders import create_dataloader -from utils.loss import YOLOClassificationLoss -from utils.optimizer import build_optimizer, get_lr +from ultralytics.engine.trainer import BaseTrainer +from ultralytics.models.yolo.classify.val import ClassificationValidator +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.modeling_yolo import YOLO11ForClassification +from ultralytics.data.loaders import create_dataloader +from ultralytics.utils.loss import YOLOClassificationLoss +from ultralytics.utils.optimizer import build_optimizer, get_lr class ClassificationTrainer(BaseTrainer): """图像分类任务专属 Trainer，实现基类的抽象方法""" diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/val.py b/src/mindnlp/ultralytics/models/yolo/classify/val.py similarity index 94% rename from src/mindnlp/models/ultralytics/models/yolo/classify/val.py rename to src/mindnlp/ultralytics/models/yolo/classify/val.py index 159a030a2..944d2a034 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/classify/val.py +++ b/src/mindnlp/ultralytics/models/yolo/classify/val.py @@ -4,9 +4,9 @@ import mindspore as ms from mindspore import ops -from engine.validator import BaseValidator -from data.loaders import create_dataloader -from utils.metrics import ClassifyMetrics +from ultralytics.engine.validator import BaseValidator +from ultralytics.data.loaders import create_dataloader +from ultralytics.utils.metrics import ClassifyMetrics LOGGER = logging.getLogger(__name__) diff --git a/src/mindnlp/models/ultralytics/models/__init__.py b/src/mindnlp/ultralytics/models/yolo/detect/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/models/__init__.py rename to src/mindnlp/ultralytics/models/yolo/detect/__init__.py diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/predict.py b/src/mindnlp/ultralytics/models/yolo/detect/predict.py similarity index 96% rename from src/mindnlp/models/ultralytics/models/yolo/detect/predict.py rename to src/mindnlp/ultralytics/models/yolo/detect/predict.py index 5d66e73f6..ed0a17503 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/detect/predict.py +++ b/src/mindnlp/ultralytics/models/yolo/detect/predict.py @@ -3,8 +3,8 @@ import cv2 import mindspore as ms -from engine.predictor import BasePredictor -from utils.ops import non_max_suppression +from ultralytics.engine.predictor import BasePredictor +from ultralytics.utils.ops import non_max_suppression # 结果封装与可视化类 class Results: diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/train.py b/src/mindnlp/ultralytics/models/yolo/detect/train.py similarity index 92% rename from src/mindnlp/models/ultralytics/models/yolo/detect/train.py rename to src/mindnlp/ultralytics/models/yolo/detect/train.py index e961b436b..854fb2f92 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/detect/train.py +++ b/src/mindnlp/ultralytics/models/yolo/detect/train.py @@ -3,12 +3,12 @@ import logging import mindspore as ms -from modeling_yolo import YOLO11ForObjectDetection -from configuration_yolo import YOLOConfig -from data.loaders import create_dataloader -from engine.trainer import BaseTrainer -from utils.optimizer import build_optimizer, get_lr -from utils.loss import v8DetectionLoss +from ultralytics.modeling_yolo import YOLO11ForObjectDetection +from ultralytics.configuration_yolo import YOLOConfig +from ultralytics.data.loaders import create_dataloader +from ultralytics.engine.trainer import BaseTrainer +from ultralytics.utils.optimizer import build_optimizer, get_lr +from ultralytics.utils.loss import v8DetectionLoss # 统一日志配置 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/val.py b/src/mindnlp/ultralytics/models/yolo/detect/val.py similarity index 94% rename from src/mindnlp/models/ultralytics/models/yolo/detect/val.py rename to src/mindnlp/ultralytics/models/yolo/detect/val.py index bd186339f..08ca4d3f3 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/detect/val.py +++ b/src/mindnlp/ultralytics/models/yolo/detect/val.py @@ -3,10 +3,10 @@ import mindspore as ms from mindspore import ops -from engine.validator import BaseValidator -from utils.metrics import DetMetrics, process_batch -from utils.ops import non_max_suppression, xywh2xyxy_np -from data.loaders import create_dataloader +from ultralytics.engine.validator import BaseValidator +from ultralytics.utils.metrics import DetMetrics, process_batch +from ultralytics.utils.ops import non_max_suppression, xywh2xyxy_np +from ultralytics.data.loaders import create_dataloader LOGGER = logging.getLogger(__name__) diff --git a/src/mindnlp/models/ultralytics/models/yolo/classify/__init__.py b/src/mindnlp/ultralytics/models/yolo/pose/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/models/yolo/classify/__init__.py rename to src/mindnlp/ultralytics/models/yolo/pose/__init__.py diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/predict.py b/src/mindnlp/ultralytics/models/yolo/pose/predict.py similarity index 97% rename from src/mindnlp/models/ultralytics/models/yolo/pose/predict.py rename to src/mindnlp/ultralytics/models/yolo/pose/predict.py index 6ee1717db..e78168781 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/pose/predict.py +++ b/src/mindnlp/ultralytics/models/yolo/pose/predict.py @@ -2,8 +2,9 @@ import numpy as np import cv2 import mindspore as ms -from engine.predictor import BasePredictor -from utils.ops import non_max_suppression + +from ultralytics.engine.predictor import BasePredictor +from ultralytics.utils.ops import non_max_suppression # 预测结果封装与渲染类 class Results: diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/train.py b/src/mindnlp/ultralytics/models/yolo/pose/train.py similarity index 94% rename from src/mindnlp/models/ultralytics/models/yolo/pose/train.py rename to src/mindnlp/ultralytics/models/yolo/pose/train.py index db6fafd3f..3ea3ee042 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/pose/train.py +++ b/src/mindnlp/ultralytics/models/yolo/pose/train.py @@ -3,11 +3,11 @@ import mindspore as ms from mindspore import nn, ops -from engine.trainer import BaseTrainer -from modeling_yolo import YOLO11ForPose, YOLOConfig -from utils.loss import v8PoseLoss -from utils.optimizer import build_optimizer, get_lr -from data.loaders import create_dataloader +from ultralytics.engine.trainer import BaseTrainer +from ultralytics.modeling_yolo import YOLO11ForPose, YOLOConfig +from ultralytics.utils.loss import v8PoseLoss +from ultralytics.utils.optimizer import build_optimizer, get_lr +from ultralytics.data.loaders import create_dataloader class PoseTrainer(BaseTrainer): """ diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/val.py b/src/mindnlp/ultralytics/models/yolo/pose/val.py similarity index 97% rename from src/mindnlp/models/ultralytics/models/yolo/pose/val.py rename to src/mindnlp/ultralytics/models/yolo/pose/val.py index 7ec6d091a..5da11385a 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/pose/val.py +++ b/src/mindnlp/ultralytics/models/yolo/pose/val.py @@ -2,9 +2,9 @@ import mindspore as ms from mindspore import ops -from engine.validator import BaseValidator -from utils.ops import non_max_suppression, xywh2xyxy_np -from utils.metrics import PoseMetrics, kpt_iou +from ultralytics.engine.validator import BaseValidator +from ultralytics.utils.ops import non_max_suppression, xywh2xyxy_np +from ultralytics.utils.metrics import PoseMetrics, kpt_iou class PoseValidator(BaseValidator): """ diff --git a/src/mindnlp/models/ultralytics/models/yolo/detect/__init__.py b/src/mindnlp/ultralytics/models/yolo/segment/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/models/yolo/detect/__init__.py rename to src/mindnlp/ultralytics/models/yolo/segment/__init__.py diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/predict.py b/src/mindnlp/ultralytics/models/yolo/segment/predict.py similarity index 98% rename from src/mindnlp/models/ultralytics/models/yolo/segment/predict.py rename to src/mindnlp/ultralytics/models/yolo/segment/predict.py index 5cffc6cce..a4d3ca00c 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/segment/predict.py +++ b/src/mindnlp/ultralytics/models/yolo/segment/predict.py @@ -3,8 +3,8 @@ import cv2 import mindspore as ms -from engine.predictor import BasePredictor -from utils.ops import non_max_suppression +from ultralytics.engine.predictor import BasePredictor +from ultralytics.utils.ops import non_max_suppression # 实例分割结果实体与可视化类 class Results: diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/train.py b/src/mindnlp/ultralytics/models/yolo/segment/train.py similarity index 91% rename from src/mindnlp/models/ultralytics/models/yolo/segment/train.py rename to src/mindnlp/ultralytics/models/yolo/segment/train.py index 8a99b8f49..068813be1 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/segment/train.py +++ b/src/mindnlp/ultralytics/models/yolo/segment/train.py @@ -4,11 +4,11 @@ import mindspore as ms # --- 核心架构组件 --- -from engine.trainer import BaseTrainer -from modeling_yolo import YOLO11ForSegmentation, YOLOConfig -from utils.loss import v8SegmentationLoss -from utils.optimizer import build_optimizer, get_lr -from data.loaders import create_dataloader +from ultralytics.engine.trainer import BaseTrainer +from ultralytics.modeling_yolo import YOLO11ForSegmentation, YOLOConfig +from ultralytics.utils.loss import v8SegmentationLoss +from ultralytics.utils.optimizer import build_optimizer, get_lr +from ultralytics.data.loaders import create_dataloader LOGGER = logging.getLogger(__name__) diff --git a/src/mindnlp/models/ultralytics/models/yolo/segment/val.py b/src/mindnlp/ultralytics/models/yolo/segment/val.py similarity index 97% rename from src/mindnlp/models/ultralytics/models/yolo/segment/val.py rename to src/mindnlp/ultralytics/models/yolo/segment/val.py index 79b1cbf40..3e3f83924 100644 --- a/src/mindnlp/models/ultralytics/models/yolo/segment/val.py +++ b/src/mindnlp/ultralytics/models/yolo/segment/val.py @@ -4,9 +4,9 @@ import mindspore as ms from mindspore import Tensor, ops -from engine.validator import BaseValidator -from utils.ops import non_max_suppression, process_mask, xywh2xyxy_np -from utils.metrics import SegmentMetrics +from ultralytics.engine.validator import BaseValidator +from ultralytics.utils.ops import non_max_suppression, process_mask, xywh2xyxy_np +from ultralytics.utils.metrics import SegmentMetrics LOGGER = logging.getLogger(__name__) diff --git a/src/mindnlp/models/ultralytics/modules.py b/src/mindnlp/ultralytics/modules.py similarity index 89% rename from src/mindnlp/models/ultralytics/modules.py rename to src/mindnlp/ultralytics/modules.py index cce68b49b..50d4ed05f 100644 --- a/src/mindnlp/models/ultralytics/modules.py +++ b/src/mindnlp/ultralytics/modules.py @@ -4,8 +4,8 @@ import mindspore.numpy as mnp from mindspore import nn, ops, Tensor, Parameter -from utils.ops import dist2bbox -from utils.tal import make_anchors +from ultralytics.utils.ops import dist2bbox +from ultralytics.utils.tal import make_anchors # 基础工具函数与算子 @@ -284,38 +284,44 @@ def construct(self, x): 训练模式：返回 (拼接后的展平Tensor, 原始特征图List) 推理模式：返回 (最终预测结果, 原始特征图List) """ - # 1. 基础卷积处理 + # 基础卷积处理 res = [] for i in range(self.nl): # 保持 [B, 64+nc, H, W] 结构 res.append(ops.concat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)) - # 2. 统一展平逻辑 (训练和推理都需要这个展平后的 x_cat 来算 Loss) + # 统一展平逻辑 y = [] for xi in res: bs, ch, h, w = xi.shape y.append(xi.view(bs, ch, -1)) - # 拼接所有尺度：[B, 144, 19200] + # 拼接所有尺度，得到形状: [B, 144, 19200] (以 640x640 输入为例) x_cat = ops.concat(y, 2) if self.training: return x_cat, res - # 3. 推理解码逻辑 (仅在 self.training=False 时执行) + # 将预测值拆分为 框回归参数(box_dist) 和 类别概率(cls_logits) box_dist, cls_logits = ops.split(x_cat, (self.reg_max * 4, self.nc), axis=1) - # DFL 积分回归 - box_decoded = self.dfl(box_dist) # 形状 [B, 4, 19200] + # DFL 积分回归 -> 得到形状 [B, 4, 19200] + box_decoded = self.dfl(box_dist) - # 获取锚点和步长 - # 这里的 x 必须是原始特征图列表 - anchors, strides = self.make_anchors(x, self.stride, 0.5) + # 获取原始锚点和步长 + anchors, strides = make_anchors(x, self.stride, 0.5) - # 解码坐标：直接传入原始 anchors，利用 ops.py 里的自适应逻辑 - dbox = dist2bbox(box_decoded, anchors, strides=strides, xywh=True, axis=1) + anchors_aligned = anchors.transpose(1, 0).expand_dims(0) + + # 同样旋转并增加维度，使其成为横向的步长向量 + strides_aligned = strides.transpose(1, 0).expand_dims(0) + + # 解码坐标：使用显式对齐后的张量进行计算 + # dbox 形状将保持为 [B, 4, 19200] + dbox = dist2bbox(box_decoded, anchors_aligned, strides=strides_aligned, xywh=True, axis=1) # 最终拼接：[B, 4+nc, 19200] + # 注意对类别概率进行 sigmoid 处理 final_pred = ops.concat((dbox, ops.sigmoid(cls_logits)), 1) return final_pred, res @@ -371,32 +377,21 @@ def construct(self, x): out = tuple(ops.concat((box_outs[i], cls_outs[i]), 1) for i in range(self.nl)) return out, kpt_outs - box_flat = ops.concat([b.view(bs, self.reg_max * 4, -1) for b in box_outs], axis=2) - cls_flat = ops.concat([c.view(bs, self.nc, -1) for c in cls_outs], axis=2) - kpt_flat = ops.concat([k.view(bs, self.nkpt, -1) for k in kpt_outs], axis=2) - - # 统一提升精度以保证解码阶段数值稳定性 - box_flat = box_flat.astype(ms.float32) - cls_flat = cls_flat.astype(ms.float32) - kpt_flat = kpt_flat.astype(ms.float32) - - cls_prob = ops.sigmoid(cls_flat) + box_flat = ops.concat([b.view(bs, self.reg_max * 4, -1) for b in box_outs], axis=2).astype(ms.float32) + cls_flat = ops.concat([c.view(bs, self.nc, -1) for c in cls_outs], axis=2).astype(ms.float32) + kpt_flat = ops.concat([k.view(bs, self.nkpt, -1) for k in kpt_outs], axis=2).astype(ms.float32) - self.anchors, self.strides = make_anchors(x, self.stride, 0.5) - anc = self.anchors.astype(ms.float32).view(1, -1, 2).transpose(0, 2, 1) - strd = self.strides.astype(ms.float32).view(1, -1, 1).transpose(0, 2, 1) + anchors, strides = make_anchors(x, self.stride, 0.5) + anc_aligned = anchors.transpose(1, 0).expand_dims(0) + strd_aligned = strides.transpose(1, 0).expand_dims(0) box_dist = self.dfl(box_flat) - lt = box_dist[:, :2, :] - rb = box_dist[:, 2:, :] - c_xy = anc + (rb - lt) / 2.0 - wh = rb + lt - dbox = ops.concat((c_xy, wh), axis=1) * strd + dbox = dist2bbox(box_dist, anc_aligned, strides=strd_aligned, xywh=True, axis=1) - pred_kpt = self.kpts_decode(kpt_flat, self.anchors, self.strides) + pred_kpt = self.kpts_decode(kpt_flat, anchors, strides) - # 按照 (Box, Cls, Pose) 顺序拼接 - y = ops.concat((dbox, cls_prob, pred_kpt), axis=1) + y = ops.concat((dbox, ops.sigmoid(cls_flat), pred_kpt), axis=1) + return y.transpose(0, 2, 1) class ProtoCell(nn.Cell): @@ -452,16 +447,16 @@ def construct(self, x): if self.training: return out, ops.concat(mc, 2), p_out - self.anchors, self.strides = self.make_anchors(out, self.stride, 0.5) - x_all = ops.concat([xi.view(bs, self.no, -1) for xi in out], 2) + anchors, strides = make_anchors(out, self.stride, 0.5) - box, cls = ops.split(x_all, (self.reg_max * 4, self.nc), 1) + anc_aligned = anchors.transpose(1, 0).expand_dims(0) + strd_aligned = strides.transpose(1, 0).expand_dims(0) - box_decoded = self.dfl(box) - box_decoded = ops.relu(box_decoded) - dbox = dist2bbox(box_decoded.transpose(0, 2, 1), - ops.expand_dims(self.anchors, 0), - self.strides, xywh=True) + x_all = ops.concat([xi.view(bs, self.no, -1) for xi in out], 2) + box, cls = ops.split(x_all, (self.reg_max * 4, self.nc), 1) + + box_decoded = self.dfl(box) # 得到 [B, 4, N] + dbox = dist2bbox(box_decoded, anc_aligned, strides=strd_aligned, xywh=True, axis=1) mc_all = ops.concat(mc, 2) final_pred = ops.concat((dbox, ops.sigmoid(cls), mc_all), 1) diff --git a/src/mindnlp/models/ultralytics/readme.md b/src/mindnlp/ultralytics/readme.md similarity index 68% rename from src/mindnlp/models/ultralytics/readme.md rename to src/mindnlp/ultralytics/readme.md index ec79f30b1..95762f7a5 100644 --- a/src/mindnlp/models/ultralytics/readme.md +++ b/src/mindnlp/ultralytics/readme.md @@ -1,32 +1,51 @@ -# mindnlp兼容ultralutics库项目运行指南 +# mindnlp兼容ultralytics库项目运行指南 本项目基于 MindSpore 框架实现了 YOLO11 的四大核心任务：图像分类、目标检测、实例分割和姿态估计。支持从头训练、加载预训练权重微调、模型验证与推理。 -1. 数据集准备 +```bash +# 创建环境 (Python 3.9) +conda create -n mindnlp_yolo python=3.10 -y +conda activate mindnlp_yolo + +# 配置 Ascend/CANN 环境（按实际安装路径调整） +source /usr/local/Ascend/ascend-toolkit/set_env.sh + +# 安装项目依赖 +pip install -r requirements.txt + +# 将本地开发的 MindNLP 挂载到环境中（版本为0.6.0） +# 请确保当前处于包含 setup.py 的根目录下 +pip install -e . +``` + + + +##1. 数据集准备 请参考 `./ultralytics/cfg/datasets` 目录下各个数据集的配置文件（`.yaml`）获取下载路径。 下载完成后，请将数据集解压并放置到以下目录： `./ultralytics/datasets/` -2. 预训练权重转换 +##2. 预训练权重转换 -我们提供了将 PyTorch 的 `.pt` 权重转化为 MindSpore 的 `.ckpt` 权重的转换脚本。请在项目根目录下运行以下命令： +我们提供了将 PyTorch 的 `.pt` 权重转化为 MindSpore 的 `.ckpt` 权重的转换脚本。执行完转换脚本后，生成的 .ckpt 文件默认就在根目录。请在项目根目录下运行以下命令： ```bash -# 转换分类任务权重 -python tools/convert.py --task classify +# 转换分类任务权重 +python tools/convert.py --task classify --scale n # 转换检测任务权重 -python tools/convert.py --task detect +python tools/convert.py --task detect --scale n # 转换分割任务权重 -python tools/convert.py --task segment +python tools/convert.py --task segment --scale n # 转换姿态估计任务权重 -python tools/convert.py --task pose +python tools/convert.py --task pose --scale n +``` -3. 任务运行指令 +##3. 任务运行指令 以下命令均需在项目根目录下执行。 （1）图像分类任务 (Classify) - +```bash # 加载权重微调 (Fine-tune) python examples/yolo/classify/run_train.py --weights ./yolo11n-cls.ckpt --save_dir ./runs/cls/train_finetune @@ -38,8 +57,9 @@ python examples/yolo/classify/run_val.py # 模型推理 (Inference) python examples/yolo/classify/inference.py +``` （2）目标检测任务 (Detect) - +```bash # 加载权重微调 (Fine-tune) python examples/yolo/detect/run_train.py --weights ./yolo11n.ckpt --save_dir ./runs/detect/train_finetune @@ -51,8 +71,9 @@ python examples/yolo/detect/run_val.py # 模型推理 (Inference) python examples/yolo/detect/inference.py +``` （3）实例分割任务 (Segment) - +```bash # 加载权重微调 (Fine-tune) python examples/yolo/segment/run_train.py --weights ./yolo11n-seg.ckpt --save_dir ./runs/segment/train_finetune @@ -64,8 +85,9 @@ python examples/yolo/segment/run_val.py # 模型推理 (Inference) python examples/yolo/segment/inference.py +``` （4）姿态估计任务 (Pose) - +```bash # 加载权重微调 (Fine-tune) python examples/yolo/pose/run_train.py --weights ./yolo11n-pose.ckpt --save_dir ./runs/pose/train_finetune @@ -76,4 +98,9 @@ python examples/yolo/pose/run_train.py --save_dir ./runs/pose/train_scratch python examples/yolo/pose/run_val.py # 模型推理 (Inference) -python examples/yolo/pose/inference.py \ No newline at end of file +python examples/yolo/pose/inference.py +``` +## 4. 用户级 API 调用示例 +```bash +python standalone.py +``` \ No newline at end of file diff --git a/src/mindnlp/ultralytics/requirements.txt b/src/mindnlp/ultralytics/requirements.txt new file mode 100644 index 000000000..7907b540f --- /dev/null +++ b/src/mindnlp/ultralytics/requirements.txt @@ -0,0 +1,26 @@ +mindspore==2.8.0 +numpy>=1.23.5,<2.0.0 +scipy +requests +pyyaml +attrs +jinja2 +ml-dtypes +packaging +importlib-metadata +decorator +cloudpickle +tornado +absl-py +accelerate + +opencv-python +pillow>=10.0.0 +tqdm +matplotlib>=3.3.0 + +transformers==4.44.2 + +#权重转化时需要 +torch>=2.0.0 +ultralytics>=8.3.0 \ No newline at end of file diff --git a/src/mindnlp/ultralytics/standalone.py b/src/mindnlp/ultralytics/standalone.py new file mode 100644 index 000000000..eb863dcb0 --- /dev/null +++ b/src/mindnlp/ultralytics/standalone.py @@ -0,0 +1,64 @@ +import os +import sys +import argparse +import logging +import gc +import mindspore as ms + +# 配置自定义日志 +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) + +LOGGER = logging.getLogger("YOLO11-Pipeline") + +current_dir = os.path.dirname(os.path.abspath(__file__)) +parent_dir = os.path.dirname(current_dir) +if parent_dir not in sys.path: + sys.path.insert(0, parent_dir) + +from ultralytics import YOLO + +if __name__ == "__main__": + + ms.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend") + + LOGGER.info(" 开始初始化 YOLO11 模型...") + + #微调 + model = YOLO("yolo11n.pt") + #从头开始训练 + #model = YOLO("yolo11.yaml") + + LOGGER.info(" 启动训练流程...") + + # 训练模型 + results = model.train( + data="coco128.yaml", + epochs=100, + imgsz=640, + batch=16, + amp=False, + val_interval=10, + ) + + # 打印最佳权重精度 + LOGGER.info(f" 训练完成！最高综合评价指标 (Fitness): {results.best_fitness:.4f}") + + # 内存清理 + gc.collect() + + # 执行推理预测 + source_img = "./datasets/coco128/images/train2017" + LOGGER.info(f" 开始推理预测: {source_img}") + + predict_results = model( + source=source_img, + imgsz=640, + conf=0.25, + iou=0.45, + save=True, + ) + + LOGGER.info(f" 推理完成！结果保存在: {predict_results[0].save_dir}") \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/tools/convert.py b/src/mindnlp/ultralytics/tools/convert.py similarity index 57% rename from src/mindnlp/models/ultralytics/tools/convert.py rename to src/mindnlp/ultralytics/tools/convert.py index 471ac13c5..3f70b932a 100644 --- a/src/mindnlp/models/ultralytics/tools/convert.py +++ b/src/mindnlp/ultralytics/tools/convert.py @@ -1,62 +1,61 @@ import os +import sys import argparse import numpy as np import mindspore as ms from mindspore import Tensor -from ultralytics import YOLO +import subprocess + +project_root = os.path.abspath(__file__) +while os.path.basename(project_root) != 'ultralytics' and project_root != '/': + project_root = os.path.dirname(project_root) +project_root = os.path.dirname(project_root) + +if project_root not in sys.path: + sys.path.insert(0, project_root) -import modeling_yolo -from configuration_yolo import YOLOConfig +from ultralytics import modeling_yolo +from ultralytics.configuration_yolo import YOLOConfig def translate_ms_to_pt(ms_name, task): """ 参数映射路由表：将 MindSpore 架构下的参数名称转换为 PyTorch 架构对应的参数名称 """ - # 基础规范化转换：对齐 BatchNorm 以及权重、偏置的命名差异 pt_name = ms_name.replace(".moving_mean", ".running_mean") \ .replace(".moving_variance", ".running_var") \ .replace(".gamma", ".weight") \ .replace(".beta", ".bias") - # 结构命名转换：对齐 C2f/Bottleneck 等内部核心模块的命名规范 pt_name = pt_name.replace(".conv1.", ".cv1.") pt_name = pt_name.replace(".conv2.", ".cv2.") pt_name = pt_name.replace(".conv3.", ".cv3.") return pt_name -def universal_convert(task="segment", scale="n"): +def universal_convert(pt_path, ckpt_path, task="detect", scale="n"): """ 通用权重转换流水线。 - 支持将 Ultralytics 官方 PyTorch (.pt) 权重转换为符合 MindNLP 规范的 MindSpore (.ckpt) 权重 + 支持将官方 PyTorch (.pt) 转换为符合 MindSpore (.ckpt) 的权重。 """ - # 1. 初始化任务及配置映射字典 model_map = { "classify": modeling_yolo.YOLO11ForClassification, "detect": modeling_yolo.YOLO11ForObjectDetection, "segment": modeling_yolo.YOLO11ForSegmentation, "pose": modeling_yolo.YOLO11ForPose } - pt_name_map = { - "classify": f"yolo11{scale}-cls.pt", - "detect": f"yolo11{scale}.pt", - "segment": f"yolo11{scale}-seg.pt", - "pose": f"yolo11{scale}-pose.pt" - } yaml_map = { "classify": "cfg/models/11/yolo11-cls.yaml", "detect": "cfg/models/11/yolo11.yaml", "segment": "cfg/models/11/yolo11-seg.yaml", "pose": "cfg/models/11/yolo11-pose.yaml" } - nc_map = {"classify": 1000, "detect": 80, "segment": 80, "pose": 1} print(f"[INFO] 启动 YOLO11-{task.upper()} 权重转换流程...") current_yaml = yaml_map[task] - # 2. 动态构建网络配置与 MindSpore 模型实例 + # 1. 动态构建网络配置与 MindSpore 模型实例 if task == "pose": cfg = YOLOConfig(yaml_path=current_yaml, scale=scale, nc=nc_map[task], kpt_shape=[17, 3]) else: @@ -65,39 +64,57 @@ def universal_convert(task="segment", scale="n"): ms_model = model_map[task](cfg) ms_model.set_train(False) - # 3. 加载 PyTorch 原生权重 - print(f"[INFO] 正在解析 PyTorch 权重文件: {pt_name_map[task]}") - pt_yolo = YOLO(pt_name_map[task]) - pt_dict = pt_yolo.model.state_dict() + # 通过子进程调用系统里的正版 PyTorch，提取纯 NumPy 数据 + npz_file = pt_path + ".npz" + if not os.path.exists(npz_file): + print(f"[INFO] 正在启动独立子进程，跨越物理隔离提取纯净权重...") + extract_script = f""" +import torch +from ultralytics import YOLO +import numpy as np + +# 此处是在干净的子进程中，载入的是官方原版环境 +print(" -> [子进程] 正在加载官方 PyTorch 模型: {pt_path}") +model = YOLO('{pt_path}') +state_dict = model.model.state_dict() + +# 剔除 PyTorch 依赖，将其彻底降维成通用 NumPy 格式 +print(" -> [子进程] 正在降维并持久化至 NumPy 数组...") +np_dict = {{k: v.cpu().numpy() for k, v in state_dict.items()}} +np.savez('{npz_file}', **np_dict) +""" + # 将子进程代码写入临时文件并执行 + with open("temp_extract.py", "w") as f: + f.write(extract_script.strip()) + + # 运行子进程提取数据 + subprocess.run(["python", "temp_extract.py"], check=True) + os.remove("temp_extract.py") + + # 2. 回到主进程：读取绝对纯净的 NumPy 字典，彻底断绝与 PyTorch 的瓜葛 + print(f"[INFO] 正在解析纯净权重容器: {npz_file}") + pt_dict = dict(np.load(npz_file)) - print("[DEBUG] 官方 PyTorch 权重键名抽样 (末尾 5 项):") - pt_keys = list(pt_dict.keys()) - for k in pt_keys[-5:]: - print(f" {k:50} | 尺寸: {list(pt_dict[k].shape)}") - print("-" * 80) - new_ms_ckpt = [] matched_count = 0 ms_params = list(ms_model.parameters_and_names()) print("[INFO] 开始执行参数映射与维度校验...") - # 4. 执行转换主循环 + # 3. 执行转换主循环 for ms_name, ms_param in ms_params: ms_shape = tuple(ms_param.shape) ms_size = ms_param.size - # 获取期望映射的 PyTorch 键名 expected_pt_name = translate_ms_to_pt(ms_name, task) if expected_pt_name in pt_dict: pt_v = pt_dict[expected_pt_name] - # 维度一致性校验 - if pt_v.numel() == ms_size: - val_np = pt_v.cpu().numpy().reshape(ms_shape) + # 注意：pt_v 现在已经是 NumPy 数组了，所以用 .size 而不是 numel() + if pt_v.size == ms_size: + val_np = pt_v.reshape(ms_shape) - # 安全策略：限制 BatchNorm 历史方差下限，防止半精度推理时出现除零溢出 if "moving_variance" in ms_name: val_np = np.maximum(val_np, 1e-5) @@ -106,23 +123,16 @@ def universal_convert(task="segment", scale="n"): print(f"[对齐成功] {ms_name:<55} <- {expected_pt_name}") matched_count += 1 - - # 内存优化：匹配成功后从字典中移除该键，以便最后审计遗留项 del pt_dict[expected_pt_name] else: print(f"[形状冲突] 参数 {ms_name} 期望尺寸 {ms_shape}，实际载入尺寸 {tuple(pt_v.shape)}") else: - # 处理常量和不可训练张量 (如 DFL 积分权重、步长锚点) - # 由于网络初始化时已经通过 Config 构建了正确的值，此处直接保留并存入 Checkpoint 即可 if "stride" in ms_name or "dfl.conv.weight" in ms_name: new_ms_ckpt.append({'name': ms_name, 'data': ms_param.data}) print(f"[保留原值] {ms_name:<55} (框架内置固定张量)") matched_count += 1 - else: - # 若非常量且未能在 PT 字典中找到映射，留交审计模块处理 - pass - # 5. 输出转换审计报告 + # 4. 输出审计并保存 print("-" * 80) print("[INFO] 权重转换审计清单") matched_names = [x['name'] for x in new_ms_ckpt] @@ -135,19 +145,23 @@ def universal_convert(task="segment", scale="n"): for fn in failed_names: print(f" - {fn}") - # 6. 持久化存储 - base_name = os.path.splitext(pt_name_map[task])[0] - save_ckpt_path = f"{base_name}.ckpt" - - ms.save_checkpoint(new_ms_ckpt, save_ckpt_path) - print(f"[INFO] 转换结束。参数对齐率: {matched_count}/{len(ms_params)}。已序列化至: {save_ckpt_path}") + ms.save_checkpoint(new_ms_ckpt, ckpt_path) + print(f"[INFO] 转换结束。参数对齐率: {matched_count}/{len(ms_params)}。已序列化至: {ckpt_path}") - return save_ckpt_path + # 清理打工完毕的临时 NumPy 文件 + if os.path.exists(npz_file): + os.remove(npz_file) + + return ckpt_path if __name__ == "__main__": parser = argparse.ArgumentParser(description="YOLO11 参数转换工具 (PyTorch to MindSpore Checkpoint)") - parser.add_argument("--task", "-t", type=str, default="pose", + + parser.add_argument("--pt_path", type=str, default=None, help="输入的 .pt 文件路径") + parser.add_argument("--ckpt_path", type=str, default=None, help="输出的 .ckpt 文件路径") + + parser.add_argument("--task", "-t", type=str, default="detect", choices=["classify", "detect", "segment", "pose"], help="目标模型的基础任务类型") parser.add_argument("--scale", "-s", type=str, default="n", @@ -155,4 +169,11 @@ def universal_convert(task="segment", scale="n"): help="指定模型的规模缩放因子") args = parser.parse_args() - universal_convert(args.task, args.scale) \ No newline at end of file + if args.pt_path is None: + suffix_map = {"classify": "-cls", "segment": "-seg", "pose": "-pose", "detect": ""} + args.pt_path = f"yolo11{args.scale}{suffix_map[args.task]}.pt" + + if args.ckpt_path is None: + args.ckpt_path = args.pt_path.replace(".pt", ".ckpt") + + universal_convert(args.pt_path, args.ckpt_path, args.task, args.scale) \ No newline at end of file diff --git a/src/mindnlp/models/ultralytics/models/yolo/pose/__init__.py b/src/mindnlp/ultralytics/utils/__init__.py similarity index 100% rename from src/mindnlp/models/ultralytics/models/yolo/pose/__init__.py rename to src/mindnlp/ultralytics/utils/__init__.py diff --git a/src/mindnlp/models/ultralytics/utils/ema.py b/src/mindnlp/ultralytics/utils/ema.py similarity index 100% rename from src/mindnlp/models/ultralytics/utils/ema.py rename to src/mindnlp/ultralytics/utils/ema.py diff --git a/src/mindnlp/models/ultralytics/utils/loss.py b/src/mindnlp/ultralytics/utils/loss.py similarity index 99% rename from src/mindnlp/models/ultralytics/utils/loss.py rename to src/mindnlp/ultralytics/utils/loss.py index ce80fa9eb..36ce52d4a 100644 --- a/src/mindnlp/models/ultralytics/utils/loss.py +++ b/src/mindnlp/ultralytics/utils/loss.py @@ -3,8 +3,8 @@ import mindspore.ops as ops from mindspore import nn, Tensor -from utils.tal import TaskAlignedAssigner, make_anchors -from utils.ops import bbox_iou, xywh2xyxy, dist2bbox, bbox2dist +from ultralytics.utils.tal import TaskAlignedAssigner, make_anchors +from ultralytics.utils.ops import bbox_iou, xywh2xyxy, dist2bbox, bbox2dist # COCO 数据集 17 个关键点的标准差（用于计算 OKS 相似度） OKS_SIGMA = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0 diff --git a/src/mindnlp/models/ultralytics/utils/metrics.py b/src/mindnlp/ultralytics/utils/metrics.py similarity index 100% rename from src/mindnlp/models/ultralytics/utils/metrics.py rename to src/mindnlp/ultralytics/utils/metrics.py diff --git a/src/mindnlp/models/ultralytics/utils/ops.py b/src/mindnlp/ultralytics/utils/ops.py similarity index 100% rename from src/mindnlp/models/ultralytics/utils/ops.py rename to src/mindnlp/ultralytics/utils/ops.py diff --git a/src/mindnlp/models/ultralytics/utils/optimizer.py b/src/mindnlp/ultralytics/utils/optimizer.py similarity index 100% rename from src/mindnlp/models/ultralytics/utils/optimizer.py rename to src/mindnlp/ultralytics/utils/optimizer.py diff --git a/src/mindnlp/models/ultralytics/utils/tal.py b/src/mindnlp/ultralytics/utils/tal.py similarity index 99% rename from src/mindnlp/models/ultralytics/utils/tal.py rename to src/mindnlp/ultralytics/utils/tal.py index 8cd109ea7..d8d1ca800 100644 --- a/src/mindnlp/models/ultralytics/utils/tal.py +++ b/src/mindnlp/ultralytics/utils/tal.py @@ -3,7 +3,7 @@ from mindspore import nn, ops # 规范引入算子 -from utils.ops import batch_iou +from ultralytics.utils.ops import batch_iou def make_anchors(feats, strides, grid_cell_offset=0.5): """ From d60ca442c6abfd584c74db5e77692803ba8daa9a Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Mon, 30 Mar 2026 22:54:25 +0800 Subject: [PATCH 4/6] fix: add dataset configuration yaml files --- .../ultralytics/cfg/datasets/coco128-seg.yaml | 101 ++++++++++++++++++ .../ultralytics/cfg/datasets/coco128.yaml | 101 ++++++++++++++++++ .../ultralytics/cfg/datasets/coco8-pose.yaml | 47 ++++++++ .../cfg/datasets/imagenette2-160.yaml | 28 +++++ 4 files changed, 277 insertions(+) create mode 100644 src/mindnlp/ultralytics/cfg/datasets/coco128-seg.yaml create mode 100644 src/mindnlp/ultralytics/cfg/datasets/coco128.yaml create mode 100644 src/mindnlp/ultralytics/cfg/datasets/coco8-pose.yaml create mode 100644 src/mindnlp/ultralytics/cfg/datasets/imagenette2-160.yaml diff --git a/src/mindnlp/ultralytics/cfg/datasets/coco128-seg.yaml b/src/mindnlp/ultralytics/cfg/datasets/coco128-seg.yaml new file mode 100644 index 000000000..cbca742ec --- /dev/null +++ b/src/mindnlp/ultralytics/cfg/datasets/coco128-seg.yaml @@ -0,0 +1,101 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO128-seg dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/segment/coco/ +# Example usage: yolo train data=coco128-seg.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco128-seg ← downloads here (7 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/coco128-seg # dataset root dir +train: images/train2017 # train images (relative to 'path') 128 images +val: images/train2017 # val images (relative to 'path') 128 images +test: # test images (optional) + +# Classes +names: + 0: person + 1: bicycle + 2: car + 3: motorcycle + 4: airplane + 5: bus + 6: train + 7: truck + 8: boat + 9: traffic light + 10: fire hydrant + 11: stop sign + 12: parking meter + 13: bench + 14: bird + 15: cat + 16: dog + 17: horse + 18: sheep + 19: cow + 20: elephant + 21: bear + 22: zebra + 23: giraffe + 24: backpack + 25: umbrella + 26: handbag + 27: tie + 28: suitcase + 29: frisbee + 30: skis + 31: snowboard + 32: sports ball + 33: kite + 34: baseball bat + 35: baseball glove + 36: skateboard + 37: surfboard + 38: tennis racket + 39: bottle + 40: wine glass + 41: cup + 42: fork + 43: knife + 44: spoon + 45: bowl + 46: banana + 47: apple + 48: sandwich + 49: orange + 50: broccoli + 51: carrot + 52: hot dog + 53: pizza + 54: donut + 55: cake + 56: chair + 57: couch + 58: potted plant + 59: bed + 60: dining table + 61: toilet + 62: tv + 63: laptop + 64: mouse + 65: remote + 66: keyboard + 67: cell phone + 68: microwave + 69: oven + 70: toaster + 71: sink + 72: refrigerator + 73: book + 74: clock + 75: vase + 76: scissors + 77: teddy bear + 78: hair drier + 79: toothbrush + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128-seg.zip diff --git a/src/mindnlp/ultralytics/cfg/datasets/coco128.yaml b/src/mindnlp/ultralytics/cfg/datasets/coco128.yaml new file mode 100644 index 000000000..42bc624e6 --- /dev/null +++ b/src/mindnlp/ultralytics/cfg/datasets/coco128.yaml @@ -0,0 +1,101 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO128 dataset https://www.kaggle.com/datasets/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/detect/coco/ +# Example usage: yolo train data=coco128.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco128 ← downloads here (7 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/coco128 # dataset root dir +train: images/train2017 # train images (relative to 'path') 128 images +val: images/train2017 # val images (relative to 'path') 128 images +test: # test images (optional) + +# Classes +names: + 0: person + 1: bicycle + 2: car + 3: motorcycle + 4: airplane + 5: bus + 6: train + 7: truck + 8: boat + 9: traffic light + 10: fire hydrant + 11: stop sign + 12: parking meter + 13: bench + 14: bird + 15: cat + 16: dog + 17: horse + 18: sheep + 19: cow + 20: elephant + 21: bear + 22: zebra + 23: giraffe + 24: backpack + 25: umbrella + 26: handbag + 27: tie + 28: suitcase + 29: frisbee + 30: skis + 31: snowboard + 32: sports ball + 33: kite + 34: baseball bat + 35: baseball glove + 36: skateboard + 37: surfboard + 38: tennis racket + 39: bottle + 40: wine glass + 41: cup + 42: fork + 43: knife + 44: spoon + 45: bowl + 46: banana + 47: apple + 48: sandwich + 49: orange + 50: broccoli + 51: carrot + 52: hot dog + 53: pizza + 54: donut + 55: cake + 56: chair + 57: couch + 58: potted plant + 59: bed + 60: dining table + 61: toilet + 62: tv + 63: laptop + 64: mouse + 65: remote + 66: keyboard + 67: cell phone + 68: microwave + 69: oven + 70: toaster + 71: sink + 72: refrigerator + 73: book + 74: clock + 75: vase + 76: scissors + 77: teddy bear + 78: hair drier + 79: toothbrush + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128.zip diff --git a/src/mindnlp/ultralytics/cfg/datasets/coco8-pose.yaml b/src/mindnlp/ultralytics/cfg/datasets/coco8-pose.yaml new file mode 100644 index 000000000..185e366ad --- /dev/null +++ b/src/mindnlp/ultralytics/cfg/datasets/coco8-pose.yaml @@ -0,0 +1,47 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license + +# COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics +# Documentation: https://docs.ultralytics.com/datasets/pose/coco8-pose/ +# Example usage: yolo train data=coco8-pose.yaml +# parent +# ├── ultralytics +# └── datasets +# └── coco8-pose ← downloads here (1 MB) + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/coco8-pose # dataset root dir +train: images/train # train images (relative to 'path') 4 images +val: images/val # val images (relative to 'path') 4 images +test: # test images (optional) + +# Keypoints +kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) +flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] + +# Classes +names: + 0: person + +# Keypoint names per class +kpt_names: + 0: + - nose + - left_eye + - right_eye + - left_ear + - right_ear + - left_shoulder + - right_shoulder + - left_elbow + - right_elbow + - left_wrist + - right_wrist + - left_hip + - right_hip + - left_knee + - right_knee + - left_ankle + - right_ankle + +# Download script/URL (optional) +download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-pose.zip diff --git a/src/mindnlp/ultralytics/cfg/datasets/imagenette2-160.yaml b/src/mindnlp/ultralytics/cfg/datasets/imagenette2-160.yaml new file mode 100644 index 000000000..3d699147a --- /dev/null +++ b/src/mindnlp/ultralytics/cfg/datasets/imagenette2-160.yaml @@ -0,0 +1,28 @@ +# Ultralytics AGPL-3.0 License - https://ultralytics.com/license +# Imagenette2-160 Dataset by fast.ai +# Example usage: python train.py --data imagenette2-160.yaml + +# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] +path: ./datasets/imagenette2-160 # 你的数据集根目录 (根据你的截图和运行路径调整) +train: train # 训练集相对路径 (会与 path 拼接) +val: val # 验证集相对路径 (会与 path 拼接) +test: # 测试集 (Imagenette 没有单独的 test，可以留空) + +# Classes +nc: 10 # 类别数量 +names: + 0: tench # 丁鱼 + 1: English springer # 英国斯宾格犬 + 2: cassette player # 录音机 + 3: chain saw # 链锯 + 4: church # 教堂 + 5: French horn # 圆号 + 6: garbage truck # 垃圾车 + 7: gas pump # 加油泵 + 8: golf ball # 高尔夫球 + 9: parachute # 降落伞 + +# Download script/URL (optional) +download: | + wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz + tar -xzvf imagenette2-160.tgz \ No newline at end of file From 3d4510d3eb8ae7b09c42e41e9090c3be985462ac Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Tue, 31 Mar 2026 12:42:49 +0800 Subject: [PATCH 5/6] fix: force add ignored data module and update readme with weights download links --- src/mindnlp/ultralytics/data/__init__.py | 0 src/mindnlp/ultralytics/data/augment.py | 102 +++++++++ src/mindnlp/ultralytics/data/dataset.py | 253 +++++++++++++++++++++++ src/mindnlp/ultralytics/data/loaders.py | 252 ++++++++++++++++++++++ src/mindnlp/ultralytics/readme.md | 12 +- 5 files changed, 618 insertions(+), 1 deletion(-) create mode 100644 src/mindnlp/ultralytics/data/__init__.py create mode 100644 src/mindnlp/ultralytics/data/augment.py create mode 100644 src/mindnlp/ultralytics/data/dataset.py create mode 100644 src/mindnlp/ultralytics/data/loaders.py diff --git a/src/mindnlp/ultralytics/data/__init__.py b/src/mindnlp/ultralytics/data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/mindnlp/ultralytics/data/augment.py b/src/mindnlp/ultralytics/data/augment.py new file mode 100644 index 000000000..cca03a566 --- /dev/null +++ b/src/mindnlp/ultralytics/data/augment.py @@ -0,0 +1,102 @@ +import cv2 +import random +import numpy as np +import mindspore.dataset.vision as vision + +class Compose: + """组合多种数据增强变换""" + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, img, labels=None): + for t in self.transforms: + if labels is None: + img = t(img) + else: + img, labels = t(img, labels) + return (img, labels) if labels is not None else img + +class RandomHSV: + """随机调整图像的色调(Hue)、饱和度(Saturation)和明度(Value)""" + def __init__(self, hgain=0.015, sgain=0.7, vgain=0.4): + self.hgain = hgain + self.sgain = sgain + self.vgain = vgain + + def __call__(self, img, labels=None): + r = np.random.uniform(-1, 1, 3) * [self.hgain, self.sgain, self.vgain] + 1 + hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_RGB2HSV)) + dtype = img.dtype + + x = np.arange(0, 256, dtype=r.dtype) + lut_hue = ((x * r[0]) % 180).astype(dtype) + lut_sat = np.clip(x * r[1], 0, 255).astype(dtype) + lut_val = np.clip(x * r[2], 0, 255).astype(dtype) + + img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))) + img = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB) + return (img, labels) if labels is not None else img + +class RandomFlip: + """ + 随机水平翻转图像及其对应的空间标签 + 支持基础目标检测边界框以及姿态估计关键点的同步翻转 + """ + def __init__(self, p=0.5, flip_idx=None): + self.p = p + # flip_idx 记录翻转后关键点对应的新索引位置 (如 COCO 17 个关键点的对称关系) + self.flip_idx = flip_idx + + def __call__(self, img, labels=None): + if random.random() < self.p: + img = np.ascontiguousarray(img[:, ::-1]) + if labels is not None: + # 翻转边界框中心点 X 坐标 (格式为 [cls, cx, cy, w, h, ...]) + if len(labels.shape) > 1 and labels.shape[1] >= 5: + labels[:, 1] = 1.0 - labels[:, 1] + + # 翻转姿态关键点 + if len(labels.shape) > 1 and labels.shape[1] > 5 and self.flip_idx is not None: + kpts = labels[:, 5:].reshape(labels.shape[0], -1, 3) + + # 仅对可见的关键点执行翻转 + mask = kpts[..., 2] > 0 + kpts[..., 0][mask] = 1.0 - kpts[..., 0][mask] + + # 交换左右对称的关键点索引 + kpts = kpts[:, self.flip_idx, :] + labels[:, 5:] = kpts.reshape(labels.shape[0], -1) + + return (img, labels) if labels is not None else img + +def get_classify_transforms(imgsz=224, is_training=True): + """构建分类任务专用变换流水线 (基于 MindSpore Vision 算子加速)""" + trans = [] + if is_training: + trans += [ + vision.RandomResizedCrop(imgsz), + vision.RandomHorizontalFlip(prob=0.5) + ] + else: + trans += [ + vision.Resize(int(imgsz * 1.14)), + vision.CenterCrop(imgsz) + ] + + trans += [ + vision.Rescale(1.0 / 255.0, 0.0), + vision.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), + vision.HWC2CHW() + ] + return trans + +def v8_transforms(imgsz=640, hyp=None, flip_idx=None): + """构建 YOLO 通用检测/姿态/分割任务变换组合""" + hyp = hyp or {} + return Compose([ + # 修复项：使用字典的 .get() 方法安全获取 YAML 参数 + RandomHSV(hgain=hyp.get('hsv_h', 0.015), + sgain=hyp.get('hsv_s', 0.7), + vgain=hyp.get('hsv_v', 0.4)), + RandomFlip(p=hyp.get('fliplr', 0.5), flip_idx=flip_idx) + ]) \ No newline at end of file diff --git a/src/mindnlp/ultralytics/data/dataset.py b/src/mindnlp/ultralytics/data/dataset.py new file mode 100644 index 000000000..f351d59ce --- /dev/null +++ b/src/mindnlp/ultralytics/data/dataset.py @@ -0,0 +1,253 @@ +import os +import cv2 +import numpy as np +from pathlib import Path + +class YOLOBaseDataset: + """数据集基类，提供通用的图像读取与空间几何变换方法""" + def __init__(self, imgsz=640): + self.imgsz = imgsz + + def load_image(self, img_path): + """读取图像并转换为 RGB 格式。""" + img = cv2.imread(str(img_path)) + if img is None: + return np.zeros((self.imgsz, self.imgsz, 3), dtype=np.uint8) + return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + def letterbox(self, img, bboxes=None, keypoints=None, segments=None, color=(114, 114, 114)): + """ + 几何变换核心逻辑：保持长宽比缩放并填充边缘 + 同步支持边界框 (Bbox)、关键点 (Keypoints) 及多边形 (Segments) 的坐标映射与归一化 + """ + shape = img.shape[:2] + h0, w0 = shape + r = min(self.imgsz / h0, self.imgsz / w0) + new_unpad = int(round(w0 * r)), int(round(h0 * r)) + dw = (self.imgsz - new_unpad[0]) / 2 + dh = (self.imgsz - new_unpad[1]) / 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, 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)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + + # 边界框坐标映射与截断保护 (防止因 Padding 导致坐标溢出 1.0) + if bboxes is not None and len(bboxes) > 0: + bboxes[:, [0, 2]] *= new_unpad[0] + bboxes[:, [1, 3]] *= new_unpad[1] + bboxes[:, 0] += dw + bboxes[:, 1] += dh + bboxes /= self.imgsz + bboxes = np.clip(bboxes, 0.0, 1.0) + + # 关键点坐标映射 + if keypoints is not None and len(keypoints) > 0: + mask = keypoints[..., 2] > 0 + keypoints[..., 0][mask] = keypoints[..., 0][mask] * new_unpad[0] + dw + keypoints[..., 1][mask] = keypoints[..., 1][mask] * new_unpad[1] + dh + keypoints[..., 0][mask] /= self.imgsz + keypoints[..., 1][mask] /= self.imgsz + keypoints[..., :2] = np.clip(keypoints[..., :2], 0.0, 1.0) + + # 实例分割多边形坐标映射 + if segments is not None and len(segments) > 0: + for i in range(len(segments)): + segments[i][:, 0] *= new_unpad[0] + segments[i][:, 1] *= new_unpad[1] + segments[i][:, 0] += dw + segments[i][:, 1] += dh + segments[i] /= self.imgsz + + result = [img] + if bboxes is not None: result.append(bboxes) + if keypoints is not None: result.append(keypoints) + if segments is not None: result.append(segments) + + return tuple(result) if len(result) > 1 else img + + +class YOLOClassifyDataset(YOLOBaseDataset): + """图像分类数据集接口""" + def __init__(self, root, imgsz=224): + super().__init__(imgsz) + self.root = Path(root) + + # 基于 train 目录建立类别索引映射，确保 train 与 val 映射关系一致 + data_root = self.root.parent + train_dir = data_root / 'train' + reference_dir = train_dir if train_dir.exists() else self.root + + self.classes = sorted([d.name for d in reference_dir.iterdir() if d.is_dir()]) + self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)} + + self.samples = [] + for cls in self.classes: + cls_dir = self.root / cls + if not cls_dir.exists(): + continue + + for img_path in cls_dir.glob("*"): + if img_path.suffix.lower() in ['.jpg', '.jpeg', '.png']: + self.samples.append((str(img_path), self.class_to_idx[cls])) + + if not self.samples: + raise RuntimeError(f"[ERROR] 目录 {self.root} 中未检测到有效图像样本。") + + def __getitem__(self, i): + path, label = self.samples[i] + img = self.load_image(path) + return img, np.array(label, dtype=np.int32) + + def __len__(self): + return len(self.samples) + + +class YOLODetectDataset(YOLOBaseDataset): + """目标检测数据集接口""" + def __init__(self, img_path, imgsz=640, transforms=None): + super().__init__(imgsz) + self.img_path = Path(img_path) + self.im_files = sorted([str(f) for f in self.img_path.rglob("*") + if f.suffix.lower() in ['.jpg', '.png', '.jpeg']]) + + # 基于目录结构约定，由图像路径推导标签路径 + self.label_files = [p.replace(f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}") + .replace(".jpg", ".txt").replace(".png", ".txt") + for p in self.im_files] + self.transforms = transforms + + def __getitem__(self, index): + img = self.load_image(self.im_files[index]) + label_path = self.label_files[index] + + if os.path.exists(label_path): + with open(label_path, 'r') as f: + content = f.read().strip() + l = np.array([x.split() for x in content.splitlines()], dtype=np.float32) if content else np.zeros((0, 5), dtype=np.float32) + else: + l = np.zeros((0, 5), dtype=np.float32) + + if self.transforms is not None and len(l) > 0: + img, l = self.transforms(img, l) + + cls = l[:, 0:1] if len(l) else np.zeros((0, 1), dtype=np.float32) + bboxes = l[:, 1:5] if len(l) else np.zeros((0, 4), dtype=np.float32) + + img, bboxes = self.letterbox(img, bboxes=bboxes) + return img, cls, bboxes + + def __len__(self): + return len(self.im_files) + + +class YOLOSegmentDataset(YOLOBaseDataset): + """实例分割数据集接口""" + def __init__(self, img_path, imgsz=640, transforms=None): + super().__init__(imgsz) + self.img_path = Path(img_path) + self.im_files = sorted([str(f) for f in self.img_path.rglob("*") + if f.suffix.lower() in ['.jpg', '.png', '.jpeg']]) + self.label_files = [p.replace(f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}") + .replace(".jpg", ".txt").replace(".png", ".txt") + for p in self.im_files] + self.transforms = transforms + + def __getitem__(self, index): + img = self.load_image(self.im_files[index]) + label_path = self.label_files[index] + + cls_list, segments, raw_labels = [], [], [] + + if os.path.exists(label_path): + with open(label_path, 'r') as f: + for line in f.read().strip().splitlines(): + parts = list(map(float, line.split())) + if len(parts) >= 5: + raw_labels.append(parts) + + # 当前针对分割任务仅执行像素级的数据增强（如 HSV 变换），不执行空间几何变换 + if self.transforms is not None and len(raw_labels) > 0: + img = self.transforms(img) + + for parts in raw_labels: + cls_list.append([parts[0]]) + poly = np.array(parts[1:], dtype=np.float32).reshape(-1, 2) + segments.append(poly) + + if len(segments) > 0: + img, segments = self.letterbox(img, segments=segments) + else: + img = self.letterbox(img) + + bboxes, masks = [], [] + # 由经过归一化缩放后的多边形坐标生成掩码拓扑和边界框 + for poly in segments: + x1, y1 = poly[:, 0].min(), poly[:, 1].min() + x2, y2 = poly[:, 0].max(), poly[:, 1].max() + bboxes.append([(x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1]) + + mask = np.zeros((self.imgsz, self.imgsz), dtype=np.uint8) + poly_pixel = (poly * [self.imgsz, self.imgsz]).astype(np.int32) + cv2.fillPoly(mask, [poly_pixel], 1) + masks.append(mask) + + if not cls_list: + return img, np.zeros((0, 1), dtype=np.float32), np.zeros((0, 4), dtype=np.float32), np.zeros((0, self.imgsz, self.imgsz), dtype=np.float32) + + return (img, + np.array(cls_list, dtype=np.float32), + np.array(bboxes, dtype=np.float32), + np.array(masks, dtype=np.float32)) + + def __len__(self): + return len(self.im_files) + + +class YOLOPoseDataset(YOLOBaseDataset): + """ + 姿态估计数据集接口 + 预期标签格式: [class_id, x, y, w, h, kpt1_x, kpt1_y, kpt1_v, ...] + """ + def __init__(self, img_path, imgsz=640, nkpt=17, ndim=3, transforms=None): + super().__init__(imgsz) + self.img_path = Path(img_path) + self.im_files = sorted([str(f) for f in self.img_path.rglob("*") + if f.suffix.lower() in ['.jpg', '.png', '.jpeg']]) + self.label_files = [p.replace(f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}") + .replace(".jpg", ".txt").replace(".png", ".txt") + for p in self.im_files] + self.nkpt = nkpt + self.ndim = ndim + self.transforms = transforms + + def __getitem__(self, index): + img = self.load_image(self.im_files[index]) + label_path = self.label_files[index] + + if os.path.exists(label_path): + with open(label_path, 'r') as f: + content = f.read().strip() + labels = np.array([x.split() for x in content.splitlines()], dtype=np.float32) if content else np.zeros((0, 5 + self.nkpt * self.ndim), dtype=np.float32) + else: + labels = np.zeros((0, 5 + self.nkpt * self.ndim), dtype=np.float32) + + if self.transforms is not None and len(labels) > 0: + img, labels = self.transforms(img, labels) + + if len(labels): + cls = labels[:, 0:1] + bboxes = labels[:, 1:5] + keypoints = labels[:, 5:].reshape(-1, self.nkpt, self.ndim) + else: + cls = np.zeros((0, 1), dtype=np.float32) + bboxes = np.zeros((0, 4), dtype=np.float32) + keypoints = np.zeros((0, self.nkpt, self.ndim), dtype=np.float32) + + img, bboxes, keypoints = self.letterbox(img, bboxes=bboxes, keypoints=keypoints) + + return img, cls, bboxes, keypoints + + def __len__(self): + return len(self.im_files) \ No newline at end of file diff --git a/src/mindnlp/ultralytics/data/loaders.py b/src/mindnlp/ultralytics/data/loaders.py new file mode 100644 index 000000000..352180088 --- /dev/null +++ b/src/mindnlp/ultralytics/data/loaders.py @@ -0,0 +1,252 @@ +import os +import cv2 +import glob +import numpy as np +import mindspore.dataset as ds +from pathlib import Path +from .dataset import (YOLOClassifyDataset, YOLODetectDataset, + YOLOSegmentDataset, YOLOPoseDataset) +from .augment import get_classify_transforms, v8_transforms + +# Collate 聚合函数：针对各类任务生成特定的 batch_idx +def yolo_collate_fn(imgs, clss, bboxes, batch_info=None): + """检测任务批量重组映射函数""" + batch_imgs = np.stack(imgs, axis=0) + batch_imgs = batch_imgs.transpose(0, 3, 1, 2) + batch_imgs = batch_imgs.astype(np.float32) + batch_cls, batch_bboxes, batch_idx = [], [], [] + + for i in range(len(clss)): + n = clss[i].shape[0] + if n > 0: + batch_cls.append(clss[i]) + batch_bboxes.append(bboxes[i]) + batch_idx.append(np.full((n, 1), i, dtype=np.float32)) + + if len(batch_cls) > 0: + batch_cls = np.concatenate(batch_cls, axis=0) + batch_bboxes = np.concatenate(batch_bboxes, axis=0) + batch_idx = np.concatenate(batch_idx, axis=0) + else: + batch_cls = np.zeros((0, 1), dtype=np.float32) + batch_bboxes = np.zeros((0, 4), dtype=np.float32) + batch_idx = np.zeros((0, 1), dtype=np.float32) + + return batch_imgs, batch_cls, batch_bboxes, batch_idx + + +def pose_collate_fn(imgs, clss, bboxes, kpts, batch_info=None): + """姿态估计任务批量重组映射函数""" + batch_imgs = np.stack(imgs, axis=0) + batch_imgs = batch_imgs.transpose(0, 3, 1, 2) + batch_imgs = batch_imgs.astype(np.float32) + batch_cls, batch_bboxes, batch_kpts, batch_idx = [], [], [], [] + + for i in range(len(clss)): + n = clss[i].shape[0] + if n > 0: + batch_cls.append(clss[i]) + batch_bboxes.append(bboxes[i]) + batch_kpts.append(kpts[i]) + batch_idx.append(np.full((n, 1), i, dtype=np.float32)) + + if len(batch_cls) > 0: + batch_cls = np.concatenate(batch_cls, axis=0) + batch_bboxes = np.concatenate(batch_bboxes, axis=0) + batch_kpts = np.concatenate(batch_kpts, axis=0) + batch_idx = np.concatenate(batch_idx, axis=0) + else: + batch_cls = np.zeros((0, 1), dtype=np.float32) + batch_bboxes = np.zeros((0, 4), dtype=np.float32) + batch_kpts = np.zeros((0, 17, 3), dtype=np.float32) + batch_idx = np.zeros((0, 1), dtype=np.float32) + + return batch_imgs, batch_cls, batch_bboxes, batch_kpts, batch_idx + + +def segment_collate_fn(imgs, clss, bboxes, masks, batch_info=None): + """实例分割任务批量重组映射函数""" + batch_imgs = np.stack(imgs, axis=0) + if batch_imgs.shape[-1] == 3: + batch_imgs = batch_imgs.transpose(0, 3, 1, 2) + batch_imgs = batch_imgs.astype(np.float32) + batch_cls, batch_bboxes, batch_masks, batch_idx = [], [], [], [] + + for i in range(len(clss)): + n = clss[i].shape[0] + if n > 0: + batch_cls.append(clss[i]) + batch_bboxes.append(bboxes[i]) + batch_masks.append(masks[i]) + batch_idx.append(np.full((n, 1), i, dtype=np.int32)) + + if len(batch_cls) > 0: + batch_cls = np.concatenate(batch_cls, axis=0) + batch_bboxes = np.concatenate(batch_bboxes, axis=0) + batch_masks = np.concatenate(batch_masks, axis=0) + batch_idx = np.concatenate(batch_idx, axis=0) + else: + batch_cls = np.zeros((0, 1), dtype=np.float32) + batch_bboxes = np.zeros((0, 4), dtype=np.float32) + batch_masks = np.zeros((0, 160, 160), dtype=np.float32) + batch_idx = np.zeros((0, 1), dtype=np.float32) + + return batch_imgs.astype(np.float32), batch_cls, batch_bboxes, batch_masks, batch_idx + + +# 数据加载器组装器 +def create_dataloader(path, imgsz=640, batch_size=16, task='classify', is_training=True, num_workers=8, hyp=None): + """构建多任务兼容的 MindSpore 数据预处理流水线""" + + # 针对姿态任务所需的 COCO 对称点映射字典 + coco_flip_idx = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] + + # 动态注入数据增强策略 (验证阶段则保持原图特征提取) + common_transforms = v8_transforms(imgsz=imgsz, hyp=hyp, flip_idx=coco_flip_idx) if is_training else None + + if task == 'classify': + dataset_generator = YOLOClassifyDataset(path, imgsz=imgsz) + column_names = ["image", "label"] + elif task == 'detect': + dataset_generator = YOLODetectDataset(path, imgsz=imgsz, transforms=common_transforms) + source_column_names = ["image", "cls", "bboxes"] + elif task == 'segment': + dataset_generator = YOLOSegmentDataset(path, imgsz=imgsz, transforms=common_transforms) + source_column_names = ["image", "cls", "bboxes", "masks"] + elif task == 'pose': + dataset_generator = YOLOPoseDataset(path, imgsz=imgsz, transforms=common_transforms) + source_column_names = ["image", "cls", "bboxes", "keypoints"] + else: + raise ValueError(f"[ERROR] 尚不支持的任务类型: {task}") + + # 封装为 MindSpore GeneratorDataset 并激活多进程加载引擎 + dataset = ds.GeneratorDataset( + source=dataset_generator, + column_names=source_column_names if task != 'classify' else column_names, + shuffle=is_training, + num_parallel_workers=num_workers if num_workers > 0 else 1, + python_multiprocessing=True + ) + + if task == 'classify': + trans = get_classify_transforms(imgsz=imgsz, is_training=is_training) + dataset = dataset.map( + operations=trans, + input_columns="image", + num_parallel_workers=num_workers + ) + dataset = dataset.project(column_names) + dataset = dataset.batch(batch_size, drop_remainder=is_training, num_parallel_workers=num_workers) + + elif task == 'detect': + dataset = dataset.project(source_column_names) + dataset = dataset.batch( + batch_size, + per_batch_map=yolo_collate_fn, + output_columns=["image", "cls", "bboxes", "batch_idx"], + drop_remainder=is_training + ) + + elif task == 'segment': + dataset = dataset.project(source_column_names) + dataset = dataset.batch( + batch_size, + per_batch_map=segment_collate_fn, + input_columns=source_column_names, + output_columns=["image", "cls", "bboxes", "masks", "batch_idx"], + drop_remainder=is_training, + num_parallel_workers=num_workers + ) + + elif task == 'pose': + dataset = dataset.project(source_column_names) + dataset = dataset.batch( + batch_size, + per_batch_map=pose_collate_fn, + output_columns=["image", "cls", "bboxes", "keypoints", "batch_idx"], + drop_remainder=is_training + ) + + return dataset + + +# 推理阶段专用精简加载器 +def letterbox_classify(im, new_shape=224): + """分类图像的预处理：等比例缩放短边后，进行中心裁剪""" + h, w = im.shape[:2] + r = new_shape / min(h, w) + h_new, w_new = int(h * r), int(w * r) + im = cv2.resize(im, (w_new, h_new), interpolation=cv2.INTER_LINEAR) + + top = (h_new - new_shape) // 2 + left = (w_new - new_shape) // 2 + return im[top : top + new_shape, left : left + new_shape] + +def letterbox_pad(img, new_shape=(640, 640), color=(114, 114, 114)): + """检测/分割/姿态图像的预处理：等比例缩放后，进行边缘填充 (Padding)""" + shape = img.shape[:2] + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] + dw, dh = dw / 2, dh / 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, 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)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img + +class LoadImages: + """ + 适用于本地部署的前向推理轻量化加载器 + 基于迭代器模式设计，内存占用小，支持多任务预处理路由 + """ + def __init__(self, path, imgsz=640, task='detect'): + p = str(Path(path).absolute()) + if os.path.isdir(p): + self.files = sorted(glob.glob(os.path.join(p, '*.*'))) + elif os.path.isfile(p): + self.files = [p] + else: + raise FileNotFoundError(f"[ERROR] 未能索引至有效的文件路径: {p}") + + self.imgsz = imgsz + self.task = task # 记录当前任务类型 + self.count = 0 + self.nf = len(self.files) + + def __iter__(self): + self.count = 0 + return self + + def __next__(self): + if self.count == self.nf: + raise StopIteration + + path = self.files[self.count] + self.count += 1 + + img = cv2.imread(path) + if img is None: + print(f"[WARNING] 无法正常解析图像源文件: {path}，予以跳过。") + return self.__next__() + + # 保留一份原始图像，供后续画框、保存使用 + img0 = img.copy() + + # BGR 转换为 RGB + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + # 核心逻辑路由：根据任务类型执行不同的几何预处理 + if self.task == 'classify': + img = letterbox_classify(img, new_shape=self.imgsz) + else: + img = letterbox_pad(img, new_shape=(self.imgsz, self.imgsz)) + + return path, img, img0 + + def __len__(self): + return self.nf \ No newline at end of file diff --git a/src/mindnlp/ultralytics/readme.md b/src/mindnlp/ultralytics/readme.md index 95762f7a5..3a39bbaaf 100644 --- a/src/mindnlp/ultralytics/readme.md +++ b/src/mindnlp/ultralytics/readme.md @@ -103,4 +103,14 @@ python examples/yolo/pose/inference.py ## 4. 用户级 API 调用示例 ```bash python standalone.py -``` \ No newline at end of file +``` + +## 5. 预训练权重下载 (Pre-trained Weights) +若预处理脚本自动下载权重文件失败，请点击下方链接手动下载官方 PyTorch 权重文件（.pt），并使用本文档第 2 节中的 convert.py 脚本将其转换为 MindSpore 可用的 .ckpt 权重文件： +目标检测 (Detect): https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt + +图像分类 (Classify): https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-cls.pt + +实例分割 (Segment): https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-seg.pt + +姿态估计 (Pose): https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-pose.pt \ No newline at end of file From 315d393c608eaa6305ec17c0b2e7ca448c3b8e57 Mon Sep 17 00:00:00 2001 From: kittentruck <771228437@qq.com> Date: Fri, 24 Apr 2026 21:28:52 +0800 Subject: [PATCH 6/6] Update ultralytics compatibility files for MindSpore NLP --- src/mindnlp/ultralytics/engine/trainer.py | 52 +++++++++++-------- src/mindnlp/ultralytics/engine/validator.py | 23 +++++++- src/mindnlp/ultralytics/models/model.py | 4 +- .../ultralytics/models/yolo/classify/val.py | 10 +++- .../ultralytics/models/yolo/pose/val.py | 3 +- .../ultralytics/models/yolo/segment/val.py | 5 +- src/mindnlp/ultralytics/readme.md | 2 +- src/mindnlp/ultralytics/standalone.py | 2 +- src/mindnlp/ultralytics/tools/convert.py | 6 ++- 9 files changed, 73 insertions(+), 34 deletions(-) diff --git a/src/mindnlp/ultralytics/engine/trainer.py b/src/mindnlp/ultralytics/engine/trainer.py index 1cd02ec4d..b20227a50 100644 --- a/src/mindnlp/ultralytics/engine/trainer.py +++ b/src/mindnlp/ultralytics/engine/trainer.py @@ -108,29 +108,35 @@ def train(self): loss_val = float(loss.asnumpy()) if hasattr(loss, "asnumpy") else float(loss) print(f"Epoch [{epoch}/{self.epochs-1}] Step [{step}/{steps_per_epoch}] | Loss: {loss_val:.4f}") + # 每轮先保存最新权重，确保训练产物可用 + ms.save_checkpoint(self.ema.ema_model, str(self.save_dir / "last.ckpt")) + # 验证与保存阶段 - if (epoch + 1) % self.args.val_interval == 0 or epoch == self.epochs - 1: - print(f"\n[INFO] 开始执行 Epoch {epoch} 验证程序...") - self.model.set_train(False) - - validator = self.get_validator() - stats = validator(self.ema.ema_model) - - print("-" * 50) - print(f"[评估报告] Epoch {epoch}") - for k, v in stats.items(): - if k.startswith('metrics/'): - metric_name = k.replace('metrics/', '') - print(f" - {metric_name:<15} : {float(v):.5f}") - - fitness_f = float(stats.get('fitness', 0.0)) - print(f"[INFO] 当前模型综合评价指标 (Fitness): {fitness_f:.5f}") - print("-" * 50 + "\n") - - self._save_checkpoint(epoch, fitness_f) - - self.model.set_train(True) - self.train_step.set_train(True) + if (epoch + 1) % self.args.val_interval == 0: + try: + print(f"\n[INFO] 开始执行 Epoch {epoch} 验证程序...") + self.model.set_train(False) + + validator = self.get_validator() + stats = validator(self.ema.ema_model) + + print("-" * 50) + print(f"[评估报告] Epoch {epoch}") + for k, v in stats.items(): + if k.startswith('metrics/'): + metric_name = k.replace('metrics/', '') + print(f" - {metric_name:<15} : {float(v):.5f}") + + fitness_f = float(stats.get('fitness', 0.0)) + print(f"[INFO] 当前模型综合评价指标 (Fitness): {fitness_f:.5f}") + print("-" * 50 + "\n") + + self._save_checkpoint(epoch, fitness_f) + except Exception as e: + print(f"[WARNING] Epoch {epoch} 验证阶段执行失败，已跳过 best.ckpt 更新。错误信息: {e}") + finally: + self.model.set_train(True) + self.train_step.set_train(True) def _save_checkpoint(self, epoch, fitness): """权重序列化保存逻辑""" @@ -169,4 +175,4 @@ def preprocess_batch(self, batch): if not isinstance(batch["batch_idx"], ms.Tensor): batch["batch_idx"] = ms.Tensor(batch["batch_idx"], ms.int32) - return batch \ No newline at end of file + return batch diff --git a/src/mindnlp/ultralytics/engine/validator.py b/src/mindnlp/ultralytics/engine/validator.py index 1d803116b..640961056 100644 --- a/src/mindnlp/ultralytics/engine/validator.py +++ b/src/mindnlp/ultralytics/engine/validator.py @@ -1,3 +1,4 @@ +import os import time import logging import yaml @@ -47,6 +48,26 @@ def __call__(self, model): model.set_train(False) self.init_metrics(model) + if self.dataloader is None: + if self.args is None: + raise ValueError("验证阶段缺少数据集配置参数 data，无法构建 dataloader。") + + if isinstance(self.args, dict): + data_path = self.args.get("data") + batch_size = self.args.get("batch", self.args.get("batch_size", 16)) + else: + data_path = getattr(self.args, "data", None) + batch_size = getattr(self.args, "batch", getattr(self.args, "batch_size", 16)) + + if not data_path: + raise ValueError("验证阶段缺少数据集配置参数 data，无法构建 dataloader。") + + with open(data_path, "r", encoding="utf-8") as f: + data_cfg = yaml.safe_load(f) + + val_path = os.path.join(data_cfg.get('path', ''), data_cfg.get('val', 'val')) + self.dataloader = self.get_dataloader(val_path, batch_size=batch_size) + bar = tqdm(self.dataloader.create_dict_iterator(), desc="Validating", total=self.dataloader.get_dataset_size()) @@ -120,4 +141,4 @@ def get_stats(self): def print_results(self): speed_str = " | ".join([f"{k}: {v:.1f}ms" for k, v in self.get_stats().get('speed', {}).items()]) - LOGGER.info(f"推理测速: {speed_str}") \ No newline at end of file + LOGGER.info(f"推理测速: {speed_str}") diff --git a/src/mindnlp/ultralytics/models/model.py b/src/mindnlp/ultralytics/models/model.py index cb05b7676..39df94928 100644 --- a/src/mindnlp/ultralytics/models/model.py +++ b/src/mindnlp/ultralytics/models/model.py @@ -177,7 +177,7 @@ def val(self, **kwargs): else: raise ValueError(f"[MindNLP YOLO] 暂不支持的任务类型: {self.task}") - return validator(model=self.model) + return validator(model=pass_model) def __call__(self, source=None, **kwargs): """ @@ -198,4 +198,4 @@ def __call__(self, source=None, **kwargs): else: raise ValueError(f"[MindNLP YOLO] 暂不支持的任务类型: {self.task}") - return predictor(source=source, model=self.model) \ No newline at end of file + return predictor(source=source, model=self.model) diff --git a/src/mindnlp/ultralytics/models/yolo/classify/val.py b/src/mindnlp/ultralytics/models/yolo/classify/val.py index 944d2a034..539701bf8 100644 --- a/src/mindnlp/ultralytics/models/yolo/classify/val.py +++ b/src/mindnlp/ultralytics/models/yolo/classify/val.py @@ -19,8 +19,16 @@ def __init__(self, dataloader=None, save_dir=None, args=None): # 确保任务类型正确声明 if self.args is None: self.args = SimpleNamespace(task="classify", half=False) + elif isinstance(self.args, dict): + args_dict = dict(self.args) + args_dict["task"] = "classify" + if "half" not in args_dict: + args_dict["half"] = False + self.args = SimpleNamespace(**args_dict) else: self.args.task = "classify" + if not hasattr(self.args, "half"): + self.args.half = False self.names = None @@ -85,4 +93,4 @@ def print_results(self): LOGGER.info( f"验证结果 | Top-1 Acc: {stats.get('metrics/accuracy_top1', 0):.4f} | " f"Top-5 Acc: {stats.get('metrics/accuracy_top5', 0):.4f}" - ) \ No newline at end of file + ) diff --git a/src/mindnlp/ultralytics/models/yolo/pose/val.py b/src/mindnlp/ultralytics/models/yolo/pose/val.py index 5da11385a..de9e62a93 100644 --- a/src/mindnlp/ultralytics/models/yolo/pose/val.py +++ b/src/mindnlp/ultralytics/models/yolo/pose/val.py @@ -3,6 +3,7 @@ from mindspore import ops from ultralytics.engine.validator import BaseValidator +from ultralytics.data.loaders import create_dataloader from ultralytics.utils.ops import non_max_suppression, xywh2xyxy_np from ultralytics.utils.metrics import PoseMetrics, kpt_iou @@ -216,4 +217,4 @@ def match_predictions(self, pred_cls, target_cls, iou, iou_thresholds=None): tp[match_data[:, 0].astype(int), j] = True - return tp \ No newline at end of file + return tp diff --git a/src/mindnlp/ultralytics/models/yolo/segment/val.py b/src/mindnlp/ultralytics/models/yolo/segment/val.py index 3e3f83924..ac2bf8818 100644 --- a/src/mindnlp/ultralytics/models/yolo/segment/val.py +++ b/src/mindnlp/ultralytics/models/yolo/segment/val.py @@ -5,6 +5,7 @@ from mindspore import Tensor, ops from ultralytics.engine.validator import BaseValidator +from ultralytics.data.loaders import create_dataloader from ultralytics.utils.ops import non_max_suppression, process_mask, xywh2xyxy_np from ultralytics.utils.metrics import SegmentMetrics @@ -86,7 +87,7 @@ def update_metrics(self, preds, batch): batch_idx = batch["batch_idx"].view(-1).asnumpy() all_gt_cls = batch["cls"].view(-1).asnumpy() all_gt_bboxes = batch["bboxes"].asnumpy() - all_gt_masks = batch["masks"] + all_gt_masks = batch["masks"].asnumpy() imgsz = getattr(self.args, 'imgsz', 640) @@ -202,4 +203,4 @@ def calculate_tp(self, preds, pred_cls, gts, gt_cls, is_mask=False): match_matrix = (iou_np >= threshold) & cc_np tp[:, i] = match_matrix.any(axis=1) - return tp \ No newline at end of file + return tp diff --git a/src/mindnlp/ultralytics/readme.md b/src/mindnlp/ultralytics/readme.md index 3a39bbaaf..b9dc9dacf 100644 --- a/src/mindnlp/ultralytics/readme.md +++ b/src/mindnlp/ultralytics/readme.md @@ -3,7 +3,7 @@ 本项目基于 MindSpore 框架实现了 YOLO11 的四大核心任务：图像分类、目标检测、实例分割和姿态估计。支持从头训练、加载预训练权重微调、模型验证与推理。 ```bash -# 创建环境 (Python 3.9) +# 创建环境 (Python 3.10) conda create -n mindnlp_yolo python=3.10 -y conda activate mindnlp_yolo diff --git a/src/mindnlp/ultralytics/standalone.py b/src/mindnlp/ultralytics/standalone.py index eb863dcb0..1d2081184 100644 --- a/src/mindnlp/ultralytics/standalone.py +++ b/src/mindnlp/ultralytics/standalone.py @@ -36,7 +36,7 @@ # 训练模型 results = model.train( data="coco128.yaml", - epochs=100, + epochs=10, imgsz=640, batch=16, amp=False, diff --git a/src/mindnlp/ultralytics/tools/convert.py b/src/mindnlp/ultralytics/tools/convert.py index 3f70b932a..213d65e1e 100644 --- a/src/mindnlp/ultralytics/tools/convert.py +++ b/src/mindnlp/ultralytics/tools/convert.py @@ -88,7 +88,9 @@ def universal_convert(pt_path, ckpt_path, task="detect", scale="n"): f.write(extract_script.strip()) # 运行子进程提取数据 - subprocess.run(["python", "temp_extract.py"], check=True) + env = os.environ.copy() + env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + subprocess.run([sys.executable, "temp_extract.py"], check=True, env=env) os.remove("temp_extract.py") # 2. 回到主进程：读取绝对纯净的 NumPy 字典，彻底断绝与 PyTorch 的瓜葛 @@ -176,4 +178,4 @@ def universal_convert(pt_path, ckpt_path, task="detect", scale="n"): if args.ckpt_path is None: args.ckpt_path = args.pt_path.replace(".pt", ".ckpt") - universal_convert(args.pt_path, args.ckpt_path, args.task, args.scale) \ No newline at end of file + universal_convert(args.pt_path, args.ckpt_path, args.task, args.scale)