491 lines
13 KiB
Python
491 lines
13 KiB
Python
"""
|
||
图片预处理工具 - 提升数字识别效果
|
||
|
||
功能说明:
|
||
对数字图片进行多种预处理，以提升YOLO模型的识别效果。
|
||
支持多种预处理方法，可单独使用或组合使用。
|
||
|
||
主要特性:
|
||
- 多种预处理方法（6种）
|
||
- 支持批量处理
|
||
- 可保持彩色或转为灰度
|
||
- 实时进度显示
|
||
- 可预览处理效果
|
||
- 自动创建输出目录
|
||
|
||
预处理方法详解:
|
||
1. auto (自动增强):
|
||
- 去噪 + 锐化
|
||
- 适合一般场景
|
||
|
||
2. clahe (对比度限制自适应直方图均衡化):
|
||
- 增强局部对比度
|
||
- 突出数字边缘
|
||
- 推荐用于低对比度图片 ⭐
|
||
|
||
3. binary (自适应二值化):
|
||
- 将图片转为黑白
|
||
- 适合文档类图片
|
||
- 可能丢失信息，谨慎使用
|
||
|
||
4. denoise (去噪):
|
||
- 去除图片噪点
|
||
- 保持边缘清晰
|
||
- 适合噪声较大的图片
|
||
|
||
5. sharpen (锐化):
|
||
- 增强边缘和细节
|
||
- 使数字更清晰
|
||
- 可能放大噪声
|
||
|
||
6. combined (组合方法):
|
||
- CLAHE + 去噪 + 锐化
|
||
- 综合效果最好
|
||
- 处理时间较长
|
||
|
||
重要提示:
|
||
- 训练和预测必须使用相同的预处理方法！
|
||
- 建议使用 --keep-color 保持彩色，避免训练/预测不一致
|
||
- clahe + keep-color 是推荐的最佳组合 ⭐
|
||
|
||
使用场景:
|
||
场景1: 预处理训练数据
|
||
python scripts/preprocess_images.py \
|
||
--input digit-validation/images \
|
||
--output digit-validation-processed \
|
||
--method clahe \
|
||
--keep-color
|
||
|
||
场景2: 预处理验证数据
|
||
python scripts/preprocess_images.py \
|
||
--input valid \
|
||
--output valid-processed \
|
||
--method clahe \
|
||
--keep-color
|
||
|
||
场景3: 预览效果（处理前3张）
|
||
python scripts/preprocess_images.py \
|
||
--input valid \
|
||
--output test-output \
|
||
--method clahe \
|
||
--show-preview
|
||
|
||
场景4: 测试不同方法
|
||
for method in auto clahe binary denoise sharpen combined; do
|
||
python scripts/preprocess_images.py \
|
||
--input valid \
|
||
--output valid-${method} \
|
||
--method ${method} \
|
||
--keep-color
|
||
done
|
||
|
||
输出:
|
||
- 处理后的图片（与输入文件名相同）
|
||
- 图片质量分析报告
|
||
- 处理统计信息
|
||
|
||
性能:
|
||
- 处理速度: ~0.1s/张（CPU）
|
||
- 支持格式: JPG, JPEG, PNG, BMP
|
||
- 保持原图尺寸不变
|
||
|
||
依赖环境:
|
||
- opencv-python >= 4.0.0
|
||
- numpy
|
||
- tqdm（进度条）
|
||
|
||
作者: Gavin Chan
|
||
版本: 1.0
|
||
日期: 2025-10-30
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
from typing import Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="预处理数字图片以提升识别效果")
|
||
parser.add_argument(
|
||
"--input",
|
||
type=Path,
|
||
required=True,
|
||
help="输入图片文件夹路径"
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
type=Path,
|
||
required=True,
|
||
help="输出图片文件夹路径"
|
||
)
|
||
parser.add_argument(
|
||
"--method",
|
||
type=str,
|
||
default="auto",
|
||
choices=["auto", "clahe", "binary", "denoise", "sharpen", "combined"],
|
||
help="预处理方法"
|
||
)
|
||
parser.add_argument(
|
||
"--keep-color",
|
||
action="store_true",
|
||
help="保持彩色图片（默认转为灰度）"
|
||
)
|
||
parser.add_argument(
|
||
"--show-preview",
|
||
action="store_true",
|
||
help="显示处理前后对比（仅处理前3张）"
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def enhance_contrast_clahe(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
使用CLAHE（自适应直方图均衡化）增强对比度
|
||
"""
|
||
if len(image.shape) == 3:
|
||
# 彩色图片：在LAB空间处理
|
||
lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
||
l, a, b = cv2.split(lab)
|
||
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
||
l = clahe.apply(l)
|
||
lab = cv2.merge([l, a, b])
|
||
return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
|
||
else:
|
||
# 灰度图片：直接处理
|
||
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
||
return clahe.apply(image)
|
||
|
||
|
||
def denoise_image(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
去噪处理
|
||
"""
|
||
if len(image.shape) == 3:
|
||
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
|
||
else:
|
||
return cv2.fastNlMeansDenoising(image, None, 10, 7, 21)
|
||
|
||
|
||
def sharpen_image(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
锐化图片
|
||
"""
|
||
kernel = np.array([[-1, -1, -1],
|
||
[-1, 9, -1],
|
||
[-1, -1, -1]])
|
||
return cv2.filter2D(image, -1, kernel)
|
||
|
||
|
||
def adaptive_binarization(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
自适应二值化
|
||
"""
|
||
if len(image.shape) == 3:
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
else:
|
||
gray = image
|
||
|
||
# 自适应阈值
|
||
binary = cv2.adaptiveThreshold(
|
||
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||
cv2.THRESH_BINARY, 11, 2
|
||
)
|
||
return binary
|
||
|
||
|
||
def morphology_operations(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
形态学操作：闭运算和开运算
|
||
"""
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
|
||
|
||
# 闭运算：填充小孔
|
||
closing = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel)
|
||
|
||
# 开运算：去除小噪点
|
||
opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
|
||
|
||
return opening
|
||
|
||
|
||
def preprocess_auto(image: np.ndarray, keep_color: bool = False) -> np.ndarray:
|
||
"""
|
||
自动预处理（推荐）
|
||
"""
|
||
# 1. 去噪
|
||
denoised = denoise_image(image)
|
||
|
||
# 2. 对比度增强
|
||
enhanced = enhance_contrast_clahe(denoised)
|
||
|
||
if keep_color:
|
||
# 保持彩色
|
||
# 3. 轻微锐化
|
||
sharpened = sharpen_image(enhanced)
|
||
return sharpened
|
||
else:
|
||
# 转为灰度
|
||
if len(enhanced.shape) == 3:
|
||
gray = cv2.cvtColor(enhanced, cv2.COLOR_BGR2GRAY)
|
||
else:
|
||
gray = enhanced
|
||
|
||
# 3. 轻微锐化
|
||
sharpened = sharpen_image(gray)
|
||
|
||
return sharpened
|
||
|
||
|
||
def preprocess_combined(image: np.ndarray) -> np.ndarray:
|
||
"""
|
||
组合预处理（强化版）
|
||
"""
|
||
# 1. 去噪
|
||
denoised = denoise_image(image)
|
||
|
||
# 2. 转灰度
|
||
if len(denoised.shape) == 3:
|
||
gray = cv2.cvtColor(denoised, cv2.COLOR_BGR2GRAY)
|
||
else:
|
||
gray = denoised
|
||
|
||
# 3. 对比度增强
|
||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||
enhanced = clahe.apply(gray)
|
||
|
||
# 4. 自适应二值化
|
||
binary = cv2.adaptiveThreshold(
|
||
enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||
cv2.THRESH_BINARY, 11, 2
|
||
)
|
||
|
||
# 5. 形态学操作
|
||
result = morphology_operations(binary)
|
||
|
||
return result
|
||
|
||
|
||
def preprocess_image(
|
||
image: np.ndarray,
|
||
method: str = "auto",
|
||
keep_color: bool = False
|
||
) -> np.ndarray:
|
||
"""
|
||
根据指定方法预处理图片
|
||
"""
|
||
if method == "auto":
|
||
return preprocess_auto(image, keep_color)
|
||
elif method == "clahe":
|
||
return enhance_contrast_clahe(image)
|
||
elif method == "binary":
|
||
return adaptive_binarization(image)
|
||
elif method == "denoise":
|
||
return denoise_image(image)
|
||
elif method == "sharpen":
|
||
return sharpen_image(image)
|
||
elif method == "combined":
|
||
return preprocess_combined(image)
|
||
else:
|
||
return image
|
||
|
||
|
||
def process_folder(
|
||
input_dir: Path,
|
||
output_dir: Path,
|
||
method: str = "auto",
|
||
keep_color: bool = False,
|
||
show_preview: bool = False
|
||
) -> None:
|
||
"""
|
||
处理文件夹中的所有图片
|
||
"""
|
||
# 创建输出目录
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 获取所有图片文件
|
||
image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.JPG", "*.JPEG", "*.PNG", "*.BMP"]
|
||
image_files = []
|
||
for ext in image_extensions:
|
||
image_files.extend(input_dir.glob(ext))
|
||
|
||
image_files = sorted(image_files)
|
||
|
||
if not image_files:
|
||
print(f"在 {input_dir} 中没有找到图片文件")
|
||
return
|
||
|
||
print(f"找到 {len(image_files)} 张图片")
|
||
print(f"预处理方法: {method}")
|
||
print(f"保持彩色: {keep_color}")
|
||
print("-" * 80)
|
||
|
||
preview_count = 0
|
||
|
||
for image_path in tqdm(image_files, desc="预处理图片"):
|
||
# 读取图片
|
||
image = cv2.imread(str(image_path))
|
||
if image is None:
|
||
print(f"警告：无法读取图片 {image_path}")
|
||
continue
|
||
|
||
# 预处理
|
||
processed = preprocess_image(image, method, keep_color)
|
||
|
||
# 保存处理后的图片
|
||
output_path = output_dir / image_path.name
|
||
cv2.imwrite(str(output_path), processed)
|
||
|
||
# 显示预览
|
||
if show_preview and preview_count < 3:
|
||
print(f"\n预览: {image_path.name}")
|
||
show_comparison(image, processed, image_path.name)
|
||
preview_count += 1
|
||
|
||
print(f"\n✓ 处理完成！输出目录: {output_dir}")
|
||
|
||
# 统计信息
|
||
print(f"\n处理统计:")
|
||
print(f" 输入图片: {len(image_files)}")
|
||
print(f" 输出图片: {len(list(output_dir.glob('*')))} ")
|
||
print(f" 预处理方法: {method}")
|
||
|
||
|
||
def show_comparison(original: np.ndarray, processed: np.ndarray, title: str) -> None:
|
||
"""
|
||
显示处理前后对比（需要图形界面）
|
||
"""
|
||
try:
|
||
import matplotlib.pyplot as plt
|
||
|
||
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
|
||
|
||
# 原图
|
||
if len(original.shape) == 3:
|
||
axes[0].imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
|
||
else:
|
||
axes[0].imshow(original, cmap='gray')
|
||
axes[0].set_title(f'原图 - {title}')
|
||
axes[0].axis('off')
|
||
|
||
# 处理后
|
||
if len(processed.shape) == 3:
|
||
axes[1].imshow(cv2.cvtColor(processed, cv2.COLOR_BGR2RGB))
|
||
else:
|
||
axes[1].imshow(processed, cmap='gray')
|
||
axes[1].set_title(f'处理后 - {title}')
|
||
axes[1].axis('off')
|
||
|
||
plt.tight_layout()
|
||
plt.show()
|
||
except ImportError:
|
||
print(" (matplotlib未安装，跳过预览)")
|
||
except Exception as e:
|
||
print(f" (预览失败: {e})")
|
||
|
||
|
||
def analyze_image_quality(input_dir: Path) -> None:
|
||
"""
|
||
分析图片质量并给出预处理建议
|
||
"""
|
||
image_files = list(input_dir.glob("*.jpg")) + list(input_dir.glob("*.jpeg")) + \
|
||
list(input_dir.glob("*.png")) + list(input_dir.glob("*.JPG")) + \
|
||
list(input_dir.glob("*.JPEG")) + list(input_dir.glob("*.PNG"))
|
||
|
||
if not image_files:
|
||
print("没有找到图片文件")
|
||
return
|
||
|
||
print(f"分析 {len(image_files)} 张图片的质量...")
|
||
print("-" * 80)
|
||
|
||
brightness_values = []
|
||
contrast_values = []
|
||
noise_levels = []
|
||
|
||
for img_path in image_files[:5]: # 分析前5张
|
||
img = cv2.imread(str(img_path))
|
||
if img is None:
|
||
continue
|
||
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
|
||
|
||
# 亮度
|
||
brightness = np.mean(gray)
|
||
brightness_values.append(brightness)
|
||
|
||
# 对比度（标准差）
|
||
contrast = np.std(gray)
|
||
contrast_values.append(contrast)
|
||
|
||
# 噪声估计（拉普拉斯方差）
|
||
laplacian = cv2.Laplacian(gray, cv2.CV_64F)
|
||
noise = laplacian.var()
|
||
noise_levels.append(noise)
|
||
|
||
avg_brightness = np.mean(brightness_values)
|
||
avg_contrast = np.mean(contrast_values)
|
||
avg_noise = np.mean(noise_levels)
|
||
|
||
print(f"平均亮度: {avg_brightness:.2f} (0-255)")
|
||
print(f"平均对比度: {avg_contrast:.2f}")
|
||
print(f"平均噪声水平: {avg_noise:.2f}")
|
||
print("-" * 80)
|
||
|
||
# 给出建议
|
||
print("\n预处理建议:")
|
||
if avg_brightness < 100:
|
||
print(" • 图片偏暗，建议使用 --method clahe 增强对比度")
|
||
elif avg_brightness > 180:
|
||
print(" • 图片偏亮，建议使用 --method clahe 增强对比度")
|
||
else:
|
||
print(" • 亮度正常")
|
||
|
||
if avg_contrast < 40:
|
||
print(" • 对比度较低，建议使用 --method clahe 或 combined")
|
||
else:
|
||
print(" • 对比度正常")
|
||
|
||
if avg_noise > 500:
|
||
print(" • 噪声较高，建议使用 --method denoise 或 combined")
|
||
else:
|
||
print(" • 噪声水平可接受")
|
||
|
||
print("\n推荐使用: --method auto （自动综合处理）")
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
|
||
# 检查输入目录
|
||
if not args.input.exists():
|
||
raise FileNotFoundError(f"输入目录不存在: {args.input}")
|
||
|
||
# 分析图片质量
|
||
print("=" * 80)
|
||
print("图片质量分析")
|
||
print("=" * 80)
|
||
analyze_image_quality(args.input)
|
||
print()
|
||
|
||
# 处理图片
|
||
print("=" * 80)
|
||
print("开始预处理")
|
||
print("=" * 80)
|
||
process_folder(
|
||
args.input,
|
||
args.output,
|
||
args.method,
|
||
args.keep_color,
|
||
args.show_preview
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|