基于 YOLO11 的古籍页面自动上色:从数据标注到 ONNX 部署的完整实战草稿,
目录
- 一、项目背景与效果
- 二、技术方案总览
- 三、环境准备
- 四、数据标注(LabelMe)
- 五、数据格式转换:LabelMe(JSON) → YOLO(TXT)
- 六、模型训练
- 七、ONNX 导出
- 八、推理部署:预处理 + 检测/分割封装类
- 九、上色主流程
- 十、常见问题与调优
一、项目背景与效果
古籍数字化过程中,大量珍贵资源只有黑白扫描件。项目目标是利用 AI 实现灰度图像 → 彩色版本的自动转换,拆成三个子任务:
|
子任务 |
内容 |
技术选型 |
|---|---|---|
|
文字精准着色 |
保留墨色风格,统一渲染为标准黑色 |
实例分割(需逐像素边界) |
|
印章智能识别 |
检测印章区域,还原朱砂红色彩 |
目标检测 + 模板贴图 |
|
边框线条修复 |
识别页面边框/表格线,渲染为红色 |
检测/分割规则着色 |
为什么印章用检测、文字用分割? 这是"检测 vs 分割怎么选"的教科书案例:印章着色只需知道位置和大小,贴一张标准红章图即可,用检测足够;文字着色要求颜色严格限制在笔画内部、不能溢出到纸张,必须逐像素分割。
最终效果:生成的彩色古籍图片可用于高质量印刷,还原纸张纹理与色彩细节。
二、技术方案总览
┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐
│ 01 数据标注 │ → │ 02 模型训练 │ → │ 03 迭代优化 │ → │ 04 端到端应用│ → │ 05 成果输出 │
└──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘
LabelMe手工 YOLO11n检测 模型预标注 Python脚本 彩色图片
标注少量图片 +YOLO11n-seg 人工修正扩充 批量自动处理 可用于印刷
分割训练 迭代训练
其中第 03 步是工业界标准打法(伪标签自训练):用初版模型对海量未标注图片预标注,人工只做审核修正后回流训练集。真实场景下古籍图片上百万张,不可能全量人工标注——先标几千张训初版模型,标注成本下降一个数量级。注意风险:修正不到位会把错误标签洗进数据集(模型自我强化错误),所以修正环节的质量控制比速度更重要。
三、环境准备
pip install ultralytics onnx onnxruntime opencv-python numpy pillow pyyaml tqdm pymupdf labelme
有 GPU 的机器建议同时安装 CUDA 版 onnxruntime:pip install onnxruntime-gpu,推理时自动走 GPU。
四、数据标注(LabelMe)
LabelMe 是一款开源图像标注工具(pip install labelme,终端输入 labelme 启动),支持两种本项目用到的标注模式:
|
模式 |
工具 |
用途 |
本项目 |
|---|---|---|---|
|
目标检测标注 |
Create RectBox(矩形框) |
定位物体大致区域 |
印章标注,标签名 |
|
实例分割标注 |
Create Polygons(多边形) |
精确勾勒目标轮廓 |
文字标注,标签名 |
印章标注五步走:打开图片 → 选 Create RectBox → 框选印章(紧贴边缘)→ 输入标签 seal → 保存为 JSON。
文字标注五步走:打开图片 → 选 Create Polygons → 沿文字列边缘点选多边形(每竖列单独一个多边形,顶点不必贴到每个笔画,沿外缘 6~10 个点即可,但宁可贴文字外侧、不要圈进过多背景)→ 输入标签 word → 保存 JSON。
标注规范(务必遵守):
- 标签命名全局统一(全小写、单一词),中英文/大小写一旦确定不要改,否则模型会当成两个类别;
- 一张图多个目标逐个标注,存进同一个 JSON;
- 效率技巧:LabelMe 会记住上次标签名,后续按回车复用。
五、数据格式转换:LabelMe(JSON) → YOLO(TXT)
5.1 两种格式的本质差异
|
|
LabelMe(JSON) |
YOLO(TXT) |
|---|---|---|
|
存储方式 |
基于像素的绝对坐标 |
归一化相对坐标(0~1) |
|
标注形式 |
多边形点序列 或 矩形框 (xmin,ymin)+(xmax,ymax) |
一行一个目标:检测 |
|
特点 |
直观易读,数值随图片尺寸变化,不能直接训练 |
与图片尺寸无关,YOLO 训练标准格式 |
转换三步:读 JSON 取 shapes → 坐标除以图像宽高归一化 → 检测转 (cx,cy,w,h)、分割直接写顶点序列,类别 ID 按 yaml 中 names 顺序。
5.2 检测格式转换脚本 labelme2yolo_det.py
功能完备的转换类:类别自动提取(字母序排序保证可复现)、8:2 划分训练/验证集(固定种子)、中文路径图像读取、越界裁剪与零尺寸过滤、自动生成 data.yaml。
import json
import os
import glob
import random
import shutil
from pathlib import Path
from tqdm import tqdm
import cv2
import numpy as np
class LabelMeToYOLODetectorConverter:
"""LabelMe格式转YOLO目标检测格式转换器
输出: (x_center, y_center, w, h) 归一化坐标 + 标准目录结构 + data.yaml
"""
def __init__(self, input_dir, output_dir, train_ratio=0.8, random_seed=42):
self.input_dir = Path(input_dir)
self.output_dir = Path(output_dir)
self.train_ratio = train_ratio
self.random_seed = random_seed
self.classes_dict = {}
random.seed(self.random_seed) # 固定种子:每次转换,同一批图永远进同一侧集合
def _read_image_with_chinese_path(self, image_path):
"""支持中文路径的图像读取(cv2.imread 在 Windows 下不支持中文路径,会返回 None)"""
try:
raw_data = np.fromfile(image_path, dtype=np.uint8) # 按字节流读取,绕过路径编码
image = cv2.imdecode(raw_data, cv2.IMREAD_COLOR) # 从内存字节解码
return image
except Exception as e:
print(f"读取图像失败 {image_path}: {e}")
return None
def _clip_coordinate_to_boundary(self, coordinate, max_value):
"""坐标裁剪到 [0, max_value-1],防止标注越界导致归一化值 > 1"""
return max(0, min(coordinate, max_value - 1))
def _get_bbox_from_points(self, points, image_width, image_height):
"""从点集(矩形/多边形通吃)计算边界框并归一化,返回 x_center, y_center, w, h"""
xs = [p[0] for p in points]
ys = [p[1] for p in points]
x1, y1, x2, y2 = min(xs), min(ys), max(xs), max(ys)
x1 = self._clip_coordinate_to_boundary(x1, image_width)
y1 = self._clip_coordinate_to_boundary(y1, image_height)
x2 = self._clip_coordinate_to_boundary(x2, image_width)
y2 = self._clip_coordinate_to_boundary(y2, image_height)
bbox_width, bbox_height = x2 - x1, y2 - y1
if bbox_width <= 0 or bbox_height <= 0: # 零尺寸过滤
return None
x_center = (x1 + x2) / 2 / image_width
y_center = (y1 + y2) / 2 / image_height
norm_width = bbox_width / image_width
norm_height = bbox_height / image_height
return [x_center, y_center, norm_width, norm_height]
def _extract_classes_from_annotations(self):
"""从所有JSON中提取类别集合,按字母序排序生成 id(保证多次运行结果一致)"""
classes_set = set()
json_files = glob.glob(str(self.input_dir / "*.json"))
for json_file in tqdm(json_files, desc="提取类别信息"):
try:
with open(json_file, "r", encoding="utf-8") as file:
annotation_data = json.load(file)
for shape in annotation_data.get("shapes", []):
label = shape.get("label", "").strip()
if label and label != "bad": # 排除空标签和异常标签
classes_set.add(label)
except Exception as e:
print(f"处理文件 {json_file} 时出错: {e}")
continue
sorted_classes = sorted(list(classes_set))
self.classes_dict = {c: i for i, c in enumerate(sorted_classes)}
print(f"检测到 {len(self.classes_dict)} 个类别: {self.classes_dict}")
return self.classes_dict
def _create_output_directories(self):
"""创建YOLO标准目录结构 images/{train,val} + labels/{train,val}"""
directories = [
self.output_dir / "images" / "train",
self.output_dir / "images" / "val",
self.output_dir / "labels" / "train",
self.output_dir / "labels" / "val",
]
for directory in directories:
directory.mkdir(parents=True, exist_ok=True)
print(f"创建目录: {directory}")
def _split_dataset_files(self, json_files):
"""80% 训练 / 20% 验证,随机打乱 + 固定种子可复现"""
shuffled_files = json_files.copy()
random.shuffle(shuffled_files)
split_index = int(len(shuffled_files) * self.train_ratio)
return shuffled_files[:split_index], shuffled_files[split_index:]
def _find_image_file(self, json_path):
"""LabelMe 保存的 JSON 与图像同名,按常见扩展名逐个尝试"""
json_path = Path(json_path)
for ext in [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]:
image_path = json_path.with_suffix(ext)
if image_path.exists():
return image_path
return None
def _convert_single_annotation(self, json_path, output_label_dir, output_image_dir):
"""转换单个文件:JSON -> txt标签 + 拷贝图像"""
try:
with open(json_path, "r", encoding="utf-8") as file:
annotation_data = json.load(file)
image_path = self._find_image_file(json_path)
if not image_path:
print(f"未找到图像文件: {json_path}")
return False
image = self._read_image_with_chinese_path(image_path)
if image is None:
return False
image_height, image_width = image.shape[:2]
txt_path = output_label_dir / (Path(json_path).stem + ".txt")
with open(txt_path, "w", encoding="utf-8") as txt_file:
for shape in annotation_data["shapes"]:
label = shape["label"].strip()
if label not in self.classes_dict:
continue
label_index = self.classes_dict[label]
bbox = self._get_bbox_from_points(
shape["points"], image_width, image_height)
if bbox is None:
print(f"无效边界框 - 文件: {json_path}, 标签: {label}")
continue
x_center, y_center, width, height = bbox
txt_file.write(
f"{label_index} {x_center:.6f} {y_center:.6f} "
f"{width:.6f} {height:.6f}\n")
shutil.copy2(image_path, output_image_dir)
return True
except Exception as e:
print(f"转换失败 {json_path}: {e}")
return False
def _generate_yolo_config_file(self):
"""生成 data.yaml"""
config_content = f"""# YOLO目标检测数据集配置文件
path: {os.path.abspath(self.output_dir)}
train: images/train
val: images/val
test:
nc: {len(self.classes_dict)}
names:
"""
for class_name, class_id in sorted(
self.classes_dict.items(), key=lambda x: x[1]):
config_content += f" {class_id}: {class_name}\n"
config_file_path = self.output_dir / "data.yaml"
with open(config_file_path, "w", encoding="utf-8") as file:
file.write(config_content)
print(f"已生成YOLO配置文件: {config_file_path}")
return config_file_path
def convert(self):
print("开始LabelMe到YOLO目标检测格式转换...")
self._extract_classes_from_annotations()
if not self.classes_dict:
print("错误: 未找到任何有效的类别")
return False
self._create_output_directories()
json_files = glob.glob(str(self.input_dir / "*.json"))
if not json_files:
print("错误: 未找到JSON文件")
return False
print(f"找到 {len(json_files)} 个JSON文件")
train_files, val_files = self._split_dataset_files(json_files)
print(f"数据集分割 - 训练集: {len(train_files)}, 验证集: {len(val_files)}")
train_success = sum(
self._convert_single_annotation(
f, self.output_dir / "labels" / "train",
self.output_dir / "images" / "train")
for f in tqdm(train_files, desc="转换训练集"))
val_success = sum(
self._convert_single_annotation(
f, self.output_dir / "labels" / "val",
self.output_dir / "images" / "val")
for f in tqdm(val_files, desc="转换验证集"))
self._generate_yolo_config_file()
total = len(json_files)
print("=" * 60)
print(f"训练集: {train_success}/{len(train_files)} "
f"验证集: {val_success}/{len(val_files)} "
f"总体成功率: {(train_success + val_success) / total * 100:.1f}%")
print(f"类别数量: {len(self.classes_dict)} 输出目录: {self.output_dir}")
print("=" * 60)
return True
def main():
CONFIG = {
"input_dir": "./datasets/melabel/seal", # LabelMe数据目录(存放JSON)
"output_dir": "./datasets/yolo_seal_det", # YOLO输出目录
"train_ratio": 0.8,
"random_seed": 42,
}
converter = LabelMeToYOLODetectorConverter(**CONFIG)
if not converter.convert():
print("转换过程中出现错误,请检查输入数据和配置")
return 1
return 0
if __name__ == "__main__":
exit(main())
5.3 分割格式转换脚本 labelme2yolo_seg.py
分割标签与检测的差别只在"每行的坐标部分":检测是 (cx,cy,w,h),分割是归一化多边形顶点序列 x1 y1 x2 y2 ... xn yn。基于上面的类派生一个分割版:
class LabelMeToYOLOSegConverter(LabelMeToYOLODetectorConverter):
"""LabelMe格式转YOLO实例分割格式转换器
输出: [类别ID, x1,y1, x2,y2, ...] 全部归一化
"""
def _convert_single_annotation(self, json_path, output_label_dir, output_image_dir):
try:
with open(json_path, "r", encoding="utf-8") as file:
annotation_data = json.load(file)
image_path = self._find_image_file(json_path)
if not image_path:
print(f"未找到图像文件: {json_path}")
return False
image = self._read_image_with_chinese_path(image_path)
if image is None:
return False
image_height, image_width = image.shape[:2]
txt_path = output_label_dir / (Path(json_path).stem + ".txt")
with open(txt_path, "w", encoding="utf-8") as txt_file:
for shape in annotation_data["shapes"]:
label = shape["label"].strip()
if label not in self.classes_dict:
continue
points = shape["points"]
if len(points) < 3: # 多边形至少3个顶点
continue
# 顶点归一化(逐点除以宽高),并裁剪到边界内
norm_points = []
for x, y in points:
nx = self._clip_coordinate_to_boundary(x, image_width) / image_width
ny = self._clip_coordinate_to_boundary(y, image_height) / image_height
norm_points.extend([f"{nx:.6f}", f"{ny:.6f}"])
txt_file.write(f"{self.classes_dict[label]} " + " ".join(norm_points) + "\n")
shutil.copy2(image_path, output_image_dir)
return True
except Exception as e:
print(f"转换失败 {json_path}: {e}")
return False
在main()中把LabelMeToYOLODetectorConverter换成LabelMeToYOLOSegConverter,输入目录指向文字标注 JSON、输出目录如./datasets/yolo_word_seg即可。
5.4 转换结果可视化自检(强烈建议)
转换后随机挑一张图把标注画回去目视检查——转换对不对看一眼就知道:
from PIL import Image, ImageDraw, ImageFont
import yaml
import random
def get_class_color(class_id, classes_count, seed=42):
"""固定种子生成类别颜色:同一类别每次颜色一致,便于核对"""
random.seed(seed + class_id)
return tuple(random.randint(60, 255) for _ in range(3))
label_path = "datasets/yolo_seal_det/labels/train/0090_0569.txt"
image_path = "datasets/yolo_seal_det/images/train/0090_0569.png"
yaml_path = "datasets/yolo_seal_det/data.yaml"
with open(label_path, "r", encoding="utf-8") as f:
label_lines = [line.strip() for line in f.readlines() if line.strip()]
img = Image.open(image_path).convert("RGB")
img_width, img_height = img.size
with open(yaml_path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
class_names, total_classes = cfg["names"], len(cfg["names"])
font = ImageFont.truetype("chinese_cht.ttf", 30, encoding="utf-8")
draw = ImageDraw.Draw(img)
for line in label_lines:
parts = line.split()
class_id = int(parts[0])
x_center, y_center, width, height = map(float, parts[1:5])
if class_id < 0 or class_id >= total_classes: # 防御:ID 越界跳过而非中断
print(f"警告:无效的类别ID {class_id},跳过该标注")
continue
# 归一化 (cx,cy,w,h) -> 像素坐标(反归一化)
x_min = int((x_center - width / 2) * img_width)
y_min = int((y_center - height / 2) * img_height)
x_max = int((x_center + width / 2) * img_width)
y_max = int((y_center + height / 2) * img_height)
color = get_class_color(class_id, total_classes)
draw.rectangle((x_min, y_min, x_max, y_max), outline=color, width=3)
draw.text((x_min, max(y_min - 30, 0)), class_names[class_id],
fill="white", font=font)
img.show()
六、模型训练
6.1 训练印章检测模型(YOLO11n)
data.yaml 配置(转换脚本已自动生成,核心是 names 顺序必须与转换时的类别 ID 一致):
path: ./datasets/yolo_seal_det
train: images/train
val: images/val
nc: 1
names:
0: seal
两种等价训练方式:
# 方式一:CLI 命令行
yolo detect train data=datasets/yolo_seal_det/data.yaml model=yolo11n.pt epochs=100 imgsz=640
# 方式二:Python API
from ultralytics import YOLO
model = YOLO("yolo11n.pt") # Nano版:体积小、速度快,适合快速迭代部署
model.train(data="./cfg/seal_det.yaml", epochs=100, imgsz=640)
6.2 训练文字分割模型(YOLO11n-seg)
from ultralytics import YOLO
model = YOLO("yolo11n-seg.pt") # -seg 后缀 = 分割权重
model.train(data="./cfg/word_seg.yaml", epochs=100, imgsz=640)
训练监控要点:
- 分割训练要同时看
box_mAP和seg_mAP(分割掩膜精度),后者才是文字着色质量的关键; - 配合
patience=100自动早停、cos_lr=True余弦退火; - batch size 从显存能放下的最大值试起(OOM 就减半);epoch 不预设死,看验证 loss 回升就停。
训练产物在 runs/detect/train/weights/best.pt 和 runs/segment/train/weights/best.pt。
七、ONNX 导出
from ultralytics import YOLO
# 检测模型
YOLO("runs/detect/train/weights/best.pt").export(format="onnx", imgsz=640)
# 分割模型
YOLO("runs/segment/train/weights/best.pt").export(format="onnx", imgsz=640)
分割 ONNX 模型有两个输出:outputs[0]检测头(8400 个候选框 × (4+nc+32)),outputs[1]原型掩膜(32×160×160)。这个结构是理解分割后处理的关键。
八、推理部署:预处理 + 检测/分割封装类
部署环境只需 onnxruntime + opencv + numpy,不需要安装 ultralytics,可放到无 GPU 的工控机上跑。
8.1 输入预处理 image_resize
古籍扫描件分辨率动辄上万像素,先归一到工作分辨率并二值化:
import cv2
import numpy as np
def image_resize(src_img):
"""调整图片大小 + 硬阈值二值化
- 超大图等比缩到 3032x2148 量级(min保证等比,不拉伸变形)
- 二值化消除纸张泛黄、墨色深浅等噪声,最大化印章/文字对比度
"""
h, w = src_img.shape[:2]
if h > 5000:
scale = min(h / 3032, w / 2148)
n_h, n_w = round(h / scale), round(w / scale)
img = cv2.resize(src_img, (n_w, n_h))
elif h < 2500:
n_h, n_w = 3032, 2148
img = cv2.resize(src_img, (n_w, n_h))
else:
img = src_img
img = np.where(img > 127, 255, 0).astype(np.uint8) # >127 置白,否则置黑
return img
8.2 印章检测器 SealDetector(含 NMS 完整后处理)
import onnxruntime as ort
import numpy as np
import cv2
class SealDetector:
"""印章检测器:YOLO11n ONNX 封装。输入 BGR 图,返回 [(box, score, class_id), ...]"""
def __init__(self, onnx_model, confidence_thres=0.5, iou_thres=0.5):
self.onnx_model = onnx_model
self.confidence_thres = confidence_thres
self.iou_thres = iou_thres
self.session = ort.InferenceSession(
self.onnx_model,
providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
self.model_inputs = self.session.get_inputs()
input_shape = self.model_inputs[0].shape # [1, 3, 640, 640]
self.input_height = input_shape[2]
self.input_width = input_shape[3]
def preprocess(self, img):
self.img = img
self.img_height, self.img_width = img.shape[:2]
img = cv2.resize(img, (self.input_width, self.input_height))
image_data = np.array(img) / 255.0
image_data = np.transpose(image_data, (2, 0, 1)) # HWC -> CHW
image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
return image_data
def postprocess(self, output):
outputs = np.transpose(np.squeeze(output[0])) # (8400, 4+nc) -> 按行遍历候选框
rows = outputs.shape[0]
boxes, scores, class_ids = [], [], []
# 直接 resize 预处理(无 letterbox),坐标映射就是简单等比乘法
x_factor = self.img_width / self.input_width
y_factor = self.img_height / self.input_height
for i in range(rows):
classes_scores = outputs[i][4:]
max_score = np.amax(classes_scores)
if max_score >= self.confidence_thres:
class_id = np.argmax(classes_scores)
x, y, w, h = outputs[i][:4]
# 中心点格式 -> 左上角格式,并缩放回原始图像尺寸
left = round((x - w / 2) * x_factor)
top = round((y - h / 2) * y_factor)
boxes.append([left, top, round(w * x_factor), round(h * y_factor)])
scores.append(max_score)
class_ids.append(int(class_id))
# NMS 非极大值抑制
indices = cv2.dnn.NMSBoxes(
boxes, scores, self.confidence_thres, self.iou_thres)
return [(boxes[i], scores[i], class_ids[i]) for i in indices]
def __call__(self, image):
img_data = self.preprocess(image)
outputs = self.session.run(None, {self.model_inputs[0].name: img_data})
return self.postprocess(outputs)
8.3 文字分割器 WordSegModel(原型掩膜机制)
YOLO 分割不直接输出整张掩膜,而是输出 32 个原型掩膜 protos + 每个框 32 个系数,掩膜 = 系数 × 原型的线性组合:
masks_in (N,32) @ protos_flat (32, 160*160) -> (N, 160, 160) -> 裁剪/上采样/合并
class WordSegModel:
"""文字分割器:YOLO11n-seg ONNX 封装。输入 BGR 子图,返回 (文字掩膜, 文字框列表)
mask: (h, w) uint8,255=文字;word_boxes: (N, 4) xyxy
"""
def __init__(self, onnx_model, conf_thres=0.3, iou_thres=0.5):
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.session = ort.InferenceSession(
onnx_model,
providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
self.model_inputs = self.session.get_inputs()
input_shape = self.model_inputs[0].shape
self.input_height, self.input_width = input_shape[2], input_shape[3]
# ---------- 预处理:letterbox(等比缩放 + 灰边,文字对形变敏感,用标准 letterbox) ----------
def preprocess(self, img):
self.img = img
self.img_height, self.img_width = img.shape[:2]
new_shape = (self.input_height, self.input_width)
shape = img.shape[:2]
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
new_unpad = (round(shape[1] * r), round(shape[0] * r))
dw = (new_shape[1] - new_unpad[0]) / 2
dh = (new_shape[0] - new_unpad[1]) / 2
if shape[::-1] != new_unpad:
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = round(dh - 0.1), round(dh + 0.1)
left, right = round(dw - 0.1), round(dw + 0.1)
img = cv2.copyMakeBorder(img, top, bottom, left, right,
cv2.BORDER_CONSTANT, value=(114, 114, 114))
# 记录反 letterbox 参数
self.r, self.dw, self.dh = r, round(dw), round(dh)
image_data = np.array(img) / 255.0
image_data = np.transpose(image_data, (2, 0, 1))
return np.expand_dims(image_data, axis=0).astype(np.float32)
# ---------- zero-mask 裁剪:掩膜限制在各自检测框内,消除框外幽灵掩膜 ----------
@staticmethod
def crop_mask(masks, boxes):
masks = masks.astype(np.float32)
boxes = boxes.astype(np.int32)
n, h, w = masks.shape
for i in range(n):
x1, y1, x2, y2 = boxes[i]
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
masks[i, :y1, :] = 0
masks[i, y2:, :] = 0
masks[i, :, :x1] = 0
masks[i, :, x2:] = 0
return masks
def process_mask(self, protos, masks_in, bboxes, shape, upsample=True):
"""原型掩膜线性组合 -> 裁剪 -> 上采样 -> 二值化"""
c, mh, mw = protos.shape # e.g. (32, 160, 160)
protos_flat = protos.astype(np.float32).reshape(c, -1)
masks = masks_in @ protos_flat # (N, mh*mw)
masks = masks.reshape(-1, mh, mw) # (N, mh, mw)
# 原图尺度 bbox 缩放到 proto 尺度
ratios = np.array([mw / shape[1], mh / shape[0],
mw / shape[1], mh / shape[0]])
masks = self.crop_mask(masks, bboxes * ratios)
if upsample: # 上采样回输入图尺寸
n = masks.shape[0]
masks_up = np.zeros((n, shape[0], shape[1]), dtype=np.float32)
for i in range(n):
masks_up[i] = cv2.resize(masks[i], (shape[1], shape[0]),
interpolation=cv2.INTER_LINEAR)
masks = masks_up
return (masks > 0.0).astype(np.uint8)
def postprocess(self, outputs):
protos = outputs[1][0] # 原型掩膜 (32, 160, 160)
pred = np.transpose(np.squeeze(outputs[0])) # (8400, 4+1+nc+32)
pred = pred[pred[:, 4] > self.conf_thres]
if len(pred) == 0:
return None, None
boxes, scores, class_ids = [], [], []
x_factor = self.input_width / self.img_width
y_factor = self.input_height / self.img_height
for row in pred:
x, y, w, h = row[:4]
boxes.append([round((x - w / 2) * x_factor), round((y - h / 2) * y_factor),
round((x + w / 2) * x_factor), round((y + h / 2) * y_factor)])
scores.append(float(row[4]))
class_ids.append(int(np.argmax(row[5:-32])))
indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_thres, self.iou_thres)
if len(indices) == 0:
return None, None
keep = np.array(indices).flatten()
boxes = np.array(boxes)[keep]
masks_coef = pred[keep][:, -32:] # 32维掩膜系数
masks = self.process_mask(
protos, masks_coef, boxes,
(self.img_height, self.img_width), upsample=True)
# 合并所有文字实例为单张掩膜
mask_combined = (masks.sum(axis=0) > 0).astype(np.uint8) * 255
return mask_combined, boxes
def __call__(self, image):
img_data = self.preprocess(image)
outputs = self.session.run(None, {self.model_inputs[0].name: img_data})
mask, word_boxes = self.postprocess(outputs)
if mask is not None:
mask = mask[self.dh:self.img_height + self.dh or None,
self.dw:self.img_width + self.dh or None] # 粗裁灰边
return mask, word_boxes
形状推导记忆点:masks_in (N,32) @ protos_flat (32, 25600) → (N, 160, 160),其中25600 = 160×160是 proto 空间分辨率(640 输入 4 倍下采样)。
九、上色主流程
9.1 整体流程
PDF页面(300dpi渲染) → 灰度 → image_resize(缩放+二值化)
→ SealDetector 整页检测(4类标签)
→ 切上/下半页(各为一个以页宽为边长的正方形)
→ 逐印章分流:
label 3 版框 → 框内跑 WordSegModel → 掩膜着色(文字变黑)
label 0/1 印章 → 记录(label, x, y, scale) → 贴朱砂章模板
label 2 → 框内 B 通道提亮(→175)
→ 保存上/下半页彩色 PNG
为什么切上下半页:古籍页面长宽比约 3:2,而分割/着色按正方形块处理效果更好。切分后印章归属用中心点纵坐标与 resize_h/2 比较判断。注意坐标系换算:下半页局部坐标 = 原图坐标 − (resize_h − resize_w),漏掉这个偏移贴图会偏出页面。
9.2 完整代码
import io
import fitz # pymupdf
import numpy as np
import cv2
from PIL import Image
# ---------- 配置 ----------
seal_det_model = SealDetector("weights/seal_det_best.onnx", 0.5, 0.5)
word_seg_model = WordSegModel("weights/word_seg_best.onnx", 0.3, 0.5)
SEAL_TEMPLATE = {0: cv2.imread("templates/seal_small.png", -1), # 小印章 296x296
1: cv2.imread("templates/seal_big.png", -1)} # 大印章 740x732
SEAL0_SHAPE, SEAL1_SHAPE = [296, 296], [740, 732]
DIV = 10 # 框外扩像素,防贴边
def paste_seal(img, label, x, y, scale):
"""按记录的 scale 缩放印章模板并 alpha 混合贴图"""
tpl = SEAL_TEMPLATE[label]
base_shape = SEAL0_SHAPE if label == 0 else SEAL1_SHAPE
n_w, n_h = round(base_shape[0] * scale), round(base_shape[1] * scale)
tpl = cv2.resize(tpl, (n_w, n_h))
x, y = int(x), int(y)
h, w = img.shape[:2]
# 边界裁剪
x1, y1 = max(x, 0), max(y, 0)
x2, y2 = min(x + n_w, w), min(y + n_h, h)
if x2 <= x1 or y2 <= y1:
return
roi = tpl[y1 - y:y2 - y, x1 - x:x2 - x]
alpha = roi[:, :, 3:4] / 255.0
img[y1:y2, x1:x2, :3] = (roi[:, :, :3] * alpha +
img[y1:y2, x1:x2, :3] * (1 - alpha)).astype(np.uint8)
def colorize_page(page, pdf_idx, page_number):
# 1. PDF 页 -> 图像(300dpi 印刷级分辨率)
pix = page.get_pixmap(dpi=300)
src_img = np.array(Image.open(io.BytesIO(pix.tobytes("png"))))
if len(src_img.shape) == 3:
src_img = cv2.cvtColor(src_img, cv2.COLOR_RGB2GRAY)
img = image_resize(src_img)
resize_h, resize_w = img.shape[:2]
infer_image = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB).astype(np.uint8)
# 2. 印章检测
boxes = seal_det_model(infer_image)
# 3. 切上/下半页(古籍页 3:2,各取一个以页宽为边长的正方形)
up_img = infer_image[:resize_w, :]
down_img = infer_image[resize_h - resize_w:, :]
half_h = resize_h // 2
up_record = {"box": [], "seal_info": []} # box=版框四角点, seal_info=[label,x,y,scale]
down_record = {"box": [], "seal_info": []}
# 4. 逐印章分流(上半页与下半页,y 坐标系不同)
for box, score, label in boxes:
seal_x0, seal_y0, w, h = box
if label == 2: # 提亮区:B通道 暗像素→175
tgt = up_img if seal_y0 + h / 2 < half_h else down_img
ty = seal_y0 if tgt is up_img else seal_y0 - (resize_h - resize_w)
tgt[:, seal_x0:seal_x0 + w, 2] = np.where(
tgt[:, seal_x0:seal_x0 + w, 2] < 127, 175,
tgt[:, seal_x0:seal_x0 + w, 2])
continue
if label in (0, 1): # 大小印章:中心对齐 + 记录 scale
base = SEAL0_SHAPE if label == 0 else SEAL1_SHAPE
scale = (w / base[0] + h / base[1]) / 2
n_w, n_h = round(base[0] * scale), round(base[1] * scale)
nx = seal_x0 + (w - n_w) // 2
ny = seal_y0 + (h - n_h) // 2
rec = up_record if seal_y0 + h / 2 < half_h else down_record
if rec is down_record:
ny -= resize_h - resize_w # 下半页坐标平移!
rec["seal_info"].append([label, nx, ny, scale])
continue
if label == 3 and w >= 1350 and h >= 950: # 版框(尺寸门槛过滤伪框)
pts = [[seal_x0 - DIV, seal_y0 - DIV],
[seal_x0 + w + DIV, seal_y0 - DIV],
[seal_x0 + w + DIV, seal_y0 + h],
[seal_x0 - DIV, seal_y0 + h]]
if seal_y0 + h / 2 < half_h:
up_record["box"] = pts
else:
pts = [[p[0], p[1] - (resize_h - resize_w)] for p in pts] # y 平移
down_record["box"] = pts
# 5. 版框内文字分割 + 着色
for record, canvas in ((up_record, up_img), (down_record, down_img)):
if not record["box"]:
continue
(x0, y0), _, (x1, y1), _ = record["box"]
x0, y0 = max(x0 - DIV, 0), max(y0 - DIV, 0)
x1, y1 = min(x1 + DIV, canvas.shape[1]), min(y1 + DIV, canvas.shape[0])
mask, word_boxes = word_seg_model(canvas[y0:y1, x0:x1])
if mask is None:
continue
mask = WordSegModel.crop_mask(mask[None], word_boxes)[0] # 保险裁剪
# 文字变黑:掩膜内 R 通道强制 0,掩膜外反相;B 通道统一提亮 175
mask0 = np.where(mask > 127, 0, 255 - canvas[y0:y1, x0:x1, 0])
canvas[y0:y1, x0:x1, 2] = np.where(
mask0 > 127, 175, canvas[y0:y1, x0:x1, 2])
# 6. 贴章
for label, sx, sy, scale in record["seal_info"]:
paste_seal(canvas, label, sx, sy, scale)
# 7. 保存
cv2.imwrite(f"output/{pdf_idx}_{page_number:04d}_up.png",
cv2.cvtColor(up_img, cv2.COLOR_RGB2BGR))
cv2.imwrite(f"output/{pdf_idx}_{page_number:04d}_down.png",
cv2.cvtColor(down_img, cv2.COLOR_RGB2BGR))
def batch_process(pdf_path, pdf_idx):
doc = fitz.open(pdf_path)
for page_number, page in enumerate(doc):
colorize_page(page, pdf_idx, page_number)
if __name__ == "__main__":
batch_process("sikuquanshu.pdf", pdf_idx=0)
9.3 关键设计解读
|
设计 |
原因 |
|---|---|
|
|
过滤小字号、页边距等伪"版框",数字是项目调出来的工程经验 |
|
|
防止印章/文字贴边被裁掉 |
|
分块推理(版框内跑分割) |
版框外全是空白,整页推理浪费且易误检 |
|
文字二值化后再推理 |
消除纸张泛黄/墨色深浅噪声,对比度最大化 |
|
|
掩膜内文字 R 通道强制 0(纯黑),掩膜外反相增强对比 |
|
B 通道统一 175 |
整体色调渲染,与红章合成时视觉更协调 |
十、常见问题与调优
Q1:cv2.imread 读中文路径返回 None?
用 np.fromfile + cv2.imdecode 按字节流读取(本文转换脚本已实现)。
Q2:训练好的模型换机器部署,GPU/CPU 怎么兼容?
available = ort.get_available_providers()
providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in available]
session = ort.InferenceSession(model, providers=providers or available)
Q3:分割掩膜出现在检测框外(幽灵掩膜)?
用 zero-mask 裁剪:掩膜先缩放到 proto 尺度,按 bbox 把框外像素置 0,再上采样。
Q4:如何确定 batch size 和 epoch?
batch 从显存能放下的最大值试起(OOM 减半),大 batch 需同比例调大 lr;epoch 看验证 loss,回升即早停(patience=50~100),fine-tune 通常 100~300 轮起步。
Q5:分割推理坐标对不上?
检查三件事:① 预处理是 resize 还是 letterbox——letterbox 必须做逆变换(减 padding、除 r);② 上/下半页切分后坐标是否平移;③ 掩膜上采样尺寸是否等于当前画布尺寸。
Q6:标注量太少怎么办?
用迭代闭环:初版模型预标注海量图片 → 人工修正 → 回流重训。注意修正质量,错误标签会被模型自我强化。
Q7:要印刷输出,分辨率不够?
推理用工作分辨率提速,但着色结果应映射回原尺寸:保存每个实例的多边形顶点(原图尺度)而非 640 下的位图,最终在原图上色。
浙公网安备 33010602011771号