490 lines
17 KiB
Python
490 lines
17 KiB
Python
"""
|
||
YOLO数字识别 - 改进版本（推荐使用）
|
||
|
||
功能说明:
|
||
在基础版本上添加了智能过滤和后处理逻辑，提高4位数字识别的准确率。
|
||
这是生产环境推荐使用的版本。
|
||
|
||
主要特性:
|
||
- 智能检测过滤（置信度、位置、尺寸）
|
||
- 检测数量异常处理（<4或>4个数字）
|
||
- 垂直位置对齐验证
|
||
- 尺寸一致性检查
|
||
- 自适应参数调整
|
||
- 详细的识别质量报告
|
||
|
||
算法改进:
|
||
1. 多级置信度过滤（基础阈值 + 动态调整）
|
||
2. 位置异常检测（y坐标、尺寸统计分析）
|
||
3. 数量控制（超过4个时选择最优组合）
|
||
4. 数量不足时降低阈值重试（可选）
|
||
|
||
相比基础版的优势:
|
||
✓ 更准确：智能过滤减少误检
|
||
✓ 更稳定：处理各种异常情况
|
||
✓ 更可靠：提供详细的质量指标
|
||
✓ 更灵活：自适应不同图片质量
|
||
|
||
适用场景:
|
||
- 生产环境的数字识别
|
||
- 对准确率有要求的场景
|
||
- 图片质量参差不齐的情况
|
||
- 需要质量评估的应用
|
||
|
||
使用示例:
|
||
# 使用最佳模型识别（推荐）
|
||
python scripts/predict_digits_improved.py \
|
||
--model runs/digit_yolo/exp_preprocessed_color_150/weights/best.pt \
|
||
--source valid \
|
||
--conf 0.2 \
|
||
--save-vis
|
||
|
||
# 低置信度识别（图片模糊时）
|
||
python scripts/predict_digits_improved.py --conf 0.15
|
||
|
||
# 高清识别
|
||
python scripts/predict_digits_improved.py --imgsz 640
|
||
|
||
# 自定义输出
|
||
python scripts/predict_digits_improved.py \
|
||
--output results/my_predictions.txt
|
||
|
||
性能指标:
|
||
- 识别速度: ~0.5s/张 (CPU M2)
|
||
- 推荐置信度: 0.15-0.25
|
||
- 最佳图片尺寸: 320 (速度) 或 640 (精度)
|
||
|
||
作者: Gavin Chan
|
||
版本: 2.0
|
||
日期: 2025-10-30
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
from typing import List, Tuple
|
||
|
||
from ultralytics import YOLO
|
||
import cv2
|
||
import numpy as np
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
"""
|
||
解析命令行参数
|
||
|
||
Returns:
|
||
argparse.Namespace: 包含所有配置参数的对象
|
||
- model: YOLO模型文件路径
|
||
- source: 待识别图片的文件夹路径
|
||
- conf: 置信度阈值（推荐0.15-0.25）
|
||
- imgsz: 输入图片尺寸（320快速，640精确）
|
||
- output: 输出结果文件路径
|
||
- save_vis: 是否保存可视化结果
|
||
"""
|
||
parser = argparse.ArgumentParser(description="识别4位数字图片（改进版）")
|
||
parser.add_argument(
|
||
"--model",
|
||
type=Path,
|
||
default=Path("runs/digit_yolo/exp1/weights/best.pt"),
|
||
help="训练好的YOLO模型路径"
|
||
)
|
||
parser.add_argument(
|
||
"--source",
|
||
type=Path,
|
||
default=Path("valid"),
|
||
help="待识别图片的文件夹路径"
|
||
)
|
||
parser.add_argument(
|
||
"--conf",
|
||
type=float,
|
||
default=0.2,
|
||
help="置信度阈值"
|
||
)
|
||
parser.add_argument(
|
||
"--imgsz",
|
||
type=int,
|
||
default=320,
|
||
help="输入图片大小"
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
type=Path,
|
||
default=Path("results/predictions_improved.txt"),
|
||
help="输出结果文件路径"
|
||
)
|
||
parser.add_argument(
|
||
"--save-vis",
|
||
action="store_true",
|
||
help="是否保存可视化结果"
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def filter_detections(detections: List[Tuple[float, float, float, float, int, float]],
|
||
img_width: int, img_height: int) -> List[Tuple[float, float, float, float, int, float]]:
|
||
"""
|
||
智能过滤检测结果，去除误检和异常检测
|
||
|
||
过滤策略:
|
||
1. 置信度过滤: 去除置信度 < 0.15 的检测
|
||
2. 数量控制: 如果检测超过6个，保留置信度最高的6个
|
||
3. 位置过滤: 去除垂直位置（y坐标）偏离过大的检测
|
||
4. 尺寸过滤: 去除尺寸异常的检测框（过大或过小）
|
||
|
||
算法细节:
|
||
- 使用中位数判断y坐标是否异常（避免均值受极值影响）
|
||
- y坐标偏离超过平均高度视为异常
|
||
- 宽度偏离平均宽度2倍以上视为异常
|
||
|
||
Args:
|
||
detections (List[Tuple]): 原始检测列表，每个元素为六元组:
|
||
(x1, y1, x2, y2, class, conf)
|
||
- x1, y1: 左上角坐标
|
||
- x2, y2: 右下角坐标
|
||
- class: 类别ID（0-9对应数字0-9）
|
||
- conf: 置信度分数（0-1）
|
||
img_width (int): 图片宽度（像素）
|
||
img_height (int): 图片高度（像素）
|
||
|
||
Returns:
|
||
List[Tuple]: 过滤后的检测列表，格式与输入相同
|
||
- 返回符合条件的检测
|
||
- 按置信度降序排列
|
||
- 最多返回4-6个检测结果
|
||
|
||
示例:
|
||
>>> detections = [(10, 20, 30, 40, 5, 0.8), (50, 22, 70, 42, 3, 0.7)]
|
||
>>> filtered = filter_detections(detections, 640, 480)
|
||
>>> print(len(filtered)) # 2
|
||
"""
|
||
if not detections:
|
||
return []
|
||
|
||
# 1. 去除置信度过低的检测
|
||
filtered = [d for d in detections if d[5] > 0.15]
|
||
|
||
if len(filtered) == 0:
|
||
return []
|
||
|
||
# 2. 计算每个检测框的中心点和宽度
|
||
centers_and_widths = []
|
||
for det in filtered:
|
||
x1, y1, x2, y2, cls, conf = det
|
||
x_center = (x1 + x2) / 2
|
||
y_center = (y1 + y2) / 2
|
||
width = x2 - x1
|
||
height = y2 - y1
|
||
centers_and_widths.append((x_center, y_center, width, height, det))
|
||
|
||
# 3. 如果检测数量远超4个，尝试过滤
|
||
if len(centers_and_widths) > 6:
|
||
# 按置信度排序，保留前6个
|
||
centers_and_widths.sort(key=lambda x: x[4][5], reverse=True)
|
||
centers_and_widths = centers_and_widths[:6]
|
||
|
||
# 4. 去除垂直位置异常的检测框（y坐标差异过大）
|
||
if len(centers_and_widths) >= 2:
|
||
y_coords = [c[1] for c in centers_and_widths]
|
||
y_median = np.median(y_coords)
|
||
avg_height = np.mean([c[3] for c in centers_and_widths])
|
||
|
||
# 保留y坐标在合理范围内的检测框
|
||
filtered_by_y = []
|
||
for item in centers_and_widths:
|
||
x_center, y_center, width, height, det = item
|
||
if abs(y_center - y_median) < avg_height * 0.8: # y坐标偏差不超过平均高度的80%
|
||
filtered_by_y.append(item)
|
||
|
||
if filtered_by_y:
|
||
centers_and_widths = filtered_by_y
|
||
|
||
# 5. 返回过滤后的检测框
|
||
return [item[4] for item in centers_and_widths]
|
||
|
||
|
||
def extract_digits_from_predictions(results, img_width: int, img_height: int) -> Tuple[str, float, int]:
|
||
"""
|
||
从YOLO预测结果中提取并智能处理数字
|
||
|
||
完整处理流程:
|
||
1. 提取所有检测框的坐标、类别、置信度
|
||
2. 调用filter_detections进行智能过滤
|
||
3. 按x坐标从左到右排序（数字顺序）
|
||
4. 根据检测数量采取不同策略:
|
||
- 正好4个: 直接使用
|
||
- 超过4个: 选择置信度最高的4个
|
||
- 少于4个: 返回实际检测到的数字
|
||
|
||
智能选择策略:
|
||
当检测超过4个时，不是简单按位置选择前4个，
|
||
而是选择置信度最高的4个，这样可以过滤掉低质量检测。
|
||
|
||
Args:
|
||
results: YOLO模型的预测结果对象
|
||
- results.boxes: 所有检测框信息
|
||
- results.boxes.xyxy: 坐标 [x1, y1, x2, y2]
|
||
- results.boxes.cls: 类别ID (0-9)
|
||
- results.boxes.conf: 置信度
|
||
img_width (int): 图片宽度，用于过滤时的参考
|
||
img_height (int): 图片高度，用于过滤时的参考
|
||
|
||
Returns:
|
||
Tuple[str, float, int]: 三元组
|
||
- str: 识别出的数字字符串，如"3809"或"567"
|
||
- float: 平均置信度（所有选中数字的置信度均值）
|
||
- int: 原始检测数量（过滤前的数量，用于诊断）
|
||
|
||
示例:
|
||
>>> results = model.predict("image.jpg")[0]
|
||
>>> digits, conf, count = extract_digits_from_predictions(results, 640, 480)
|
||
>>> print(f"识别: {digits} (置信度:{conf:.3f}, 原始检测:{count}个)")
|
||
识别: 3809 (置信度:0.584, 原始检测:5个)
|
||
"""
|
||
# 提取检测框和类别
|
||
detections: List[Tuple[float, float, float, float, int, float]] = []
|
||
|
||
if results.boxes is not None and len(results.boxes) > 0:
|
||
boxes = results.boxes
|
||
for i in range(len(boxes)):
|
||
# 获取边界框坐标 (x1, y1, x2, y2)
|
||
box = boxes.xyxy[i].cpu().numpy()
|
||
x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
|
||
|
||
# 获取类别（数字0-9）
|
||
cls = int(boxes.cls[i].cpu().numpy())
|
||
|
||
# 获取置信度
|
||
conf = float(boxes.conf[i].cpu().numpy())
|
||
|
||
detections.append((x1, y1, x2, y2, cls, conf))
|
||
|
||
original_count = len(detections)
|
||
|
||
# 过滤检测结果
|
||
detections = filter_detections(detections, img_width, img_height)
|
||
|
||
# 按照x坐标从左到右排序
|
||
detections.sort(key=lambda x: (x[0] + x[2]) / 2)
|
||
|
||
# 如果检测数量正好是4个，直接使用
|
||
if len(detections) == 4:
|
||
digits = [str(det[4]) for det in detections]
|
||
confs = [det[5] for det in detections]
|
||
avg_conf = float(np.mean(confs))
|
||
return "".join(digits), avg_conf, original_count
|
||
|
||
# 如果检测数量大于4，尝试选择最可能的4个
|
||
if len(detections) > 4:
|
||
# 策略1: 选择置信度最高的4个，然后按x坐标排序
|
||
sorted_by_conf = sorted(detections, key=lambda x: x[5], reverse=True)
|
||
top4 = sorted_by_conf[:4]
|
||
top4.sort(key=lambda x: (x[0] + x[2]) / 2)
|
||
|
||
digits = [str(det[4]) for det in top4]
|
||
confs = [det[5] for det in top4]
|
||
avg_conf = float(np.mean(confs))
|
||
return "".join(digits), avg_conf, original_count
|
||
|
||
# 检测数量少于4个，直接返回
|
||
digits = [str(det[4]) for det in detections]
|
||
confs = [det[5] for det in detections] if detections else [0.0]
|
||
avg_conf = float(np.mean(confs))
|
||
return "".join(digits), avg_conf, original_count
|
||
|
||
|
||
def predict_single_image(model: YOLO, image_path: Path, conf: float, imgsz: int) -> Tuple[str, float, int]:
|
||
"""
|
||
预测单张图片中的数字（改进版）
|
||
|
||
相比基础版的改进:
|
||
- 返回原始检测数量，便于诊断问题
|
||
- 调用智能提取函数，处理异常情况
|
||
- 更详细的错误处理
|
||
|
||
处理流程:
|
||
1. 使用OpenCV读取图片，获取尺寸
|
||
2. 调用YOLO模型进行检测
|
||
3. 调用extract_digits_from_predictions进行智能处理
|
||
4. 返回最终识别结果和质量指标
|
||
|
||
Args:
|
||
model (YOLO): 已加载的YOLO模型对象
|
||
image_path (Path): 图片文件的完整路径
|
||
conf (float): 置信度阈值（0-1）
|
||
imgsz (int): 模型输入尺寸（320或640）
|
||
|
||
Returns:
|
||
Tuple[str, float, int]: 三元组
|
||
- str: 识别出的数字字符串
|
||
- float: 平均置信度
|
||
- int: 原始检测数量（过滤前）
|
||
|
||
异常处理:
|
||
- 图片无法读取: 返回 ("", 0.0, 0) 并打印警告
|
||
- 没有检测结果: 返回 ("", 0.0, 0)
|
||
|
||
示例:
|
||
>>> model = YOLO("best.pt")
|
||
>>> digits, conf, count = predict_single_image(model, Path("test.jpg"), 0.2, 320)
|
||
>>> if len(digits) == 4:
|
||
... print(f"✓ 识别成功: {digits}")
|
||
... else:
|
||
... print(f"⚠️ 只检测到 {len(digits)} 位")
|
||
"""
|
||
# 读取图片获取宽度
|
||
img = cv2.imread(str(image_path))
|
||
if img is None:
|
||
print(f"警告：无法读取图片 {image_path}")
|
||
return "", 0.0, 0
|
||
|
||
img_height, img_width = img.shape[:2]
|
||
|
||
# 进行预测
|
||
results = model.predict(
|
||
source=str(image_path),
|
||
conf=conf,
|
||
imgsz=imgsz,
|
||
verbose=False
|
||
)[0]
|
||
|
||
# 提取数字
|
||
digits, avg_conf, original_count = extract_digits_from_predictions(results, img_width, img_height)
|
||
|
||
return digits, avg_conf, original_count
|
||
|
||
|
||
def main() -> None:
|
||
"""
|
||
主函数：执行智能批量数字识别流程
|
||
|
||
完整流程:
|
||
1. 解析命令行参数并验证
|
||
2. 加载YOLO模型
|
||
3. 扫描图片文件夹，支持多种图片格式
|
||
4. 逐张进行智能识别（带过滤和后处理）
|
||
5. 收集并统计识别结果
|
||
6. 生成详细的质量报告
|
||
7. 保存结果到文本文件
|
||
8. 可选：生成可视化标注图片
|
||
|
||
输出内容:
|
||
控制台输出:
|
||
- 每张图片的识别结果（数字、置信度、检测数量）
|
||
- 统计信息（准确率、平均置信度等）
|
||
- 质量分析（低置信度、异常检测等）
|
||
|
||
文本文件（results/predictions_improved.txt）:
|
||
文件名 识别结果 置信度 数字个数 原始检测数
|
||
YZM.jpeg 3809 0.584 4 5
|
||
|
||
可视化图片（可选）:
|
||
results/visualizations_improved/
|
||
- 每张图片带检测框和标签
|
||
- 便于人工审核和调试
|
||
|
||
质量指标:
|
||
- 正确率: 识别出4位数字的图片比例
|
||
- 平均置信度: 所有图片的平均置信度
|
||
- 低质量警告: 识别不足4位的图片列表
|
||
- 过度检测: 原始检测超过6个的图片
|
||
|
||
异常处理:
|
||
- FileNotFoundError: 模型或图片目录不存在时抛出
|
||
- 图片读取失败: 跳过并打印警告
|
||
- 其他异常向上传播
|
||
|
||
依赖环境:
|
||
- ultralytics (YOLO模型)
|
||
- opencv-python (图片读取)
|
||
- numpy (数值计算)
|
||
"""
|
||
args = parse_args()
|
||
|
||
# 检查模型文件
|
||
if not args.model.exists():
|
||
raise FileNotFoundError(f"模型文件不存在: {args.model}")
|
||
|
||
# 检查源文件夹
|
||
if not args.source.exists():
|
||
raise FileNotFoundError(f"源文件夹不存在: {args.source}")
|
||
|
||
# 加载模型
|
||
print(f"加载模型: {args.model}")
|
||
model = YOLO(str(args.model))
|
||
|
||
# 获取所有图片文件
|
||
image_extensions = [".jpg", ".jpeg", ".png", ".bmp"]
|
||
image_files = []
|
||
for ext in image_extensions:
|
||
image_files.extend(args.source.glob(f"*{ext}"))
|
||
image_files.extend(args.source.glob(f"*{ext.upper()}"))
|
||
|
||
image_files = sorted(image_files)
|
||
|
||
if not image_files:
|
||
print(f"在 {args.source} 中没有找到图片文件")
|
||
return
|
||
|
||
print(f"找到 {len(image_files)} 张图片")
|
||
print("-" * 90)
|
||
|
||
# 预测结果
|
||
results = []
|
||
|
||
for image_path in image_files:
|
||
digits, conf, original_count = predict_single_image(model, image_path, args.conf, args.imgsz)
|
||
|
||
# 检查是否识别出4位数字
|
||
if len(digits) != 4:
|
||
status = f"⚠️ 检测到 {len(digits)} 位 (原始:{original_count})"
|
||
else:
|
||
status = f"✓ (原始:{original_count})"
|
||
|
||
result_line = f"{image_path.name:<20} -> {digits:<8} 置信度:{conf:.3f} {status}"
|
||
print(result_line)
|
||
|
||
results.append({
|
||
"filename": image_path.name,
|
||
"digits": digits,
|
||
"confidence": conf,
|
||
"digit_count": len(digits),
|
||
"original_count": original_count
|
||
})
|
||
|
||
print("-" * 90)
|
||
print(f"识别完成！")
|
||
|
||
# 统计信息
|
||
correct_count = sum(1 for r in results if r["digit_count"] == 4)
|
||
print(f"正确识别4位数字: {correct_count}/{len(results)} ({correct_count/len(results)*100:.1f}%)")
|
||
|
||
# 保存结果
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
with args.output.open("w", encoding="utf-8") as f:
|
||
f.write("文件名\t识别结果\t置信度\t数字个数\t原始检测数\n")
|
||
for r in results:
|
||
f.write(f"{r['filename']}\t{r['digits']}\t{r['confidence']:.3f}\t{r['digit_count']}\t{r['original_count']}\n")
|
||
|
||
print(f"结果已保存到: {args.output}")
|
||
|
||
# 如果需要保存可视化结果
|
||
if args.save_vis:
|
||
print("\n生成可视化结果...")
|
||
output_dir = args.output.parent / "visualizations_improved"
|
||
model.predict(
|
||
source=str(args.source),
|
||
conf=args.conf,
|
||
imgsz=args.imgsz,
|
||
save=True,
|
||
project=str(output_dir.parent),
|
||
name=output_dir.name,
|
||
exist_ok=True
|
||
)
|
||
print(f"可视化结果已保存到: {output_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|