使用SuperGradients中的YOLO-NAS-Pose模型进行姿态估计评估-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00594/article/details/148508986

使用SuperGradients中的YOLO-NAS-Pose模型进行姿态估计评估

super-gradients Deci-AI/super-gradients: 是一个用于深度学习模型优化的库。适合用于需要训练和优化深度学习模型的研究者和开发者。特点是可以提供高效的优化算法和工具，支持多种深度学习框架。项目地址: https://gitcode.com/gh_mirrors/su/super-gradients

评估方法概述

在计算机视觉领域，姿态估计是一项重要的任务，它需要精确地检测人体关键点并确定其空间位置。SuperGradients提供了YOLO-NAS-Pose模型，这是一个高效的人体姿态估计解决方案。本文将详细介绍如何使用PyCocoTools对YOLO-NAS-Pose模型进行离线评估。

评估协议差异

SuperGradients(SG)和COCO官方评估协议在图像预处理方面存在关键差异：

SG评估方式：
- 训练/验证时将图像统一调整为固定尺寸(默认为640x640)
- 使用保持长宽比的缩放+填充方式
- 在调整后的图像分辨率上计算AP/AR指标
COCO官方评估：
- 使用原始图像分辨率进行评估
- 不进行任何图像尺寸调整

这种差异通常会导致SG报告的指标比COCO官方评估低约1AP。因此，为了获得与学术研究一致的评估结果，我们需要使用COCO官方评估协议。

环境准备与模型加载

首先需要安装必要的软件包：

!pip install -qq super_gradients==3.7.1

然后加载预训练的YOLO-NAS-Pose模型：

from super_gradients.common.object_names import Models
from super_gradients.training import models

model = models.get(
    Models.YOLO_NAS_POSE_L,
    pretrained_weights="coco_pose"
).cuda()

也可以加载自定义训练好的模型：

model = models.get(
    Models.YOLO_NAS_POSE_L,
    num_classes=17,
    checkpoint_path="/path/to/your/checkpoint.pth",
)

准备COCO验证数据

评估需要COCO2017验证集的数据和标注。首先设置数据路径：

import os
COCO_DATA_DIR = os.environ.get("COCO_ROOT_DIR", "/path/to/coco2017")

然后加载图像和标注：

from pycocotools.coco import COCO

images_path = os.path.join(COCO_DATA_DIR, "images/val2017")
image_files = [os.path.join(images_path, x) for x in os.listdir(images_path)]

gt_annotations_path = os.path.join(COCO_DATA_DIR, "annotations/person_keypoints_val2017.json")
gt = COCO(gt_annotations_path)

模型预测

使用加载的模型对验证集进行预测：

predictions = model.predict(
    image_files, conf=0.01, iou=0.7, 
    pre_nms_max_predictions=300, 
    post_nms_max_predictions=20, 
    fuse_model=False
)

结果转换与评估

将预测结果转换为COCO评估工具要求的格式：

import copy
import json_tricks as json
import collections
import numpy as np
import tempfile

def predictions_to_coco(predictions, image_files):
    # 实现预测结果到COCO格式的转换
    ...

coco_pred = predictions_to_coco(predictions, image_files)

使用COCO评估工具计算最终指标：

with tempfile.TemporaryDirectory() as td:
    res_file = os.path.join(td, "keypoints_coco2017_results.json")
    
    with open(res_file, "w") as f:
        json.dump(coco_pred, f)
    
    coco_dt = copy.deepcopy(gt)
    coco_dt = coco_dt.loadRes(res_file)
    
    coco_evaluator = COCOeval(gt, coco_dt, iouType="keypoints")
    coco_evaluator.evaluate()
    coco_evaluator.accumulate()
    coco_evaluator.summarize()

评估指标解读

COCO评估工具会输出多个关键指标：

平均精度(AP)指标：
- AP @0.5..0.95: 0.682 (在不同IoU阈值下的平均精度)
- AP @0.5: 0.891 (IoU阈值为0.5时的精度)
- AP @0.75: 0.752 (IoU阈值为0.75时的精度)
平均召回率(AR)指标：
- AR @0.5..0.95: 0.735 (在不同IoU阈值下的平均召回率)
- AR @0.5: 0.924 (IoU阈值为0.5时的召回率)
- AR @0.75: 0.799 (IoU阈值为0.75时的召回率)

这些指标全面反映了模型在不同严格程度下的性能表现，为模型优化和比较提供了客观依据。