347 lines
10 KiB
Python
347 lines
10 KiB
Python
"""
|
||
YOLO数字识别 - 基础版本
|
||
|
||
功能说明:
|
||
使用训练好的YOLO模型识别图片中的4位阿拉伯数字。
|
||
这是基础版本，提供简单的数字检测和识别功能。
|
||
|
||
主要特性:
|
||
- 批量处理图片文件夹
|
||
- 支持自定义置信度阈值
|
||
- 从左到右排序数字
|
||
- 生成可视化结果（可选）
|
||
- 输出识别结果到文本文件
|
||
|
||
算法流程:
|
||
1. 加载YOLO模型
|
||
2. 对每张图片进行目标检测
|
||
3. 提取检测到的数字（0-9）
|
||
4. 按x坐标从左到右排序
|
||
5. 组合成完整数字串
|
||
|
||
适用场景:
|
||
- 快速测试模型效果
|
||
- 简单的数字识别任务
|
||
- 作为改进版的基准对比
|
||
|
||
注意事项:
|
||
- 不包含智能过滤，可能识别出非4位数字
|
||
- 对于复杂场景建议使用 predict_digits_improved.py
|
||
|
||
使用示例:
|
||
# 基础使用
|
||
python scripts/predict_digits.py
|
||
|
||
# 自定义参数
|
||
python scripts/predict_digits.py \
|
||
--model runs/digit_yolo/exp1/weights/best.pt \
|
||
--source valid \
|
||
--conf 0.25 \
|
||
--save-vis
|
||
|
||
# 高清识别
|
||
python scripts/predict_digits.py --imgsz 640 --conf 0.2
|
||
|
||
作者: Gavin Chan
|
||
版本: 1.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
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
"""
|
||
解析命令行参数
|
||
|
||
Returns:
|
||
argparse.Namespace: 包含所有命令行参数的对象
|
||
- model: YOLO模型文件路径
|
||
- source: 待识别图片的文件夹路径
|
||
- conf: 置信度阈值（0-1之间）
|
||
- imgsz: 输入图片尺寸
|
||
- 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.25,
|
||
help="置信度阈值"
|
||
)
|
||
parser.add_argument(
|
||
"--imgsz",
|
||
type=int,
|
||
default=320,
|
||
help="输入图片大小"
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
type=Path,
|
||
default=Path("results/predictions.txt"),
|
||
help="输出结果文件路径"
|
||
)
|
||
parser.add_argument(
|
||
"--save-vis",
|
||
action="store_true",
|
||
help="是否保存可视化结果"
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def extract_digits_from_predictions(results, img_width: int) -> str:
|
||
"""
|
||
从YOLO预测结果中提取数字并按位置排序
|
||
|
||
处理流程:
|
||
1. 遍历所有检测框
|
||
2. 提取边界框的x坐标中心点
|
||
3. 获取每个检测框的类别（0-9）和置信度
|
||
4. 按x坐标从左到右排序
|
||
5. 组合成完整的数字字符串
|
||
|
||
Args:
|
||
results: YOLO模型的预测结果对象
|
||
- results.boxes: 检测框信息
|
||
- results.boxes.xyxy: 边界框坐标 [x1, y1, x2, y2]
|
||
- results.boxes.cls: 类别ID（0-9对应数字0-9）
|
||
- results.boxes.conf: 置信度分数
|
||
img_width: 图片宽度（像素），用于坐标归一化（当前版本未使用）
|
||
|
||
Returns:
|
||
str: 识别出的数字字符串，如 "1234"，可能不足或超过4位
|
||
|
||
示例:
|
||
>>> results = model.predict("image.jpg")[0]
|
||
>>> digits = extract_digits_from_predictions(results, 640)
|
||
>>> print(digits) # "3809"
|
||
"""
|
||
# 提取检测框和类别
|
||
detections: List[Tuple[float, int]] = [] # (x_center, digit_class)
|
||
|
||
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()
|
||
x_center = (box[0] + box[2]) / 2
|
||
|
||
# 获取类别（数字0-9）
|
||
cls = int(boxes.cls[i].cpu().numpy())
|
||
|
||
# 获取置信度
|
||
conf = float(boxes.conf[i].cpu().numpy())
|
||
|
||
detections.append((x_center, cls, conf))
|
||
|
||
# 按照x坐标从左到右排序
|
||
detections.sort(key=lambda x: x[0])
|
||
|
||
# 提取数字
|
||
digits = [str(det[1]) for det in detections]
|
||
|
||
# 组合成4位数字字符串
|
||
result = "".join(digits)
|
||
|
||
return result
|
||
|
||
|
||
def predict_single_image(model: YOLO, image_path: Path, conf: float, imgsz: int) -> Tuple[str, float]:
|
||
"""
|
||
预测单张图片中的数字
|
||
|
||
处理流程:
|
||
1. 使用OpenCV读取图片获取尺寸信息
|
||
2. 调用YOLO模型进行目标检测
|
||
3. 提取并排序检测到的数字
|
||
4. 计算平均置信度作为质量指标
|
||
|
||
Args:
|
||
model (YOLO): 已加载的YOLO模型对象
|
||
image_path (Path): 图片文件的完整路径
|
||
conf (float): 置信度阈值（0-1），低于此值的检测将被过滤
|
||
imgsz (int): 模型输入图片大小，如320或640
|
||
|
||
Returns:
|
||
Tuple[str, float]: 二元组
|
||
- str: 识别出的数字字符串，如"1234"或"567"（可能不足4位）
|
||
- float: 所有检测框的平均置信度，范围0-1
|
||
|
||
异常处理:
|
||
- 如果图片无法读取，返回 ("", 0.0) 并打印警告
|
||
- 如果没有检测到任何数字，返回 ("", 0.0)
|
||
|
||
示例:
|
||
>>> model = YOLO("best.pt")
|
||
>>> digits, conf = predict_single_image(model, Path("test.jpg"), 0.25, 320)
|
||
>>> print(f"识别结果: {digits}, 置信度: {conf:.3f}")
|
||
识别结果: 3809, 置信度: 0.584
|
||
"""
|
||
# 读取图片获取宽度
|
||
img = cv2.imread(str(image_path))
|
||
if img is None:
|
||
print(f"警告：无法读取图片 {image_path}")
|
||
return "", 0.0
|
||
|
||
img_height, img_width = img.shape[:2]
|
||
|
||
# 进行预测
|
||
results = model.predict(
|
||
source=str(image_path),
|
||
conf=conf,
|
||
imgsz=imgsz,
|
||
verbose=False
|
||
)[0]
|
||
|
||
# 提取数字
|
||
digits = extract_digits_from_predictions(results, img_width)
|
||
|
||
# 计算平均置信度
|
||
avg_conf = 0.0
|
||
if results.boxes is not None and len(results.boxes) > 0:
|
||
confs = results.boxes.conf.cpu().numpy()
|
||
avg_conf = float(confs.mean())
|
||
|
||
return digits, avg_conf
|
||
|
||
|
||
def main() -> None:
|
||
"""
|
||
主函数：执行批量数字识别流程
|
||
|
||
完整流程:
|
||
1. 解析命令行参数
|
||
2. 验证模型文件和图片目录是否存在
|
||
3. 加载YOLO模型
|
||
4. 遍历所有图片文件进行识别
|
||
5. 统计识别结果（正确率、置信度等）
|
||
6. 保存结果到文本文件
|
||
7. 可选：生成带标注的可视化图片
|
||
|
||
输出格式:
|
||
控制台输出:
|
||
- 每张图片的识别结果
|
||
- 统计信息（正确率等）
|
||
- 文件保存路径
|
||
|
||
文本文件（results/predictions.txt）:
|
||
文件名 识别结果 置信度 数字个数
|
||
YZM.jpeg 3809 0.584 4
|
||
...
|
||
|
||
异常处理:
|
||
- FileNotFoundError: 模型或图片目录不存在
|
||
- 其他异常会向上传播
|
||
|
||
注意:
|
||
- 需要预先安装 ultralytics 和 opencv-python
|
||
- 模型文件需要是训练好的 .pt 格式
|
||
- 支持的图片格式: .jpg, .jpeg, .png, .bmp
|
||
"""
|
||
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("-" * 80)
|
||
|
||
# 预测结果
|
||
results = []
|
||
|
||
for image_path in image_files:
|
||
digits, conf = predict_single_image(model, image_path, args.conf, args.imgsz)
|
||
|
||
# 检查是否识别出4位数字
|
||
if len(digits) != 4:
|
||
status = f"⚠️ 检测到 {len(digits)} 位数字"
|
||
else:
|
||
status = "✓"
|
||
|
||
result_line = f"{image_path.name:<20} -> {digits:<6} (置信度: {conf:.3f}) {status}"
|
||
print(result_line)
|
||
|
||
results.append({
|
||
"filename": image_path.name,
|
||
"digits": digits,
|
||
"confidence": conf,
|
||
"digit_count": len(digits)
|
||
})
|
||
|
||
print("-" * 80)
|
||
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数字个数\n")
|
||
for r in results:
|
||
f.write(f"{r['filename']}\t{r['digits']}\t{r['confidence']:.3f}\t{r['digit_count']}\n")
|
||
|
||
print(f"结果已保存到: {args.output}")
|
||
|
||
# 如果需要保存可视化结果
|
||
if args.save_vis:
|
||
print("\n生成可视化结果...")
|
||
output_dir = args.output.parent / "visualizations"
|
||
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()
|