PP-OCRv5_mobile_det转换为onnx预测框差异过大

[dlq986]

请问用自定义数据集微调PP-OCRv5_mobile_det后，使用PaddleOCR的api进行预测，效果很好：

将模型导出为onnx模型之后，自己做后处理后的预测框差别过大，只标注到了中心区域，部署到安卓端后是同样的效果，请问这个是什么原因导致的？如果后处理逻辑有问题，正确的后处理逻辑是什么样？

onnx 的Python 预测脚本如下：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ONNX模型预测脚本 - 用于车牌检测
使用PP-OCRv5_mobile_det ONNX模型进行文本检测
"""

import cv2
import numpy as np
import onnxruntime as ort
import json
from pathlib import Path
from typing import List, Tuple, Dict, Union


class DBPostProcess:
    """DBNet后处理类"""
    
    def __init__(self, thresh=0.3, box_thresh=0.6, max_candidates=1000, unclip_ratio=1.5):
        self.thresh = thresh
        self.box_thresh = box_thresh
        self.max_candidates = max_candidates
        self.unclip_ratio = unclip_ratio
        self.min_size = 3
    
    def __call__(self, pred, shape_list):
        """
        后处理预测结果
        :param pred: 模型输出的预测图
        :param shape_list: 原始图像形状列表
        :return: 检测框和分数
        """
        pred = pred[:, 0, :, :]
        segmentation = pred > self.thresh
        
        boxes_batch = []
        scores_batch = []
        
        for batch_index in range(pred.shape[0]):
            src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
            mask = segmentation[batch_index]
            boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask, src_w, src_h)
            boxes_batch.append(boxes)
            scores_batch.append(scores)
        
        return boxes_batch, scores_batch
    
    def boxes_from_bitmap(self, pred, bitmap, dest_width, dest_height):
        """从bitmap中提取检测框"""
        import cv2
        from shapely.geometry import Polygon
        import pyclipper
        
        height, width = bitmap.shape
        contours, _ = cv2.findContours(
            (bitmap * 255).astype(np.uint8),
            cv2.RETR_LIST,
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        num_contours = min(len(contours), self.max_candidates)
        boxes = []
        scores = []
        
        for index in range(num_contours):
            contour = contours[index]
            points, sside = self.get_mini_boxes(contour)
            if sside < self.min_size:
                continue
            
            points = np.array(points)
            score = self.box_score_fast(pred, points.reshape(-1, 2))
            if score < self.box_thresh:
                continue
            
            # unclip
            box = self.unclip(points, self.unclip_ratio)
            if len(box) == 0 or len(box) > 1:
                continue
            box = box.reshape(-1, 1, 2)
            
            box, sside = self.get_mini_boxes(box)
            if sside < self.min_size + 2:
                continue
            
            box = np.array(box)
            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
            
            boxes.append(box.tolist())
            scores.append(score)
        
        return boxes, scores
    
    def unclip(self, box, unclip_ratio):
        """扩展检测框"""
        import pyclipper  # type: ignore
        from shapely.geometry import Polygon
        
        poly = Polygon(box)
        distance = poly.area * unclip_ratio / poly.length
        offset = pyclipper.PyclipperOffset()  # type: ignore
        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)  # type: ignore
        expanded = np.array(offset.Execute(distance))
        return expanded
    
    def get_mini_boxes(self, contour):
        """获取最小外接矩形"""
        bounding_box = cv2.minAreaRect(contour)
        points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
        
        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
        if points[1][1] > points[0][1]:
            index_1 = 0
            index_4 = 1
        else:
            index_1 = 1
            index_4 = 0
        
        if points[3][1] > points[2][1]:
            index_2 = 2
            index_3 = 3
        else:
            index_2 = 3
            index_3 = 2
        
        box = [points[index_1], points[index_2], points[index_3], points[index_4]]
        return box, min(bounding_box[1])
    
    def box_score_fast(self, bitmap, box):
        """计算框的置信度分数"""
        h, w = bitmap.shape[:2]
        box = box.copy()
        xmin = np.clip(np.floor(box[:, 0].min()).astype(int), 0, w - 1)
        xmax = np.clip(np.ceil(box[:, 0].max()).astype(int), 0, w - 1)
        ymin = np.clip(np.floor(box[:, 1].min()).astype(int), 0, h - 1)
        ymax = np.clip(np.ceil(box[:, 1].max()).astype(int), 0, h - 1)
        
        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        box[:, 0] = box[:, 0] - xmin
        box[:, 1] = box[:, 1] - ymin
        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)  # type: ignore
        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]


class ONNXDetector:
    """ONNX检测模型预测器"""
    
    def __init__(self, model_path: str):
        """
        初始化ONNX检测器
        :param model_path: ONNX模型文件路径
        """
        self.model_path = model_path
        
        # 创建ONNX Runtime session
        self.session = ort.InferenceSession(
            model_path,
            providers=['CPUExecutionProvider']  # 可以改为 'CUDAExecutionProvider' 使用GPU
        )
        
        # 获取输入输出信息
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
        # 图像预处理参数（根据inference.yml配置）
        self.mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3)).astype(np.float32)
        self.std = np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3)).astype(np.float32)
        
        # 后处理器
        self.postprocessor = DBPostProcess(
            thresh=0.3,
            box_thresh=0.6,
            max_candidates=1000,
            unclip_ratio=1.5
        )
        
        print(f"模型加载成功: {model_path}")
        print(f"输入名称: {self.input_name}")
        print(f"输出名称: {self.output_name}")
    
    def resize_image(self, img: np.ndarray, max_side_len=736) -> Tuple[np.ndarray, Tuple]:
        """
        调整图像大小，保持宽高比
        :param img: 输入图像
        :param max_side_len: 最大边长
        :return: 调整后的图像和形状信息
        """
        h, w = img.shape[:2]
        
        # 计算缩放比例
        ratio = 1.0
        if max(h, w) > max_side_len:
            if h > w:
                ratio = float(max_side_len) / h
            else:
                ratio = float(max_side_len) / w
        
        resize_h = int(h * ratio)
        resize_w = int(w * ratio)
        
        # 调整为32的倍数
        resize_h = max(int(round(resize_h / 32) * 32), 32)
        resize_w = max(int(round(resize_w / 32) * 32), 32)
        
        # 调整图像大小
        if resize_h != h or resize_w != w:
            img = cv2.resize(img, (resize_w, resize_h))
        
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)
        
        return img, (h, w, ratio_h, ratio_w)
    
    def preprocess(self, img: np.ndarray) -> Tuple[np.ndarray, List]:
        """
        图像预处理
        :param img: 输入图像 (BGR格式)
        :return: 预处理后的图像和形状信息
        """
        # 调整图像大小
        img, shape_info = self.resize_image(img)
        
        # 归一化
        img = img.astype(np.float32) / 255.0
        img = (img - self.mean) / self.std
        
        # HWC -> CHW
        img = img.transpose(2, 0, 1)
        
        # 添加batch维度
        img = np.expand_dims(img, axis=0)
        
        return img, [shape_info]
    
    def predict(self, image_path: str) -> Dict:
        """
        对单张图像进行预测
        :param image_path: 图像文件路径
        :return: 预测结果字典
        """
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # 预处理
        input_data, shape_list = self.preprocess(img)
        
        # 推理
        outputs = self.session.run([self.output_name], {self.input_name: input_data})
        pred = outputs[0]
        
        # 后处理
        boxes_batch, scores_batch = self.postprocessor(pred, shape_list)
        
        # 整理结果
        result = {
            "input_path": image_path,
            "dt_polys": boxes_batch[0] if boxes_batch else [],
            "dt_scores": scores_batch[0] if scores_batch else []
        }
        
        return result
    
    def visualize(self, image_path: str, result: Dict, output_path: Union[str, None] = None):
        """
        可视化检测结果
        :param image_path: 原始图像路径
        :param result: 预测结果
        :param output_path: 输出图像路径（如果为None，则自动生成）
        """
        # 读取原始图像
        img = cv2.imread(image_path)
        
        # 绘制检测框
        for box, score in zip(result['dt_polys'], result['dt_scores']):
            box = np.array(box).astype(np.int32).reshape(-1, 2)
            cv2.polylines(img, [box], True, (0, 255, 0), 2)
            
            # 添加分数标签
            cv2.putText(
                img,
                f"{score:.2f}",
                tuple(box[0]),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 0, 255),
                1
            )
        
        # 保存结果
        if output_path is None:
            output_path = str(Path(image_path).stem) + "_result.jpg"
        
        cv2.imwrite(output_path, img)
        print(f"可视化结果已保存到: {output_path}")
        
        return output_path


def main():
    """主函数 - 示例用法"""
    
    # 模型路径
    model_path = "./onnx_output/label-original/inference.onnx"
    
    # 测试图像路径
    image_path = "./test.jpg"
    
    # 输出目录
    output_dir = "./output/onnx_test"
    Path(output_dir).mkdir(exist_ok=True)
    
    # 创建检测器
    detector = ONNXDetector(model_path)
    
    # 预测
    print(f"\n正在处理图像: {image_path}")
    result = detector.predict(image_path)
    
    # 打印结果
    print("\n检测结果:")
    print(f"检测到 {len(result['dt_polys'])} 个文本区域")
    for i, (box, score) in enumerate(zip(result['dt_polys'], result['dt_scores'])):
        print(f"  区域 {i+1}: 置信度 {score:.4f}, 坐标 {box}")
    
    # 保存JSON结果
    json_path = f"{output_dir}/result.json"
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=4)
    print(f"\nJSON结果已保存到: {json_path}")
    
    # 可视化
    vis_path = f"{output_dir}/result_vis.jpg"
    detector.visualize(image_path, result, vis_path)
    
    print("\n预测完成！")


if __name__ == "__main__":
    main()

[dlq986]

安卓端部署的时候做了二值化处理，处理后的图像也是只有中心区域。

[dlq986]

调整了unclip的逻辑，效果差异没那么大了。