YOLONAS-pose无法获取检测目标类名问题求助
解决YoloNAS-pose无法获取检测目标类名的问题
核心原因
YoloNAS-pose的预测结果结构和纯检测的YoloNAS略有差异,其姿态分支主要针对COCO数据集的person类,但检测部分仍支持COCO全80类,只是需要正确解析结果中的类别标识。
解决方案
1. 加载正确的预训练模型
确保加载带COCO全类别检测+姿态的预训练权重,避免使用仅针对人体姿态的轻量化版本:
from super_gradients.training import models # 加载COCO预训练的YoloNAS-pose模型 model = models.get("yolo_nas_pose", pretrained_weights="coco_pose")
2. 正确解析预测结果获取类名
YoloNAS-pose的每个检测结果包含class_id字段,可通过两种方式映射得到类名:
方法一:使用模型内置类别列表
直接调用模型方法获取COCO类名,无需手动定义:
# 对目标图片执行预测 predictions = model.predict("你的图片路径.jpg") # 获取模型内置的COCO类别列表 coco_classes = model.get_class_names() # 遍历所有检测结果 for pred in predictions: class_id = pred.class_id class_name = coco_classes[class_id] print(f"检测目标:{class_name}") # 若目标是person,可额外查看姿态关键点 if class_name == "person": print(f"姿态关键点:{pred.keypoints}")
方法二:手动定义COCO类名列表(备用)
如果model.get_class_names()无法正常获取类名,可手动定义COCO的80类列表进行映射:
COCO_CLASSES = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] # 解析预测结果 for pred in predictions: class_id = pred.class_id class_name = COCO_CLASSES[class_id] print(f"检测目标类名:{class_name}")
3. 特殊情况排查
- 若检测目标非
person类,YoloNAS-pose只会输出检测框和类名,不会返回姿态关键点,属于正常现象。 - 若仍无法获取类名,可打印预测结果的所有属性(
print(pred.__dict__)),确认class_id或其他类别相关字段的位置,再调整解析逻辑。
内容的提问来源于stack exchange,提问作者vijay ganesh
相关产品推荐
相关产品推荐

