如何用YOLO NAS统计单张图像中各类检测目标的数量？

核心实现流程 #

YOLO NAS的预测结果包含检测框、置信度及类别标签，只需解析这些输出即可完成单张图像的类别计数，以下是具体步骤：

1. 加载模型与单张图像 #

先确保安装super_gradients库，再加载预训练模型并读取目标图像：

from super_gradients.training import models
import cv2

# 加载YOLO NAS模型（可替换为yolo_nas_m/yolo_nas_l）
model = models.get("yolo_nas_s", pretrained_weights="coco")

# 读取单张水果图像
image_path = "your_fruit_image.jpg"
image = cv2.imread(image_path)

2. 执行单张图像预测 #

对目标图像完成检测，获取预测结果：

# 执行预测，返回DetectionResults对象
predictions = model.predict(image)

3. 解析结果并统计类别数量 #

提取检测到的类别标签，映射为类别名称后统计数量：

from collections import Counter

# 获取模型对应的类别名称列表（自定义数据集需替换为你的类别列表）
class_names = model.get_class_names()

# 提取所有检测结果的类别标签索引
detected_label_ids = [pred.class_id for pred in predictions[0].predictions]

# 映射为类别名称并统计数量
fruit_count = Counter([class_names[id] for id in detected_label_ids])

4. 展示统计结果与标注图像 #

• 打印统计信息：

print("单张图像水果统计结果：")
for fruit, count in fruit_count.items():
    print(f"{fruit}: {count}个")

• 在图像上标注检测框与类别并展示：

# 提取检测框、类别名称、置信度
bboxes = predictions[0].predictions.bboxes_xyxy
labels = [class_names[pred.class_id] for pred in predictions[0].predictions]
confidences = predictions[0].predictions.confidence

# 在图像上绘制标注
for bbox, label, conf in zip(bboxes, labels, confidences):
    x1, y1, x2, y2 = map(int, bbox)
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(image, f"{label} {conf:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 显示标注后的图像
cv2.imshow("Fruit Detection Result", image)
cv2.waitKey(0)
cv2.destroyAllWindows()

自定义数据集适配 #

若使用自定义训练的水果检测模型，只需将class_names替换为你的数据集类别列表，例如：

class_names = ["apple", "banana", "orange", "grape"]

内容的提问来源于stack exchange，提问作者CM379