基于 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(矩形框)

定位物体大致区域

印章标注,标签名 seal

实例分割标注

Create Polygons(多边形)

精确勾勒目标轮廓

文字标注,标签名 word

印章标注五步走:打开图片 → 选 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)

一行一个目标:检测 [id, cx, cy, w, h];分割 [id, x1,y1, x2,y2, ...]

特点

直观易读,数值随图片尺寸变化,不能直接训练

与图片尺寸无关,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 关键设计解读

设计

原因

label==3 尺寸门槛 1350×950

过滤小字号、页边距等伪"版框",数字是项目调出来的工程经验

DIV=10 框外扩

防止印章/文字贴边被裁掉

分块推理(版框内跑分割)

版框外全是空白,整页推理浪费且易误检

文字二值化后再推理

消除纸张泛黄/墨色深浅噪声,对比度最大化

np.where(mask>127, 0, 255-R)

掩膜内文字 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 下的位图,最终在原图上色。

posted @ 2026-09-13 21:22  鬼门元歌  阅读(12)  评论(0)    收藏  举报