Method save_visualization

tools/infer_doc_onnx.py:1207–1233

保存可视化结果

(self, result: Dict, output_path: str)

            f.write(f'{merged_text}\n\n')

    def save_visualization(self, result: Dict, output_path: str):
        """保存可视化结果"""
        img_name, img_dir = _get_image_name_and_dir(result, output_path)
        vis_path = os.path.join(img_dir, f'{img_name}_vis.jpg')

        if '_page_image' in result:
            image = result['_page_image'].copy()
        else:
            image = cv2.imread(result['input_path'])

        for box_info in result['layout_results']['boxes']:
            x1, y1, x2, y2 = map(int, box_info['coordinate'])
            label = box_info['label']
            score = box_info['score']

            # 提取基础标签名（去除编号后缀，如 text_01 -> text）
            base_label = label.rsplit('_', 1)[0] if '_' in label and label.rsplit('_', 1)[1].isdigit() else label

            # 获取颜色，如果没有定义则使用默认红色
            color = self.colors.get(base_label, (255, 0, 0))

            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image, f'{label}: {score: .2f}', (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        cv2.imwrite(vis_path, image)
        # logger.info(f"  Saved visualization to {vis_path}")


# ==================== Main Function ====================