Refactor figure extraction to be async

app.py

@@ -244,7 +244,7 @@ async def process_pdf(
         progress(0.3, desc="Step 2/5: Extracting figures from PDF...")
         extraction_work_dir = work_dir / "figure_extraction"
         extraction_work_dir.mkdir()
-        paired_dir = run_figure_extraction(str(pdf_path), str(extraction_work_dir), progress=progress)
+        paired_dir = await run_figure_extraction(str(pdf_path), str(extraction_work_dir), progress=progress)
         if not paired_dir or not any(Path(paired_dir).iterdir()):
             raise gr.Error("Failed to extract any figures from the PDF.")
 

pragent/backend/figure_table_pipeline.py

@@ -8,9 +8,10 @@ from pragent.backend.loader import ImagePDFLoader
 from pragent.backend.yolo import extract_and_save_layout_components
 from tqdm.asyncio import tqdm
 
+import asyncio
 from typing import Any
 
-def run_figure_extraction(pdf_path: str, base_work_dir: str, progress: Any | None = None) -> str:
+async def run_figure_extraction(pdf_path: str, base_work_dir: str, progress: Any | None = None) -> str:
     """
     一个完整的、从PDF提取并配对图表的流程。
     这是被 app.py 调用的主函数。
@@ -52,9 +53,18 @@ def run_figure_extraction(pdf_path: str, base_work_dir: str, progress: Any | Non
     for i, path in enumerate(page_image_paths):
         if progress:
             progress(0.3 + (i / num_pages) * 0.2, desc=f"Analyzing page {i+1}/{num_pages}")
+        
         page_num_str = Path(path).stem
         page_crop_dir = os.path.join(cropped_results_dir, page_num_str)
-        extract_and_save_layout_components(image_path=path, model_path=model_path, save_base_dir=page_crop_dir, imgsz=640)
+        
+        await asyncio.to_thread(
+            extract_and_save_layout_components,
+            image_path=path,
+            model_path=model_path,
+            save_base_dir=page_crop_dir,
+            imgsz=640
+        )
+
     tqdm.write(f"[*] 所有裁剪结果已保存至: {cropped_results_dir}")
 
     tqdm.write(f"\n--- 步骤 3/3: 对裁剪出的组件进行配对 ---")