YOLO26如何训练自己的关键点pose数据集 | （手部关键点为案列）

重复模块m (n次迭代):
┌─────────────────────────────────────────────────────────┐
│                                                         │
│  如果 attn=True:                                        │
│    Sequential(                                          │
│        Bottleneck(self.c, self.c),                     │  ←─ 先特征提取
│        PSABlock(self.c, attn_ratio=0.5, num_heads=...) │  ←─ 后注意力增强
│    )                                                    │
│                                                         │
│  否则如果 c3k=True:                                     │
│    C3k(self.c, self.c, 2)                             │  ←─ 同YOLOv11
│                                                         │
│  否则:                                                  │
│    Bottleneck(self.c, self.c)                         │  ←─ 同YOLOv11
│                                                         │
└─────────────────────────────────────────────────────────┘

class BboxLoss(nn.Module):
    """Criterion class for computing training losses for bounding boxes."""

    def __init__(self, reg_max: int = 16):
        """Initialize the BboxLoss module with regularization maximum and DFL settings."""
        super().__init__()
        self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None

# YOLO26手势关键点检测伪代码
# 假设检测21个手势关键点（类似MediaPipe手势模型）

class YOLO26_HandPose:
    def __init__(self):
        # 初始化参数
        self.num_keypoints = 21  # 21个手势关键点
        self.num_classes = 1     # 手部类别
        self.input_size = 640    # 输入尺寸
        
    def build_backbone(self):
        """改进的骨干网络"""
        # 使用CSPNet或RepVGG风格的架构
        backbone = Sequential([
            FocusLayer(3, 32),  # 聚焦层减少计算量
            CSPBlock(32, 64, num_blocks=3),
            CSPBlock(64, 128, num_blocks=6),
            CSPBlock(128, 256, num_blocks=9),
            CSPBlock(256, 512, num_blocks=3),
            SPPF(512, 512)  # 空间金字塔池化
        ])
        return backbone
        
    def build_neck(self):
        """特征金字塔网络（改进版）"""
        # 使用PANet或BiFPN进行多尺度特征融合
        neck = BiFPN([
            FeatureMap(512, 256),
            FeatureMap(256, 128),
            FeatureMap(128, 64)
        ])
        return neck
        
    def build_head(self):
        """检测头 - 输出边界框、关键点和置信度"""
        # 每个检测头包含：
        # 1. 边界框预测 (x, y, w, h, obj_conf)
        # 2. 类别预测 (cls_conf)
        # 3. 关键点预测 (21个关键点的x,y,可见性)
        
        head_outputs = {
            'bbox': Conv2D(256, 4 + 1),  # 4个bbox坐标 + 1个obj置信度
            'cls': Conv2D(256, self.num_classes),
            'keypoints': Conv2D(256, self.num_keypoints * 3)  # 每个点(x,y,visibility)
        }
        return head_outputs
        
    def forward(self, x):
        """前向传播"""
        # 1. 特征提取
        features = self.backbone(x)
        
        # 2. 多尺度特征融合
        pyramid_features = self.neck(features)
        
        # 3. 多尺度预测
        predictions = []
        for feature in pyramid_features:
            # 在不同尺度上进行预测
            bbox_pred = self.head.bbox(feature)
            cls_pred = self.head.cls(feature)
            kpts_pred = self.head.keypoints(feature)
            
            predictions.append({
                'bbox': bbox_pred,
                'cls': cls_pred,
                'keypoints': kpts_pred
            })
        
        return predictions
        
    def decode_predictions(self, predictions, conf_threshold=0.5):
        """解码预测结果"""
        decoded_results = []
        
        for pred in predictions:
            # 应用sigmoid激活
            bbox_scores = sigmoid(pred['bbox'][..., 4:5])
            cls_scores = sigmoid(pred['cls'])
            kpts_scores = sigmoid(pred['keypoints'][..., 2::3])
            
            # 筛选高置信度检测
            mask = bbox_scores > conf_threshold
            
            if mask.any():
                # 解码边界框 (从grid坐标到原图坐标)
                bboxes = self.decode_boxes(pred['bbox'][mask])
                
                # 解码关键点
                keypoints = self.decode_keypoints(pred['keypoints'][mask])
                
                # 非极大值抑制
                indices = nms(bboxes, bbox_scores[mask])
                
                for idx in indices:
                    result = {
                        'bbox': bboxes[idx],
                        'confidence': bbox_scores[idx],
                        'class_id': argmax(cls_scores[mask][idx]),
                        'keypoints': keypoints[idx],  # shape: (21, 3) - [x, y, visibility]
                        'keypoints_confidence': kpts_scores[mask][idx]
                    }
                    decoded_results.append(result)
        
        return decoded_results
        
    def decode_keypoints(self, kpts_pred):
        """将关键点预测解码为实际坐标"""
        # kpts_pred shape: [N, 21*3]
        # 前两个通道是坐标，第三个通道是可见性
        
        batch_size = kpts_pred.shape[0]
        keypoints = kpts_pred.reshape(batch_size, self.num_keypoints, 3)
        
        # 应用sigmoid到可见性分数
        keypoints[..., 2] = sigmoid(keypoints[..., 2])
        
        return keypoints
        
    def calculate_loss(self, predictions, targets):
        """损失函数计算"""
        total_loss = 0
        
        # 1. 边界框损失 (CIoU损失)
        bbox_loss = ciou_loss(predictions['bbox'][..., :4], targets['bboxes'])
        
        # 2. 目标置信度损失 (二元交叉熵)
        obj_loss = binary_cross_entropy(
            predictions['bbox'][..., 4:5],
            targets['obj_mask']
        )
        
        # 3. 类别损失 (交叉熵)
        cls_loss = cross_entropy(
            predictions['cls'],
            targets['cls']
        )
        
        # 4. 关键点损失 (改进的Wing Loss或MSE)
        kpts_loss = wing_loss(
            predictions['keypoints'][..., :2],  # 坐标
            targets['keypoints'][..., :2]
        )
        
        # 5. 关键点可见性损失
        kpts_vis_loss = binary_cross_entropy(
            predictions['keypoints'][..., 2:3],  # 可见性
            targets['keypoints'][..., 2:3]
        )
        
        # 总损失 (加权和)
        total_loss = (
            bbox_loss * 0.05 +
            obj_loss * 0.7 +
            cls_loss * 0.3 +
            kpts_loss * 0.1 +
            kpts_vis_loss * 0.05
        )
        
        return total_loss
        
    def train_step(self, images, targets):
        """训练步骤"""
        # 前向传播
        predictions = self.forward(images)
        
        # 计算损失
        loss = self.calculate_loss(predictions, targets)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        return loss
        
    def visualize_results(self, image, results):
        """可视化检测结果"""
        for result in results:
            # 绘制边界框
            draw_bbox(image, result['bbox'], result['confidence'])
            
            # 绘制关键点
            for kp in result['keypoints']:
                x, y, visible = kp
                if visible > 0.5:  # 只绘制可见的关键点
                    draw_keypoint(image, (x, y))
            
            # 绘制关键点连接线
            connections = [
                (0,1), (1,2), (2,3), (3,4),        # 拇指
                (0,5), (5,6), (6,7), (7,8),        # 食指
                (0,9), (9,10), (10,11), (11,12),   # 中指
                (0,13), (13,14), (14,15), (15,16), # 无名指
                (0,17), (17,18), (18,19), (19,20)  # 小指
            ]
            
            for start_idx, end_idx in connections:
                start_kp = result['keypoints'][start_idx]
                end_kp = result['keypoints'][end_idx]
                
                if start_kp[2] > 0.5 and end_kp[2] > 0.5:
                    draw_line(image, start_kp[:2], end_kp[:2])
        
        return image


# 使用示例
def main():
    # 初始化模型
    model = YOLO26_HandPose()
    
    # 训练模式
    model.train()
    for epoch in range(100):
        for batch_images, batch_targets in train_loader:
            loss = model.train_step(batch_images, batch_targets)
            
    # 推理模式
    model.eval()
    image = load_image("hand_image.jpg")
    
    # 预处理
    processed_img = preprocess(image)
    
    # 推理
    with torch.no_grad():
        predictions = model.forward(processed_img)
        results = model.decode_predictions(predictions)
    
    # 可视化
    output_image = model.visualize_results(image, results)
    save_image(output_image, "result.jpg")


# 关键点定义（21个点）
KEYPOINT_NAMES = [
    'WRIST',                    # 0 - 手腕
    'THUMB_CMC', 'THUMB_MCP', 'THUMB_IP', 'THUMB_TIP',      # 1-4 - 拇指
    'INDEX_FINGER_MCP', 'INDEX_FINGER_PIP',                 # 5-6 - 食指
    'INDEX_FINGER_DIP', 'INDEX_FINGER_TIP',                 # 7-8
    'MIDDLE_FINGER_MCP', 'MIDDLE_FINGER_PIP',               # 9-10 - 中指
    'MIDDLE_FINGER_DIP', 'MIDDLE_FINGER_TIP',               # 11-12
    'RING_FINGER_MCP', 'RING_FINGER_PIP',                   # 13-14 - 无名指
    'RING_FINGER_DIP', 'RING_FINGER_TIP',                   # 15-16
    'PINKY_MCP', 'PINKY_PIP', 'PINKY_DIP', 'PINKY_TIP'      # 17-20 - 小指
]

import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26-pose.yaml')
    #model.load('yolo26n-seg.pt') # loading pretrain weights
    model.train(data='data/hand-keypoints/hand-keypoints.yaml',
                cache=False,
                imgsz=640,
                epochs=200,
                batch=32,
                close_mosaic=10,
                workers=2,
                device='0',
                # optimizer='SGD',   # using SGD,auto
                project='runs/train',
                name='exp',
                )

YOLO26-pose summary (fused): 132 layers, 3,116,730 parameters, 0 gradients, 8.3 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Pose(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 123/123 2.7it/s 46.1s
                   all       7847       7847      0.979      0.977      0.992      0.914      0.923      0.907      0.938      0.833