Эх сурвалжийг харах

feat: 重构API处理逻辑,优化结果保存和Markdown内容处理

zhch158_admin 2 сар өмнө
parent
commit
6b7c1ec47e

+ 67 - 411
zhch/ppstructurev3_single_client.py

@@ -16,377 +16,15 @@ from dotenv import load_dotenv
 load_dotenv(override=True)
 
 from utils import (
-    get_image_files_from_dir,
-    get_image_files_from_list,
-    get_image_files_from_csv,
     collect_pid_files,
-    load_images_from_pdf,
-    normalize_financial_numbers,
-    normalize_markdown_table
 )
 
-def convert_pdf_to_images(pdf_file: str, output_dir: str | None = None, dpi: int = 200) -> List[str]:
-    """
-    将PDF转换为图像文件
-    
-    Args:
-        pdf_file: PDF文件路径
-        output_dir: 输出目录
-        dpi: 图像分辨率
-        
-    Returns:
-        生成的图像文件路径列表
-    """
-    pdf_path = Path(pdf_file)
-    if not pdf_path.exists() or pdf_path.suffix.lower() != '.pdf':
-        print(f"❌ Invalid PDF file: {pdf_path}")
-        return []
-
-    # 如果没有指定输出目录,使用PDF同名目录
-    if output_dir is None:
-        output_path = pdf_path.parent / f"{pdf_path.stem}"
-    else:
-        output_path = Path(output_dir) / f"{pdf_path.stem}"
-    output_path = output_path.resolve()
-    output_path.mkdir(parents=True, exist_ok=True)
-
-    try:
-        # 使用doc_utils中的函数加载PDF图像
-        images = load_images_from_pdf(str(pdf_path), dpi=dpi)
-        
-        image_paths = []
-        for i, image in enumerate(images):
-            # 生成图像文件名
-            image_filename = f"{pdf_path.stem}_page_{i+1:03d}.png"
-            image_path = output_path / image_filename
-
-            # 保存图像
-            image.save(str(image_path))
-            image_paths.append(str(image_path))
-            
-        print(f"✅ Converted {len(images)} pages from {pdf_path.name} to images")
-        return image_paths
-        
-    except Exception as e:
-        print(f"❌ Error converting PDF {pdf_path}: {e}")
-        traceback.print_exc()
-        return []
-
-def get_input_files(args) -> List[str]:
-    """
-    获取输入文件列表,统一处理PDF和图像文件
-    
-    Args:
-        args: 命令行参数
-        
-    Returns:
-        处理后的图像文件路径列表
-    """
-    input_files = []
-    
-    # 获取原始输入文件
-    if args.input_csv:
-        raw_files = get_image_files_from_csv(args.input_csv, "fail")
-    elif args.input_file_list:
-        raw_files = get_image_files_from_list(args.input_file_list)
-    elif args.input_file:
-        raw_files = [Path(args.input_file).resolve()]
-    else:
-        input_dir = Path(args.input_dir).resolve()
-        if not input_dir.exists():
-            print(f"❌ Input directory does not exist: {input_dir}")
-            return []
-        
-        # 获取所有支持的文件(图像和PDF)
-        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
-        pdf_extensions = ['.pdf']
-        
-        raw_files = []
-        for ext in image_extensions + pdf_extensions:
-            raw_files.extend(list(input_dir.glob(f"*{ext}")))
-            raw_files.extend(list(input_dir.glob(f"*{ext.upper()}")))
-        
-        raw_files = [str(f) for f in raw_files]
-    
-    # 分别处理PDF和图像文件
-    pdf_count = 0
-    image_count = 0
-    
-    for file_path in raw_files:
-        file_path = Path(file_path)
-        
-        if file_path.suffix.lower() == '.pdf':
-            # 转换PDF为图像
-            print(f"📄 Processing PDF: {file_path.name}")
-            pdf_images = convert_pdf_to_images(
-                str(file_path), 
-                args.output_dir,
-                dpi=args.pdf_dpi
-            )
-            input_files.extend(pdf_images)
-            pdf_count += 1
-        else:
-            # 直接添加图像文件
-            if file_path.exists():
-                input_files.append(str(file_path))
-                image_count += 1
-    
-    print(f"📊 Input summary:")
-    print(f"  PDF files processed: {pdf_count}")
-    print(f"  Image files found: {image_count}")
-    print(f"  Total image files to process: {len(input_files)}")
-    
-    return input_files
-
-def convert_api_result_to_json(api_result: Dict[str, Any], 
-                              input_image_path: str, 
-                              output_dir: str, 
-                              filename: str,
-                              normalize_numbers: bool = True) -> tuple[str, Dict[str, Any]]:
-    """
-    将API返回结果转换为标准JSON格式,并支持数字标准化
-    """
-    # 获取主要数据
-    layout_parsing_results = api_result.get('layoutParsingResults', [])
-    
-    if not layout_parsing_results:
-        print("⚠️ Warning: No layoutParsingResults found in API response")
-        return {}
-    
-    # 取第一个结果(通常只有一个)
-    main_result = layout_parsing_results[0]
-    pruned_result = main_result.get('prunedResult', {})
-    
-    # 构造标准格式的JSON
-    converted_json = {
-        "input_path": input_image_path,
-        "page_index": None,
-        "model_settings": pruned_result.get('model_settings', {}),
-        "parsing_res_list": pruned_result.get('parsing_res_list', []),
-        "doc_preprocessor_res": {
-            "input_path": None,
-            "page_index": None,
-            "model_settings": pruned_result.get('doc_preprocessor_res', {}).get('model_settings', {}),
-            "angle": pruned_result.get('doc_preprocessor_res', {}).get('angle', 0)
-        },
-        "layout_det_res": {
-            "input_path": None,
-            "page_index": None,
-            "boxes": pruned_result.get('layout_det_res', {}).get('boxes', [])
-        },
-        "overall_ocr_res": {
-            "input_path": None,
-            "page_index": None,
-            "model_settings": pruned_result.get('overall_ocr_res', {}).get('model_settings', {}),
-            "dt_polys": pruned_result.get('overall_ocr_res', {}).get('dt_polys', []),
-            "text_det_params": pruned_result.get('overall_ocr_res', {}).get('text_det_params', {}),
-            "text_type": pruned_result.get('overall_ocr_res', {}).get('text_type', 'general'),
-            "textline_orientation_angles": pruned_result.get('overall_ocr_res', {}).get('textline_orientation_angles', []),
-            "text_rec_score_thresh": pruned_result.get('overall_ocr_res', {}).get('text_rec_score_thresh', 0.0),
-            "return_word_box": pruned_result.get('overall_ocr_res', {}).get('return_word_box', False),
-            "rec_texts": pruned_result.get('overall_ocr_res', {}).get('rec_texts', []),
-            "rec_scores": pruned_result.get('overall_ocr_res', {}).get('rec_scores', []),
-            "rec_polys": pruned_result.get('overall_ocr_res', {}).get('rec_polys', []),
-            "rec_boxes": pruned_result.get('overall_ocr_res', {}).get('rec_boxes', [])
-        },
-        "table_res_list": pruned_result.get('table_res_list', [])
-    }
-    
-    # 数字标准化处理
-    original_json = converted_json.copy()
-    changes_count = 0
-    
-    if normalize_numbers:
-        # 1. 标准化 parsing_res_list 中的文本内容
-        for item in converted_json.get('parsing_res_list', []):
-            if 'block_content' in item:
-                original_content = item['block_content']
-                normalized_content = original_content
-                # 根据block_label类型选择标准化方法
-                if item.get('block_label') == 'table':
-                    normalized_content = normalize_markdown_table(original_content)
-                # else:
-                #     normalized_content = normalize_financial_numbers(original_content)
-                
-                if original_content != normalized_content:
-                    item['block_content'] = normalized_content
-                    changes_count += len([1 for o, n in zip(original_content, normalized_content) if o != n])
-        
-        # 2. 标准化 table_res_list 中的HTML表格
-        for table_item in converted_json.get('table_res_list', []):
-            if 'pred_html' in table_item:
-                original_html = table_item['pred_html']
-                normalized_html = normalize_markdown_table(original_html)
-                
-                if original_html != normalized_html:
-                    table_item['pred_html'] = normalized_html
-                    changes_count += len([1 for o, n in zip(original_html, normalized_html) if o != n])
-
-        # 检查是否需要修复表格一致性(这里只做统计,实际修复可能需要更复杂的逻辑)
-               # 统计表格数量
-        parsing_res_tables_count = 0
-        table_res_list_count = 0
-        if 'parsing_res_list' in converted_json:
-            parsing_res_tables_count = len([item for item in converted_json['parsing_res_list'] 
-                                          if 'block_label' in item and item['block_label'] == 'table'])
-        if 'table_res_list' in converted_json:
-            table_res_list_count = len(converted_json["table_res_list"])
-        table_consistency_fixed = False
-        if parsing_res_tables_count != table_res_list_count:
-            warnings.warn(f"⚠️ Warning: {filename} Table count mismatch - parsing_res_list has {parsing_res_tables_count} tables, "
-                          f"but table_res_list has {table_res_list_count} tables.")
-            table_consistency_fixed = True
-            # 这里可以添加实际的修复逻辑,例如根据需要添加或删除表格项
-            # 但由于缺乏具体规则,暂时只做统计和警告
-
-        # 3. 标准化 overall_ocr_res 中的识别文本
-        # ocr_res = converted_json.get('overall_ocr_res', {})
-        # if 'rec_texts' in ocr_res:
-        #     original_texts = ocr_res['rec_texts'][:]
-        #     normalized_texts = []
-            
-        #     for text in original_texts:
-        #         normalized_text = normalize_financial_numbers(text)
-        #         normalized_texts.append(normalized_text)
-        #         if text != normalized_text:
-        #             changes_count += len([1 for o, n in zip(text, normalized_text) if o != n])
-            
-        #     ocr_res['rec_texts'] = normalized_texts
-        
-        # 添加标准化处理信息
-        converted_json['processing_info'] = {
-            "normalize_numbers": normalize_numbers,
-            "changes_applied": changes_count > 0,
-            "character_changes_count": changes_count,
-            "parsing_res_tables_count": parsing_res_tables_count,
-            "table_res_list_count": table_res_list_count,
-            "table_consistency_fixed": table_consistency_fixed
-        }
-        
-        # if changes_count > 0:
-        #     print(f"🔧 已标准化 {changes_count} 个字符(全角→半角)")
-    else:
-        converted_json['processing_info'] = {
-            "normalize_numbers": False,
-            "changes_applied": False,
-            "character_changes_count": 0
-        }
-    
-    # 保存JSON文件
-    output_path = Path(output_dir).resolve()
-    output_path.mkdir(parents=True, exist_ok=True)
-    
-    json_file_path = output_path / f"{filename}.json"
-    with open(json_file_path, 'w', encoding='utf-8') as f:
-        json.dump(converted_json, f, ensure_ascii=False, indent=2)
-    
-    # 如果启用了标准化且有变化,保存原始版本用于对比
-    if normalize_numbers and changes_count > 0:
-        original_output_path = output_path / f"{filename}_original.json"
-        with open(original_output_path, 'w', encoding='utf-8') as f:
-            json.dump(original_json, f, ensure_ascii=False, indent=2)
-    
-    return str(output_path), converted_json
-
-def save_output_images(api_result: Dict[str, Any], output_dir: str, output_filename: str) -> Dict[str, str]:
-    """
-    保存API返回的输出图像
-    
-    Args:
-        api_result: API返回的结果
-        output_dir: 输出目录
-        
-    Returns:
-        保存的图像文件路径字典
-    """
-    layout_parsing_results = api_result.get('layoutParsingResults', [])
-    if not layout_parsing_results:
-        return {}
-    
-    main_result = layout_parsing_results[0]
-    output_images = main_result.get('outputImages', {})
-    
-    output_path = Path(output_dir).resolve()
-    output_path.mkdir(parents=True, exist_ok=True)
-    
-    saved_images = {}
-    
-    for img_name, img_base64 in output_images.items():
-        try:
-            # 解码base64图像
-            img_data = base64.b64decode(img_base64)
-            
-            # 生成文件名
-            img_filename = f"{output_filename}_{img_name}.jpg"
-            img_path = output_path / img_filename
-            
-            # 保存图像
-            with open(img_path, 'wb') as f:
-                f.write(img_data)
-            
-            saved_images[img_name] = str(img_path)
-            # print(f"📷 Saved image: {img_path}")
-            
-        except Exception as e:
-            print(f"❌ Error saving image {img_name}: {e}")
-    
-    return saved_images
-
-def save_markdown_content(api_result: Dict[str, Any], output_dir: str, 
-                         filename: str, normalize_numbers: bool = True) -> str:
-    """
-    保存Markdown内容,支持数字标准化
-    """
-    layout_parsing_results = api_result.get('layoutParsingResults', [])
-    if not layout_parsing_results:
-        return ""
-    
-    main_result = layout_parsing_results[0]
-    markdown_data = main_result.get('markdown', {})
-    
-    output_path = Path(output_dir).resolve()
-    output_path.mkdir(parents=True, exist_ok=True)
-    
-    # 保存Markdown文本
-    markdown_text = markdown_data.get('text', '')
-    
-    # 数字标准化处理
-    changes_count = 0
-    if normalize_numbers and markdown_text:
-        original_markdown_text = markdown_text
-        markdown_text = normalize_markdown_table(markdown_text)
-        
-        changes_count = len([1 for o, n in zip(original_markdown_text, markdown_text) if o != n])
-        # if changes_count > 0:
-        #     print(f"🔧 Markdown中已标准化 {changes_count} 个字符(全角→半角)")
-    
-    md_file_path = output_path / f"{filename}.md"
-    with open(md_file_path, 'w', encoding='utf-8') as f:
-        f.write(markdown_text)
-    
-    # 如果启用了标准化且有变化,保存原始版本用于对比
-    if normalize_numbers and changes_count > 0:
-        original_output_path = output_path / f"{filename}_original.md"
-        with open(original_output_path, 'w', encoding='utf-8') as f:
-            f.write(original_markdown_text)
-
-    # 保存Markdown中的图像
-    markdown_images = markdown_data.get('images', {})
-    for img_path, img_base64 in markdown_images.items():
-        try:
-            img_data = base64.b64decode(img_base64)
-            full_img_path = output_path / img_path
-            full_img_path.parent.mkdir(parents=True, exist_ok=True)
-            
-            with open(full_img_path, 'wb') as f:
-                f.write(img_data)
-            
-            # print(f"🖼️ Saved Markdown image: {full_img_path}")
-            
-        except Exception as e:
-            print(f"❌ Error saving Markdown image {img_path}: {e}")
-
-    return str(md_file_path)
+from ppstructurev3_utils import (
+   get_input_files,
+   convert_pruned_result_to_json,
+   save_output_images,
+   save_markdown_content
+)
 
 def call_api_for_image(image_path: str, api_url: str, timeout: int = 300) -> Dict[str, Any]:
     """
@@ -474,53 +112,67 @@ def process_images_via_api(image_paths: List[str],
                 api_result = call_api_for_image(img_path, api_url, timeout)
                 processing_time = time.time() - start_time
                 
+                # 获取主要数据
+                layout_parsing_results = api_result.get('layoutParsingResults', [])
+    
+                if not layout_parsing_results:
+                    print("⚠️ Warning: No layoutParsingResults found in API response")
+                    return []
+
                 # 处理API返回结果
                 input_path = Path(img_path)
                 
                 # 生成输出文件名
                 output_filename = input_path.stem
-                
-                # 转换并保存标准JSON格式
-                json_output_path, converted_json = convert_api_result_to_json(
-                    api_result, 
-                    str(input_path), 
-                    output_dir,
-                    output_filename,
-                    normalize_numbers=normalize_numbers
-                )
 
-                # 保存输出图像
-                saved_images = save_output_images(api_result, str(output_dir), output_filename) 
+                # 处理结果
+                for idx, result in enumerate(layout_parsing_results):
+                    if idx > 0:
+                        raise ValueError("Multiple results found for a single image")
+                    
+                    json_content = result.get('prunedResult', {})
+                    json_output_path, converted_json = convert_pruned_result_to_json(
+                        json_content, 
+                        str(input_path), 
+                        output_dir,
+                        output_filename,
+                        normalize_numbers=normalize_numbers
+                    )
 
-                # 保存Markdown内容
-                md_output_path = save_markdown_content(
-                    api_result, 
-                    output_dir, 
-                    output_filename,
-                    normalize_numbers=normalize_numbers
-                )
-                
-                # 记录处理结果
-                all_results.append({
-                    "image_path": str(input_path),
-                    "processing_time": processing_time,
-                    "success": True,
-                    "api_url": api_url,
-                    "output_json": json_output_path,
-                    "output_md": md_output_path,
-                    "is_pdf_page": "_page_" in input_path.name,  # 标记是否为PDF页面
-                    "processing_info": converted_json.get('processing_info', {})
-                })
-                
-                # 更新进度条
-                success_count = sum(1 for r in all_results if r.get('success', False))
-                
-                pbar.update(1)
-                pbar.set_postfix({
-                    'time': f"{processing_time:.2f}s",
-                    'success': f"{success_count}/{len(all_results)}",
-                    'rate': f"{success_count/len(all_results)*100:.1f}%"
-                })
+                    # 保存输出图像
+                    img_content = result.get('outputImages', {})
+                    saved_images = save_output_images(img_content, str(output_dir), output_filename) 
+
+                    # 保存Markdown内容
+                    markdown_content = result.get('markdown', {})
+                    md_output_path = save_markdown_content(
+                        markdown_content, 
+                        output_dir, 
+                        output_filename,
+                        normalize_numbers=normalize_numbers
+                    )
+                    
+                    # 记录处理结果
+                    all_results.append({
+                        "image_path": str(input_path),
+                        "processing_time": processing_time,
+                        "success": True,
+                        "api_url": api_url,
+                        "output_json": json_output_path,
+                        "output_md": md_output_path,
+                        "is_pdf_page": "_page_" in input_path.name,  # 标记是否为PDF页面
+                        "processing_info": converted_json.get('processing_info', {})
+                    })
+                    
+                    # 更新进度条
+                    success_count = sum(1 for r in all_results if r.get('success', False))
+                    
+                    pbar.update(1)
+                    pbar.set_postfix({
+                        'time': f"{processing_time:.2f}s",
+                        'success': f"{success_count}/{len(all_results)}",
+                        'rate': f"{success_count/len(all_results)*100:.1f}%"
+                    })
                 
             except Exception as e:
                 print(f"Error processing {Path(img_path).name}: {e}", file=sys.stderr)
@@ -676,12 +328,16 @@ if __name__ == "__main__":
         
         # 默认配置
         default_config = {
-            "input_file": "/Users/zhch158/workspace/data/至远彩色印刷工业有限公司/data_PPStructureV3_Results/2023年度报告母公司/2023年度报告母公司_page_027.png",
+            # "input_file": "/Users/zhch158/workspace/data/至远彩色印刷工业有限公司/data_PPStructureV3_Results/2023年度报告母公司/2023年度报告母公司_page_027.png",
+            # "input_file": "/home/ubuntu/zhch/data/至远彩色印刷工业有限公司/PPStructureV3_Results/2023年度报告母公司/2023年度报告母公司_page_027.png",
+            "input_file": "/home/ubuntu/zhch/data/至远彩色印刷工业有限公司/2023年度报告母公司.pdf",
+            "output_dir": "/home/ubuntu/zhch/data/至远彩色印刷工业有限公司/PPStructureV3_Results",
+            "collect_results": f"/home/ubuntu/zhch/data/至远彩色印刷工业有限公司/PPStructureV3_Results/processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv",
             # "input_dir": "../../OmniDocBench/OpenDataLab___OmniDocBench/images",
-            "output_dir": "./OmniDocBench_API_Results",
+            # "output_dir": "./OmniDocBench_API_Results",
+            # "collect_results": f"./OmniDocBench_API_Results/processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv",
             "api_url": "http://10.192.72.11:8111/layout-parsing",
             "timeout": "300",
-            "collect_results": f"./OmniDocBench_API_Results/processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv",
         }
         
         # 构造参数