460 lines
17 KiB
Python
460 lines
17 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
Ultralytics YOLO 测试脚本
|
|||
|
|
- validate 模式：在 dataset.yaml 上评测 mAP、精度/召回
|
|||
|
|
- predict 模式：对图片/文件夹/视频/摄像头进行推理并保存可视化
|
|||
|
|
- test 模式：在指定 source 数据上评估模型指标（mAP、精度、召回等）
|
|||
|
|
- annotation 模式：对图像进行检测并生成LabelMe格式的JSON标注文件（保存到与图像相同目录）
|
|||
|
|
|
|||
|
|
示例：
|
|||
|
|
# 在验证集上评测
|
|||
|
|
python3 detection/test_yolo.py \
|
|||
|
|
--weights runs/train/single_yolov8s_20250908_152257/weights/best.pt \
|
|||
|
|
--data ../datasets/yolo_dataset/dataset.yaml \
|
|||
|
|
--imgsz 640 \
|
|||
|
|
--mode validate
|
|||
|
|
|
|||
|
|
# 对文件夹做推理并保存结果
|
|||
|
|
python3 detection/test_yolo.py \
|
|||
|
|
--weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
|
|||
|
|
--source ./datasets/l0_9.9_sum/images/test \
|
|||
|
|
--imgsz 640 \
|
|||
|
|
--conf 0.25 \
|
|||
|
|
--mode predict \
|
|||
|
|
--project runs/predict \
|
|||
|
|
--name demo
|
|||
|
|
|
|||
|
|
# 在指定数据上评估模型指标
|
|||
|
|
python3 detection/test_yolo.py \
|
|||
|
|
--weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
|
|||
|
|
--source ./datasets/l0_9.9_sum/images/test \
|
|||
|
|
--imgsz 640 \
|
|||
|
|
--conf 0.25 \
|
|||
|
|
--mode test \
|
|||
|
|
--project runs/test \
|
|||
|
|
--name eval_test
|
|||
|
|
|
|||
|
|
# 生成LabelMe格式的JSON标注文件，并像 predict 模式一样输出预测可视化
|
|||
|
|
python3 detection/test_yolo.py \
|
|||
|
|
--weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
|
|||
|
|
--source ./7009-7509\
|
|||
|
|
--imgsz 640 \
|
|||
|
|
--conf 0.25 \
|
|||
|
|
--mode annotation \
|
|||
|
|
--project runs/annotations\
|
|||
|
|
--name annotations14
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import argparse
|
|||
|
|
import numpy as np
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_args():
|
|||
|
|
ap = argparse.ArgumentParser(description="YOLO 模型测试/推理")
|
|||
|
|
ap.add_argument("--weights", required=True, help="权重路径，例如 runs/train/.../weights/best.pt")
|
|||
|
|
ap.add_argument("--mode", choices=["validate", "predict", "test", "annotation"], default="validate", help="测试模式")
|
|||
|
|
ap.add_argument("--data", default="", help="dataset.yaml 路径(用于 validate)")
|
|||
|
|
ap.add_argument("--source", default="", help="推理输入(用于 predict/test)：图片/文件夹/视频/rtsp/0(摄像头)")
|
|||
|
|
ap.add_argument("--imgsz", type=int, default=640, help="输入尺寸")
|
|||
|
|
ap.add_argument("--conf", type=float, default=0.25, help="置信度阈值(预测)")
|
|||
|
|
ap.add_argument("--iou", type=float, default=0.45, help="NMS IOU 阈值(预测)")
|
|||
|
|
ap.add_argument("--device", default="", help="CUDA 设备，如 '0' 或 '0,1'，留空自动")
|
|||
|
|
ap.add_argument("--project", default="runs/test", help="输出根目录")
|
|||
|
|
ap.add_argument("--name", default="", help="任务名(默认自动)")
|
|||
|
|
ap.add_argument("--save_txt", action="store_true", help="保存 YOLO txt 结果(预测)")
|
|||
|
|
ap.add_argument("--save_conf", action="store_true", help="在 txt 中保存置信度")
|
|||
|
|
ap.add_argument("--recursive", action="store_true", help="当 --source 为目录时递归查找(默认仅当前目录)")
|
|||
|
|
return ap.parse_args()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def list_images_in_dir(dir_path: Path, recursive: bool = False):
|
|||
|
|
exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
|
|||
|
|
if recursive:
|
|||
|
|
return [p for p in dir_path.rglob("*") if p.suffix.lower() in exts]
|
|||
|
|
else:
|
|||
|
|
return [p for p in dir_path.iterdir() if p.is_file() and p.suffix.lower() in exts]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def evaluate_on_source(model, source_path: str, args):
|
|||
|
|
"""
|
|||
|
|
在指定的 source 数据上评估模型指标
|
|||
|
|
"""
|
|||
|
|
from ultralytics.utils.metrics import ConfusionMatrix
|
|||
|
|
import torch
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
src = Path(source_path)
|
|||
|
|
if not src.exists():
|
|||
|
|
print(f"[错误] 源路径不存在: {src}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 收集所有图像文件
|
|||
|
|
if src.is_dir():
|
|||
|
|
image_files = list_images_in_dir(src, recursive=args.recursive)
|
|||
|
|
if not image_files:
|
|||
|
|
print(f"[错误] 目录下未找到图像: {src}")
|
|||
|
|
return None
|
|||
|
|
else:
|
|||
|
|
image_files = [src]
|
|||
|
|
|
|||
|
|
print(f"[测试] 找到 {len(image_files)} 个图像文件")
|
|||
|
|
|
|||
|
|
# 创建临时 dataset.yaml 用于评估
|
|||
|
|
import tempfile
|
|||
|
|
import yaml
|
|||
|
|
|
|||
|
|
# 假设标签文件在对应的 labels 目录中
|
|||
|
|
labels_dir = src.parent.parent / "labels" / src.name
|
|||
|
|
if not labels_dir.exists():
|
|||
|
|
print(f"[警告] 未找到标签目录: {labels_dir}")
|
|||
|
|
print("[测试] 将仅进行推理，无法计算精确指标")
|
|||
|
|
return run_inference_only(model, image_files, args)
|
|||
|
|
|
|||
|
|
# 创建临时 dataset.yaml
|
|||
|
|
temp_yaml = {
|
|||
|
|
'path': str(src.parent),
|
|||
|
|
'train': '', # 不使用
|
|||
|
|
'val': str(src.relative_to(src.parent)), # 使用 source 作为验证集
|
|||
|
|
'test': '',
|
|||
|
|
'nc': 1, # 假设单类检测
|
|||
|
|
'names': ['fish'] # 类别名称
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
|
|||
|
|
yaml.dump(temp_yaml, f)
|
|||
|
|
temp_yaml_path = f.name
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 使用 model.val 进行评估
|
|||
|
|
print("[测试] 开始评估...")
|
|||
|
|
metrics = model.val(
|
|||
|
|
data=temp_yaml_path,
|
|||
|
|
imgsz=args.imgsz,
|
|||
|
|
conf=args.conf,
|
|||
|
|
iou=args.iou,
|
|||
|
|
device=args.device if args.device else None,
|
|||
|
|
project=args.project,
|
|||
|
|
name=args.name or "test_eval",
|
|||
|
|
verbose=True,
|
|||
|
|
save_json=True, # 保存详细结果
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 打印详细指标
|
|||
|
|
print("\n" + "="*50)
|
|||
|
|
print("评估结果:")
|
|||
|
|
print("="*50)
|
|||
|
|
|
|||
|
|
# print number of images
|
|||
|
|
print(f"总图像数: {len(image_files)}")
|
|||
|
|
|
|||
|
|
if hasattr(metrics, 'box'):
|
|||
|
|
box_metrics = metrics.box
|
|||
|
|
print(f"mAP@0.5: {box_metrics.map50:.4f}")
|
|||
|
|
print(f"mAP@0.5:0.95: {box_metrics.map:.4f}")
|
|||
|
|
print(f"Precision: {box_metrics.mp:.4f}")
|
|||
|
|
print(f"Recall: {box_metrics.mr:.4f}")
|
|||
|
|
print(f"F1-Score: {2 * box_metrics.mp * box_metrics.mr / (box_metrics.mp + box_metrics.mr):.4f}")
|
|||
|
|
|
|||
|
|
if hasattr(metrics, 'speed'):
|
|||
|
|
speed = metrics.speed
|
|||
|
|
print(f"推理速度: {speed['inference']:.2f} ms/image")
|
|||
|
|
print(f"NMS速度: {speed['postprocess']:.2f} ms/image")
|
|||
|
|
|
|||
|
|
print("="*50)
|
|||
|
|
|
|||
|
|
return metrics
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[错误] 评估失败: {e}")
|
|||
|
|
print("[测试] 回退到仅推理模式...")
|
|||
|
|
return run_inference_only(model, image_files, args)
|
|||
|
|
finally:
|
|||
|
|
# 清理临时文件
|
|||
|
|
try:
|
|||
|
|
os.unlink(temp_yaml_path)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run_inference_only(model, image_files, args):
|
|||
|
|
"""
|
|||
|
|
仅进行推理，计算基本统计信息
|
|||
|
|
"""
|
|||
|
|
print("[测试] 仅推理模式 - 无法计算精确指标")
|
|||
|
|
|
|||
|
|
total_detections = 0
|
|||
|
|
total_images = len(image_files)
|
|||
|
|
confidences = []
|
|||
|
|
|
|||
|
|
for img_path in image_files:
|
|||
|
|
try:
|
|||
|
|
results = model.predict(
|
|||
|
|
source=str(img_path),
|
|||
|
|
imgsz=args.imgsz,
|
|||
|
|
conf=args.conf,
|
|||
|
|
iou=args.iou,
|
|||
|
|
device=args.device if args.device else None,
|
|||
|
|
verbose=False,
|
|||
|
|
save=True, # 同步保存Ultralytics自带的可视化结果
|
|||
|
|
project=args.project,
|
|||
|
|
name=args.name or "annotations",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if results and len(results) > 0:
|
|||
|
|
res = results[0]
|
|||
|
|
if hasattr(res, 'boxes') and res.boxes is not None:
|
|||
|
|
num_dets = len(res.boxes)
|
|||
|
|
total_detections += num_dets
|
|||
|
|
|
|||
|
|
# 收集置信度
|
|||
|
|
for box in res.boxes:
|
|||
|
|
conf = float(box.conf[0].tolist())
|
|||
|
|
confidences.append(conf)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[警告] 处理图像失败 {img_path}: {e}")
|
|||
|
|
|
|||
|
|
# 打印统计信息
|
|||
|
|
print("\n" + "="*50)
|
|||
|
|
print("推理统计:")
|
|||
|
|
print("="*50)
|
|||
|
|
print(f"总图像数: {total_images}")
|
|||
|
|
print(f"总检测数: {total_detections}")
|
|||
|
|
print(f"平均检测数: {total_detections/total_images:.2f}")
|
|||
|
|
|
|||
|
|
if confidences:
|
|||
|
|
confidences = np.array(confidences)
|
|||
|
|
print(f"平均置信度: {confidences.mean():.4f}")
|
|||
|
|
print(f"最高置信度: {confidences.max():.4f}")
|
|||
|
|
print(f"最低置信度: {confidences.min():.4f}")
|
|||
|
|
|
|||
|
|
print("="*50)
|
|||
|
|
print("[注意] 这是仅推理模式，无法计算 mAP 等精确指标")
|
|||
|
|
print(" 如需精确评估，请确保有对应的标签文件")
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_labelme_annotations(model, source_path: str, args):
|
|||
|
|
"""
|
|||
|
|
对图像进行检测并生成LabelMe格式的JSON标注文件（保存到 runs/<project>/<name>/），
|
|||
|
|
可视化结果由 Ultralytics 的 predict 保存（与 predict 模式一致）。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model: YOLO模型
|
|||
|
|
source_path: 图像目录路径
|
|||
|
|
args: 命令行参数
|
|||
|
|
"""
|
|||
|
|
import json
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
src = Path(source_path)
|
|||
|
|
if not src.exists():
|
|||
|
|
print(f"[错误] 源路径不存在: {src}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 收集所有图像文件
|
|||
|
|
if src.is_dir():
|
|||
|
|
image_files = list_images_in_dir(src, recursive=args.recursive)
|
|||
|
|
if not image_files:
|
|||
|
|
print(f"[错误] 目录下未找到图像: {src}")
|
|||
|
|
return
|
|||
|
|
else:
|
|||
|
|
image_files = [src]
|
|||
|
|
|
|||
|
|
print(f"[标注] 找到 {len(image_files)} 个图像文件")
|
|||
|
|
|
|||
|
|
# 创建输出目录 runs/<project>/<name>/
|
|||
|
|
project_dir = Path(args.project) #/ (args.name or "annotations")
|
|||
|
|
project_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
output_dir = Path(args.project) / (args.name or "annotations")
|
|||
|
|
|
|||
|
|
# 流式推理并即时写出标注，避免一次性占用大量显存/内存
|
|||
|
|
print(f"[标注] 流式预测并保存可视化到: {output_dir}")
|
|||
|
|
processed_count = 0
|
|||
|
|
failed_count = 0
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
for res in model.predict(
|
|||
|
|
source=[str(p) for p in image_files],
|
|||
|
|
imgsz=args.imgsz,
|
|||
|
|
conf=args.conf,
|
|||
|
|
iou=args.iou,
|
|||
|
|
device=args.device if args.device else None,
|
|||
|
|
verbose=False,
|
|||
|
|
save=True,
|
|||
|
|
project=args.project,
|
|||
|
|
name=args.name or "annotations",
|
|||
|
|
stream=True,
|
|||
|
|
):
|
|||
|
|
try:
|
|||
|
|
# 结果对象自带路径与原图尺寸
|
|||
|
|
img_path = Path(getattr(res, 'path', ''))
|
|||
|
|
if not img_path:
|
|||
|
|
# 回退：无法获得路径时跳过
|
|||
|
|
failed_count += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if hasattr(res, 'orig_shape') and res.orig_shape is not None:
|
|||
|
|
img_height, img_width = int(res.orig_shape[0]), int(res.orig_shape[1])
|
|||
|
|
else:
|
|||
|
|
original_img = cv2.imread(str(img_path))
|
|||
|
|
if original_img is None:
|
|||
|
|
print(f"[错误] 无法读取图像: {img_path}")
|
|||
|
|
failed_count += 1
|
|||
|
|
continue
|
|||
|
|
img_height, img_width = original_img.shape[:2]
|
|||
|
|
|
|||
|
|
labelme_data = {
|
|||
|
|
"version": "5.8.3",
|
|||
|
|
"flags": {},
|
|||
|
|
"shapes": [],
|
|||
|
|
"imagePath": img_path.name,
|
|||
|
|
"imageData": None,
|
|||
|
|
"imageHeight": img_height,
|
|||
|
|
"imageWidth": img_width
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if hasattr(res, 'boxes') and res.boxes is not None:
|
|||
|
|
for box in res.boxes:
|
|||
|
|
x1, y1, x2, y2 = [float(v) for v in box.xyxy[0].tolist()]
|
|||
|
|
conf = float(box.conf[0].tolist())
|
|||
|
|
cls = int(box.cls[0].tolist())
|
|||
|
|
shape = {
|
|||
|
|
"label": "鱿鱼",
|
|||
|
|
"points": [[x1, y1], [x2, y2]],
|
|||
|
|
"group_id": None,
|
|||
|
|
"description": f"confidence: {conf:.3f}",
|
|||
|
|
"shape_type": "rectangle",
|
|||
|
|
"flags": {},
|
|||
|
|
"mask": None
|
|||
|
|
}
|
|||
|
|
labelme_data["shapes"].append(shape)
|
|||
|
|
|
|||
|
|
json_filename = img_path.stem + ".json"
|
|||
|
|
json_path = output_dir / json_filename
|
|||
|
|
with open(json_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(labelme_data, f, ensure_ascii=False, indent=2)
|
|||
|
|
|
|||
|
|
print(f"[成功] {img_path.name} -> {json_filename} (检测到 {len(labelme_data['shapes'])} 个目标)")
|
|||
|
|
processed_count += 1
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[错误] 处理图像失败 {getattr(res, 'path', 'unknown')}: {e}")
|
|||
|
|
failed_count += 1
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[错误] 批处理预测失败: {e}")
|
|||
|
|
|
|||
|
|
# 输出统计信息
|
|||
|
|
print("\n" + "="*60)
|
|||
|
|
print("标注生成完成:")
|
|||
|
|
print("="*60)
|
|||
|
|
print(f"成功处理: {processed_count} 个图像")
|
|||
|
|
print(f"处理失败: {failed_count} 个图像")
|
|||
|
|
print(f"输出目录: {output_dir}")
|
|||
|
|
print("生成文件:")
|
|||
|
|
print(" - JSON标注文件: *.json")
|
|||
|
|
print(" - 预测可视化: 由 Ultralytics 自动保存 (与 predict 模式一致)")
|
|||
|
|
print("="*60)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
args = parse_args()
|
|||
|
|
try:
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
except Exception as e:
|
|||
|
|
print("[错误] 未找到 ultralytics，请先: pip install ultralytics")
|
|||
|
|
print(e)
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 检查matplotlib兼容性
|
|||
|
|
try:
|
|||
|
|
import matplotlib
|
|||
|
|
import numpy as np
|
|||
|
|
# 测试matplotlib和numpy的兼容性
|
|||
|
|
matplotlib.use('Agg') # 使用非交互式后端
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
print(f"[信息] Matplotlib版本: {matplotlib.__version__}")
|
|||
|
|
print(f"[信息] NumPy版本: {np.__version__}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[警告] Matplotlib/NumPy兼容性问题: {e}")
|
|||
|
|
print("[建议] 尝试升级matplotlib: pip install matplotlib --upgrade")
|
|||
|
|
# 继续执行，但可能会在后续遇到问题
|
|||
|
|
|
|||
|
|
model = YOLO(args.weights)
|
|||
|
|
|
|||
|
|
if args.mode == "validate":
|
|||
|
|
if not args.data:
|
|||
|
|
print("[错误] validate 模式需要 --data 指向 dataset.yaml")
|
|||
|
|
sys.exit(1)
|
|||
|
|
metrics = model.val(
|
|||
|
|
data=args.data,
|
|||
|
|
imgsz=args.imgsz,
|
|||
|
|
device=args.device if args.device else None,
|
|||
|
|
project=args.project,
|
|||
|
|
name=args.name or "val",
|
|||
|
|
iou=args.iou,
|
|||
|
|
verbose=True,
|
|||
|
|
)
|
|||
|
|
# metrics 包含各项指标，可按需打印
|
|||
|
|
print("评测完成。mAP50:", getattr(metrics, 'box', None).map50 if hasattr(metrics, 'box') else None)
|
|||
|
|
elif args.mode == "test":
|
|||
|
|
if not args.source:
|
|||
|
|
print("[错误] test 模式需要 --source 输入")
|
|||
|
|
sys.exit(1)
|
|||
|
|
metrics = evaluate_on_source(model, args.source, args)
|
|||
|
|
if metrics:
|
|||
|
|
print("测试完成。")
|
|||
|
|
else:
|
|||
|
|
print("测试完成（仅推理模式）。")
|
|||
|
|
elif args.mode == "annotation":
|
|||
|
|
if not args.source:
|
|||
|
|
print("[错误] annotation 模式需要 --source 输入")
|
|||
|
|
sys.exit(1)
|
|||
|
|
generate_labelme_annotations(model, args.source, args)
|
|||
|
|
print("标注生成完成。")
|
|||
|
|
else: # predict mode
|
|||
|
|
if not args.source:
|
|||
|
|
print("[错误] predict 模式需要 --source 输入")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
src = Path(args.source)
|
|||
|
|
sources = None
|
|||
|
|
if src.is_dir():
|
|||
|
|
# 严格限制在该目录（默认非递归）
|
|||
|
|
sources = list_images_in_dir(src, recursive=args.recursive)
|
|||
|
|
if not sources:
|
|||
|
|
print(f"[错误] 目录下未找到图像: {src}")
|
|||
|
|
sys.exit(1)
|
|||
|
|
else:
|
|||
|
|
# 单个文件或其他流式输入，直接传递
|
|||
|
|
sources = [str(src)]
|
|||
|
|
|
|||
|
|
last_save_dir = None
|
|||
|
|
# 使用流式推理避免一次性将所有结果保存在内存/显存中
|
|||
|
|
for res in model.predict(
|
|||
|
|
source=str(src) if src.is_dir() else [str(src)],
|
|||
|
|
imgsz=args.imgsz,
|
|||
|
|
conf=args.conf,
|
|||
|
|
iou=args.iou,
|
|||
|
|
device=args.device if args.device else None,
|
|||
|
|
project=args.project,
|
|||
|
|
name=args.name or "pred",
|
|||
|
|
save=True,
|
|||
|
|
save_txt=args.save_txt,
|
|||
|
|
save_conf=args.save_conf,
|
|||
|
|
verbose=True,
|
|||
|
|
stream=True,
|
|||
|
|
):
|
|||
|
|
last_save_dir = getattr(res, 'save_dir', None)
|
|||
|
|
print("推理完成。结果目录:", last_save_dir or args.project)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|