Spaces:
Running
Running
| import gradio as gr | |
| import tempfile | |
| from pathlib import Path | |
| import json | |
| import torch | |
| import time | |
| import re | |
| import threading | |
| from queue import Queue | |
| import zipfile | |
| from datetime import datetime | |
| from detection_ocr import DetectionOCR | |
| from repair_head import RepairHead | |
| from semantic_evaluation import SemanticEvaluationPipeline | |
| # ============================================================ | |
| # Device Configuration | |
| # ============================================================ | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"[System] Using device: {DEVICE}") | |
| # ============================================================ | |
| # Lazy Singletons | |
| # ============================================================ | |
| _ocr = None | |
| _repair = None | |
| _semantic = None | |
| def get_ocr(): | |
| global _ocr | |
| if _ocr is None: | |
| print("[INIT] Loading DetectionOCR") | |
| _ocr = DetectionOCR(device=DEVICE) | |
| return _ocr | |
| def get_repair(): | |
| global _repair | |
| if _repair is None: | |
| print("[INIT] Loading RepairHead (CPU)") | |
| _repair = RepairHead(device="cpu") | |
| return _repair | |
| def get_semantic(): | |
| global _semantic | |
| if _semantic is None: | |
| print("[INIT] Loading Semantic Evaluation Pipeline") | |
| _semantic = SemanticEvaluationPipeline( | |
| base_model="Qwen/Qwen2.5-Math-1.5B-Instruct", | |
| lora_dir=None, | |
| device="cpu", | |
| ) | |
| return _semantic | |
| # ============================================================ | |
| # Helpers | |
| # ============================================================ | |
| KNOWN_ENV = ["equation", "align", "gather", "multline", "array", "cases"] | |
| def wrap_equation(expr: str) -> str: | |
| expr = expr.strip() | |
| if any(expr.startswith(f"\\begin{{{e}}}") for e in KNOWN_ENV): | |
| return expr | |
| if expr.startswith("$$") and expr.endswith("$$"): | |
| return expr | |
| if expr.startswith("\\[") and expr.endswith("\\]"): | |
| return expr | |
| return f"\\begin{{equation}}\n{expr}\n\\end{{equation}}" | |
| def is_reasonable(expr: str) -> bool: | |
| if not expr or not isinstance(expr, str): | |
| return False | |
| if expr.lower() in ["noise", "error"]: | |
| return False | |
| if len(expr) < 3 or len(expr) > 2000: | |
| return False | |
| return True | |
| def save_progress_log(outdir: Path, stage: str, content: str): | |
| """Save progress logs to files""" | |
| log_file = outdir / f"{stage}_log.txt" | |
| with open(log_file, "a", encoding="utf-8") as f: | |
| timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| f.write(f"[{timestamp}] {content}\n") | |
| def create_visualization_pdf(outdir: Path, repaired_data: dict, semantic_results: list): | |
| """创建可视化PDF，按页面顺序显示公式""" | |
| try: | |
| from fpdf import FPDF | |
| pdf = FPDF() | |
| pdf.set_auto_page_break(auto=True, margin=15) | |
| # 添加标题页 | |
| pdf.add_page() | |
| pdf.set_font('Arial', 'B', 16) | |
| pdf.cell(0, 10, 'PDF Formula Processing Results', 0, 1, 'C') | |
| pdf.ln(10) | |
| pdf.set_font('Arial', '', 12) | |
| pdf.cell(0, 10, f'Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}', 0, 1, 'C') | |
| pdf.ln(20) | |
| # 统计信息 | |
| total_pages = len(repaired_data) | |
| total_formulas = sum(len(page_data) for page_data in repaired_data.values()) | |
| valid_formulas = sum(len([f for f in page_data if f['repaired'] not in ['noise', 'error']]) | |
| for page_data in repaired_data.values()) | |
| pdf.set_font('Arial', 'B', 14) | |
| pdf.cell(0, 10, 'Processing Statistics', 0, 1, 'L') | |
| pdf.ln(5) | |
| pdf.set_font('Arial', '', 12) | |
| pdf.cell(0, 8, f'• Total Pages: {total_pages}', 0, 1, 'L') | |
| pdf.cell(0, 8, f'• Total Formulas Detected: {total_formulas}', 0, 1, 'L') | |
| pdf.cell(0, 8, f'• Valid Formulas Repaired: {valid_formulas}', 0, 1, 'L') | |
| pdf.cell(0, 8, f'• Semantic Evaluations: {len(semantic_results)}', 0, 1, 'L') | |
| pdf.ln(10) | |
| # 按页面处理 | |
| page_keys = sorted(repaired_data.keys(), key=lambda x: int(re.search(r'page_(\d+)', x).group(1)) if re.search(r'page_(\d+)', x) else 0) | |
| for page_key in page_keys: | |
| formulas = repaired_data[page_key] | |
| valid_formulas = [f for f in formulas if f['repaired'] not in ['noise', 'error']] | |
| if not valid_formulas: | |
| continue | |
| pdf.add_page() | |
| pdf.set_font('Arial', 'B', 14) | |
| pdf.cell(0, 10, f'{page_key.upper()}', 0, 1, 'L') | |
| pdf.set_font('Arial', '', 10) | |
| pdf.cell(0, 6, f'Valid Formulas: {len(valid_formulas)} / {len(formulas)}', 0, 1, 'L') | |
| pdf.ln(5) | |
| for idx, formula_data in enumerate(valid_formulas): | |
| formula = formula_data['repaired'] | |
| original = formula_data['raw_latex'][:150] + "..." if len(formula_data['raw_latex']) > 150 else formula_data['raw_latex'] | |
| # 公式标题 | |
| pdf.set_font('Arial', 'B', 10) | |
| pdf.cell(0, 6, f'Formula {idx+1}:', 0, 1, 'L') | |
| # 原始公式 | |
| pdf.set_font('Arial', 'I', 8) | |
| pdf.cell(20, 4, 'Original:', 0, 0, 'L') | |
| pdf.set_font('Courier', '', 8) | |
| pdf.multi_cell(0, 4, f' {original}', 0, 'L') | |
| # 修复后公式 | |
| pdf.set_font('Arial', 'I', 8) | |
| pdf.cell(20, 4, 'Repaired:', 0, 0, 'L') | |
| pdf.set_font('Courier', '', 8) | |
| formula_display = formula[:150] + "..." if len(formula) > 150 else formula | |
| pdf.multi_cell(0, 4, f' {formula_display}', 0, 'L') | |
| # 查找对应的语义结果 | |
| sem_result = None | |
| for sem in semantic_results: | |
| if sem.get('original_formula') == formula_data['raw_latex'] or sem.get('fixed_formula') == formula: | |
| sem_result = sem | |
| break | |
| if sem_result: | |
| verdict = sem_result.get('semantic_verdict', 'unknown') | |
| confidence = sem_result.get('confidence', 0) | |
| pdf.set_font('Arial', '', 9) | |
| verdict_color = { | |
| 'equivalent': (0, 100, 0), # 绿色 | |
| 'contradiction': (220, 0, 0), # 红色 | |
| 'non_verifiable': (255, 140, 0), # 橙色 | |
| 'garbled': (128, 128, 128), # 灰色 | |
| }.get(verdict, (0, 0, 0)) # 黑色默认 | |
| pdf.set_text_color(*verdict_color) | |
| pdf.cell(0, 5, f'Verdict: {verdict} (Confidence: {confidence:.2f})', 0, 1, 'L') | |
| pdf.set_text_color(0, 0, 0) # 重置为黑色 | |
| if sem_result.get('explanation'): | |
| pdf.set_font('Arial', 'I', 8) | |
| explanation = sem_result['explanation'][:200] + "..." if len(sem_result['explanation']) > 200 else sem_result['explanation'] | |
| pdf.multi_cell(0, 4, f'Explanation: {explanation}', 0, 'L') | |
| pdf.ln(3) | |
| # 每5个公式后添加分隔线 | |
| if (idx + 1) % 5 == 0 and idx + 1 < len(valid_formulas): | |
| pdf.line(10, pdf.get_y(), 200, pdf.get_y()) | |
| pdf.ln(5) | |
| # 保存PDF | |
| pdf_path = outdir / "visualization_report.pdf" | |
| pdf.output(str(pdf_path)) | |
| return pdf_path | |
| except Exception as e: | |
| print(f"Error creating visualization PDF: {e}") | |
| # 如果fpdf不可用，创建简单的文本报告 | |
| try: | |
| report_path = outdir / "visualization_report.txt" | |
| with open(report_path, 'w', encoding='utf-8') as f: | |
| f.write("=" * 60 + "\n") | |
| f.write("PDF FORMULA PROCESSING RESULTS\n") | |
| f.write("=" * 60 + "\n") | |
| f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n") | |
| # 统计信息 | |
| total_pages = len(repaired_data) | |
| total_formulas = sum(len(page_data) for page_data in repaired_data.values()) | |
| valid_formulas = sum(len([f for f in page_data if f['repaired'] not in ['noise', 'error']]) | |
| for page_data in repaired_data.values()) | |
| f.write("PROCESSING STATISTICS:\n") | |
| f.write(f"• Total Pages: {total_pages}\n") | |
| f.write(f"• Total Formulas Detected: {total_formulas}\n") | |
| f.write(f"• Valid Formulas Repaired: {valid_formulas}\n") | |
| f.write(f"• Semantic Evaluations: {len(semantic_results)}\n\n") | |
| f.write("=" * 60 + "\n\n") | |
| # 按页面显示公式 | |
| page_keys = sorted(repaired_data.keys(), key=lambda x: int(re.search(r'page_(\d+)', x).group(1)) if re.search(r'page_(\d+)', x) else 0) | |
| for page_key in page_keys: | |
| formulas = repaired_data[page_key] | |
| valid_formulas_list = [f for f in formulas if f['repaired'] not in ['noise', 'error']] | |
| if not valid_formulas_list: | |
| continue | |
| f.write(f"\n{page_key.upper()}\n") | |
| f.write("-" * 40 + "\n") | |
| f.write(f"Valid Formulas: {len(valid_formulas_list)} / {len(formulas)}\n\n") | |
| for idx, formula_data in enumerate(valid_formulas_list): | |
| formula = formula_data['repaired'] | |
| original = formula_data['raw_latex'] | |
| f.write(f"Formula {idx+1}:\n") | |
| f.write(f" Original: {original}\n") | |
| f.write(f" Repaired: {formula}\n") | |
| # 查找语义结果 | |
| sem_result = None | |
| for sem in semantic_results: | |
| if sem.get('original_formula') == original or sem.get('fixed_formula') == formula: | |
| sem_result = sem | |
| break | |
| if sem_result: | |
| verdict = sem_result.get('semantic_verdict', 'unknown') | |
| confidence = sem_result.get('confidence', 0) | |
| f.write(f" Verdict: {verdict} (Confidence: {confidence:.2f})\n") | |
| if sem_result.get('explanation'): | |
| explanation = sem_result['explanation'] | |
| f.write(f" Explanation: {explanation}\n") | |
| f.write("\n") | |
| return report_path | |
| except Exception as e2: | |
| print(f"Fallback text report also failed: {e2}") | |
| return None | |
| # ============================================================ | |
| # Stream Manager for Real-time Updates | |
| # ============================================================ | |
| class StreamManager: | |
| """Manage real-time streaming to multiple outputs""" | |
| def __init__(self): | |
| self.ocr_queue = Queue() | |
| self.repair_queue = Queue() | |
| self.semantic_queue = Queue() | |
| self.summary_queue = Queue() | |
| def ocr_stream(self, message: str): | |
| self.ocr_queue.put(message) | |
| if hasattr(self, 'current_outdir'): | |
| save_progress_log(self.current_outdir, "ocr", message) | |
| def repair_stream(self, message: str): | |
| self.repair_queue.put(message) | |
| if hasattr(self, 'current_outdir'): | |
| save_progress_log(self.current_outdir, "repair", message) | |
| def semantic_stream(self, message: str): | |
| self.semantic_queue.put(message) | |
| if hasattr(self, 'current_outdir'): | |
| save_progress_log(self.current_outdir, "semantic", message) | |
| def summary_stream(self, message: str): | |
| self.summary_queue.put(message) | |
| if hasattr(self, 'current_outdir'): | |
| save_progress_log(self.current_outdir, "summary", message) | |
| def set_output_dir(self, outdir: Path): | |
| self.current_outdir = outdir | |
| # ============================================================ | |
| # Main Processing Pipeline | |
| # ============================================================ | |
| def process_pipeline(pdf_file, stream_mgr: StreamManager): | |
| """Main processing pipeline with streaming to different outputs""" | |
| if pdf_file is None: | |
| stream_mgr.summary_stream("❌ No PDF uploaded.") | |
| return None, None, None, None, None, [] | |
| try: | |
| # Create output directory | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| tmpdir = Path(tempfile.mkdtemp(prefix=f"pdf_pipeline_{timestamp}_")) | |
| outdir = tmpdir / "output" | |
| outdir.mkdir(exist_ok=True) | |
| stream_mgr.set_output_dir(outdir) | |
| stream_mgr.summary_stream(f"📁 Output directory: {outdir}") | |
| stream_mgr.summary_stream("=" * 60) | |
| # ============ STAGE 1: OCR ============ | |
| stream_mgr.ocr_stream("🚀 STAGE 1: OCR Detection Started") | |
| stream_mgr.ocr_stream("Loading OCR model...") | |
| ocr = get_ocr() | |
| stream_mgr.ocr_stream("Processing PDF pages...") | |
| ocr.process_pdf(pdf_file, output_dir=outdir) | |
| stream_mgr.ocr_stream("✅ OCR Processing Completed") | |
| stream_mgr.summary_stream("✓ Stage 1: OCR completed") | |
| # ============ STAGE 2: Merge Results ============ | |
| stream_mgr.ocr_stream("\n🔄 Merging page results...") | |
| merged = {} | |
| page_files = sorted(outdir.glob("page_*_results.json")) | |
| if not page_files: | |
| stream_mgr.ocr_stream("❌ No page results found!") | |
| stream_mgr.summary_stream("❌ Error: No page results found") | |
| return None, None, None, None, None, [] | |
| for p in page_files: | |
| m = re.search(r"(page_\d+)_results\.json", p.name) | |
| if not m: | |
| continue | |
| page = m.group(1) | |
| try: | |
| data = json.loads(p.read_text("utf-8")) | |
| formulas = data.get("formulas", []) | |
| merged[page] = { | |
| "formulas": formulas, | |
| "total": len(formulas) | |
| } | |
| stream_mgr.ocr_stream(f" {page}: {len(formulas)} formulas") | |
| except Exception as e: | |
| stream_mgr.ocr_stream(f" ⚠️ Error loading {page}: {str(e)}") | |
| merged[page] = {"formulas": [], "total": 0} | |
| # Save overall OCR results | |
| overall_json = outdir / "overall_ocr_results.json" | |
| overall_json.write_text(json.dumps(merged, indent=2), "utf-8") | |
| stream_mgr.ocr_stream(f"✅ Saved: {overall_json.name}") | |
| stream_mgr.summary_stream(f"✓ Stage 2: Merged {len(merged)} pages") | |
| # ============ STAGE 3: Repair Formulas ============ | |
| stream_mgr.repair_stream("🔧 STAGE 3: LaTeX Repair Started") | |
| stream_mgr.repair_stream("Loading Repair model...") | |
| repair = get_repair() | |
| repaired_all = {} | |
| repaired_eqs = [] | |
| total_formulas = sum(v["total"] for v in merged.values()) | |
| stream_mgr.repair_stream(f"Total formulas to repair: {total_formulas}") | |
| stream_mgr.repair_stream("=" * 40) | |
| processed = 0 | |
| for page, content in merged.items(): | |
| stream_mgr.repair_stream(f"\n📄 Processing {page}:") | |
| repaired_all[page] = [] | |
| for i, fm in enumerate(content["formulas"]): | |
| processed += 1 | |
| raw = fm.get("raw_latex", "") | |
| # Show progress every 5 formulas or first/last | |
| if processed <= 5 or processed >= total_formulas - 5 or processed % 10 == 0: | |
| raw_preview = raw[:100] + "..." if len(raw) > 100 else raw | |
| stream_mgr.repair_stream(f" [{processed}/{total_formulas}] {raw_preview}") | |
| try: | |
| repaired = repair.repair(raw) | |
| if is_reasonable(repaired): | |
| repaired_all[page].append({ | |
| "raw_latex": raw, | |
| "repaired": repaired | |
| }) | |
| repaired_eqs.append(repaired) | |
| if processed <= 3: # Show first few repairs | |
| stream_mgr.repair_stream(f" → {repaired[:80]}...") | |
| else: | |
| repaired_all[page].append({ | |
| "raw_latex": raw, | |
| "repaired": "noise" | |
| }) | |
| except Exception as e: | |
| repaired_all[page].append({ | |
| "raw_latex": raw, | |
| "repaired": f"error: {str(e)}" | |
| }) | |
| stream_mgr.repair_stream(f" ⚠️ Error: {str(e)}") | |
| # Save repaired results | |
| repaired_json = outdir / "overall_ocr_repaired.json" | |
| repaired_json.write_text(json.dumps(repaired_all, indent=2), "utf-8") | |
| stream_mgr.repair_stream(f"\n✅ Repair completed: {len(repaired_eqs)} valid formulas") | |
| stream_mgr.summary_stream(f"✓ Stage 3: Repaired {len(repaired_eqs)}/{total_formulas} formulas") | |
| # ============ STAGE 4: Semantic Evaluation ============ | |
| stream_mgr.semantic_stream("🧠 STAGE 4: Semantic Evaluation Started") | |
| stream_mgr.semantic_stream("Loading Semantic model...") | |
| semantic = get_semantic() | |
| semantic_results = [] | |
| stream_mgr.semantic_stream(f"Evaluating {len(repaired_eqs)} formulas...") | |
| stream_mgr.semantic_stream("=" * 40) | |
| for i, eq in enumerate(repaired_eqs): | |
| try: | |
| # Show progress | |
| if i < 3 or i >= len(repaired_eqs) - 3 or i % 5 == 0: | |
| eq_preview = eq[:80] + "..." if len(eq) > 80 else eq | |
| stream_mgr.semantic_stream(f"[{i+1}/{len(repaired_eqs)}] {eq_preview}") | |
| res = semantic.evaluate_formula(eq) | |
| res['original_formula'] = eq # 添加原始公式 | |
| semantic_results.append(res) | |
| # Show verdict for first few | |
| if i < 5: | |
| stream_mgr.semantic_stream(f" Verdict: {res.get('semantic_verdict', 'unknown')}") | |
| stream_mgr.semantic_stream(f" Confidence: {res.get('confidence', 0):.2f}") | |
| except Exception as e: | |
| stream_mgr.semantic_stream(f"⚠️ Error evaluating formula {i+1}: {str(e)}") | |
| semantic_results.append({ | |
| 'original_formula': eq, | |
| 'error': str(e), | |
| 'semantic_verdict': 'evaluation_error', | |
| 'confidence': 0.0 | |
| }) | |
| # Save semantic results | |
| semantic_json = outdir / "semantic_results.json" | |
| semantic_json.write_text(json.dumps(semantic_results, indent=2), "utf-8") | |
| stream_mgr.semantic_stream(f"\n✅ Semantic evaluation completed: {len(semantic_results)} formulas") | |
| stream_mgr.summary_stream(f"✓ Stage 4: Semantically evaluated {len(semantic_results)} formulas") | |
| # ============ STAGE 5: Generate Visualization PDF ============ | |
| stream_mgr.summary_stream("\n📊 Generating visualization PDF...") | |
| # Generate LaTeX file | |
| merge_tex = outdir / "merge.tex" | |
| with merge_tex.open("w", encoding="utf-8") as f: | |
| f.write("\\documentclass{article}\n") | |
| f.write("\\usepackage{amsmath, amssymb, bm}\n") | |
| f.write("\\begin{document}\n\n") | |
| f.write("% Auto-generated from PDF formula extraction pipeline\n") | |
| f.write("% Generated on: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n\n") | |
| for eq in repaired_eqs: | |
| f.write(wrap_equation(eq) + "\n\n") | |
| f.write("\\end{document}\n") | |
| stream_mgr.summary_stream("✓ Generated LaTeX file") | |
| # Create visualization PDF | |
| visualization_pdf = create_visualization_pdf(outdir, repaired_all, semantic_results) | |
| if visualization_pdf and visualization_pdf.exists(): | |
| stream_mgr.summary_stream(f"✓ Created visualization report: {visualization_pdf.name}") | |
| else: | |
| stream_mgr.summary_stream("⚠️ Could not create visualization PDF") | |
| visualization_pdf = None | |
| # ============ STAGE 6: Generate Statistics ============ | |
| stats = { | |
| "total_pages": len(merged), | |
| "total_formulas": total_formulas, | |
| "repaired_formulas": len(repaired_eqs), | |
| "semantic_evaluated": len(semantic_results), | |
| "output_directory": str(outdir), | |
| "processing_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), | |
| } | |
| stats_json = outdir / "processing_stats.json" | |
| stats_json.write_text(json.dumps(stats, indent=2), "utf-8") | |
| # Final summary | |
| stream_mgr.summary_stream("=" * 60) | |
| stream_mgr.summary_stream("🎉 PROCESSING COMPLETED SUCCESSFULLY!") | |
| stream_mgr.summary_stream("=" * 60) | |
| stream_mgr.summary_stream(f"📊 SUMMARY STATISTICS:") | |
| stream_mgr.summary_stream(f" • Pages processed: {len(merged)}") | |
| stream_mgr.summary_stream(f" • Formulas detected: {total_formulas}") | |
| stream_mgr.summary_stream(f" • Formulas repaired: {len(repaired_eqs)}") | |
| stream_mgr.summary_stream(f" • Formulas evaluated: {len(semantic_results)}") | |
| stream_mgr.summary_stream(f" • Output directory: {outdir}") | |
| stream_mgr.summary_stream("=" * 60) | |
| stream_mgr.summary_stream("📁 Generated files:") | |
| stream_mgr.summary_stream(f" 1. {overall_json.name} - OCR results") | |
| stream_mgr.summary_stream(f" 2. {repaired_json.name} - Repaired formulas") | |
| stream_mgr.summary_stream(f" 3. {semantic_json.name} - Semantic evaluation") | |
| stream_mgr.summary_stream(f" 4. {merge_tex.name} - Combined LaTeX") | |
| if visualization_pdf: | |
| stream_mgr.summary_stream(f" 5. {visualization_pdf.name} - Visualization report") | |
| stream_mgr.summary_stream(f" 6. {stats_json.name} - Processing statistics") | |
| stream_mgr.summary_stream("=" * 60) | |
| return overall_json, repaired_json, semantic_json, merge_tex, visualization_pdf, semantic_results | |
| except Exception as e: | |
| import traceback | |
| error_msg = f"❌ CRITICAL ERROR:\n{str(e)}\n\nTraceback:\n{traceback.format_exc()}" | |
| stream_mgr.summary_stream(error_msg) | |
| return None, None, None, None, None, [] | |
| # ============================================================ | |
| # Simplified Streaming Handler | |
| # ============================================================ | |
| def process_with_streaming(pdf_file): | |
| """Simplified streaming handler without threading issues""" | |
| if not pdf_file: | |
| yield "", "", "", "❌ No PDF uploaded", "Error", None, None, None, None, None, [] | |
| return | |
| # Initialize stream manager | |
| stream_mgr = StreamManager() | |
| # Initialize result variables | |
| result_container = { | |
| 'overall_json': None, | |
| 'repaired_json': None, | |
| 'semantic_json': None, | |
| 'merge_tex': None, | |
| 'visualization_pdf': None, | |
| 'sem_results': [] | |
| } | |
| # Process function | |
| def run_processing(): | |
| """Run the pipeline and store results""" | |
| results = process_pipeline(pdf_file, stream_mgr) | |
| if results: | |
| result_container['overall_json'] = results[0] | |
| result_container['repaired_json'] = results[1] | |
| result_container['semantic_json'] = results[2] | |
| result_container['merge_tex'] = results[3] | |
| result_container['visualization_pdf'] = results[4] | |
| result_container['sem_results'] = results[5] | |
| # Start processing in a thread | |
| process_thread = threading.Thread(target=run_processing) | |
| process_thread.start() | |
| # Initialize log contents | |
| ocr_log_content = "" | |
| repair_log_content = "" | |
| semantic_log_content = "" | |
| summary_log_content = "" | |
| # Stream updates while processing | |
| last_update = time.time() | |
| while process_thread.is_alive(): | |
| # Collect new messages | |
| while not stream_mgr.ocr_queue.empty(): | |
| ocr_log_content += stream_mgr.ocr_queue.get() + "\n" | |
| while not stream_mgr.repair_queue.empty(): | |
| repair_log_content += stream_mgr.repair_queue.get() + "\n" | |
| while not stream_mgr.semantic_queue.empty(): | |
| semantic_log_content += stream_mgr.semantic_queue.get() + "\n" | |
| while not stream_mgr.summary_queue.empty(): | |
| summary_log_content += stream_mgr.summary_queue.get() + "\n" | |
| # Update UI periodically | |
| current_time = time.time() | |
| if current_time - last_update > 0.1: | |
| yield ( | |
| ocr_log_content[-5000:], # Limit size | |
| repair_log_content[-5000:], | |
| semantic_log_content[-5000:], | |
| summary_log_content[-5000:], | |
| "Processing...", | |
| None, None, None, None, None, [] | |
| ) | |
| last_update = current_time | |
| time.sleep(0.05) | |
| # Wait for thread to complete | |
| process_thread.join() | |
| # Collect any remaining messages | |
| while not stream_mgr.ocr_queue.empty(): | |
| ocr_log_content += stream_mgr.ocr_queue.get() + "\n" | |
| while not stream_mgr.repair_queue.empty(): | |
| repair_log_content += stream_mgr.repair_queue.get() + "\n" | |
| while not stream_mgr.semantic_queue.empty(): | |
| semantic_log_content += stream_mgr.semantic_queue.get() + "\n" | |
| while not stream_mgr.summary_queue.empty(): | |
| summary_log_content += stream_mgr.summary_queue.get() + "\n" | |
| # Prepare download files | |
| ocr_file = str(result_container['overall_json']) if result_container['overall_json'] else None | |
| repair_file = str(result_container['repaired_json']) if result_container['repaired_json'] else None | |
| semantic_file = str(result_container['semantic_json']) if result_container['semantic_json'] else None | |
| tex_file = str(result_container['merge_tex']) if result_container['merge_tex'] else None | |
| pdf_file_dl = str(result_container['visualization_pdf']) if result_container['visualization_pdf'] else None | |
| # Final update | |
| yield ( | |
| ocr_log_content[-10000:], | |
| repair_log_content[-10000:], | |
| semantic_log_content[-10000:], | |
| summary_log_content[-10000:], | |
| "Completed", | |
| ocr_file, | |
| repair_file, | |
| semantic_file, | |
| tex_file, | |
| pdf_file_dl, | |
| result_container['sem_results'] | |
| ) | |
| # ============================================================ | |
| # Gradio UI with Multi-Stream Display | |
| # ============================================================ | |
| with gr.Blocks(title="PDF Formula Pipeline", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 📘 PDF Formula Processing Pipeline | |
| **Real-time streaming display with visualization report** | |
| """) | |
| # Input Section | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Upload PDF") | |
| pdf_input = gr.File( | |
| label="Upload PDF File", | |
| file_types=[".pdf"], | |
| height=100 | |
| ) | |
| with gr.Row(): | |
| run_btn = gr.Button("Start Processing", variant="primary", size="lg") | |
| clear_btn = gr.Button("Clear All", variant="secondary") | |
| gr.Markdown("### Processing Status") | |
| status_text = gr.Textbox( | |
| label="Current Status", | |
| value="Ready", | |
| interactive=False | |
| ) | |
| # Multi-Stream Display Section | |
| gr.Markdown("### 🎯 Real-time Processing Logs") | |
| with gr.Tabs(): | |
| with gr.TabItem("📄 OCR Process"): | |
| ocr_log = gr.Textbox( | |
| label="OCR Stage Log", | |
| lines=15, | |
| interactive=False, | |
| autoscroll=True | |
| ) | |
| with gr.TabItem("🔧 Repair Process"): | |
| repair_log = gr.Textbox( | |
| label="Repair Stage Log", | |
| lines=15, | |
| interactive=False, | |
| autoscroll=True | |
| ) | |
| with gr.TabItem("🧠 Semantic Process"): | |
| semantic_log = gr.Textbox( | |
| label="Semantic Stage Log", | |
| lines=15, | |
| interactive=False, | |
| autoscroll=True | |
| ) | |
| with gr.TabItem("📊 Summary"): | |
| summary_log = gr.Textbox( | |
| label="Processing Summary", | |
| lines=20, | |
| interactive=False, | |
| autoscroll=True | |
| ) | |
| # Results Download Section | |
| gr.Markdown("### 📥 Download Results") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| ocr_json_dl = gr.File(label="OCR Results (JSON)", visible=True) | |
| repaired_json_dl = gr.File(label="Repaired Formulas (JSON)", visible=True) | |
| with gr.Column(scale=1): | |
| semantic_json_dl = gr.File(label="Semantic Results (JSON)", visible=True) | |
| merge_tex_dl = gr.File(label="Combined LaTeX (.tex)", visible=True) | |
| with gr.Row(): | |
| visualization_pdf_dl = gr.File(label="Visualization Report", visible=True) | |
| # Formula Details Section | |
| gr.Markdown("### 🔍 Formula Details") | |
| with gr.Row(): | |
| formula_idx = gr.Number( | |
| label="Formula Index (0 to preview)", | |
| minimum=0, | |
| maximum=1000, | |
| step=1, | |
| precision=0, | |
| value=0 | |
| ) | |
| preview_btn = gr.Button("Show Details", variant="secondary") | |
| formula_details = gr.JSON( | |
| label="Formula Details", | |
| visible=True | |
| ) | |
| # Hidden state for semantic results | |
| semantic_state = gr.State([]) | |
| # Preview function | |
| def show_formula_details(idx, sem_results): | |
| if not sem_results or idx < 0 or idx >= len(sem_results): | |
| return {} | |
| result = sem_results[int(idx)] | |
| # 格式化输出 | |
| formatted_result = { | |
| "Formula": result.get('original_formula', 'N/A'), | |
| "Semantic Verdict": result.get('semantic_verdict', 'unknown'), | |
| "Confidence": f"{result.get('confidence', 0):.2f}" if isinstance(result.get('confidence'), (int, float)) else "N/A", | |
| "Source": result.get('source', 'N/A'), | |
| "Syntactically Correct": result.get('syntactically_correct', 'N/A'), | |
| "Mathematically Correct": result.get('mathematically_correct', 'N/A'), | |
| } | |
| if result.get('explanation'): | |
| formatted_result["Explanation"] = result['explanation'] | |
| if result.get('suggested_fix'): | |
| formatted_result["Suggested Fix"] = result['suggested_fix'] | |
| if result.get('error'): | |
| formatted_result["Error"] = result['error'] | |
| return formatted_result | |
| # Clear function | |
| def clear_all(): | |
| return [ | |
| None, # pdf_input | |
| "", # ocr_log | |
| "", # repair_log | |
| "", # semantic_log | |
| "", # summary_log | |
| "Ready", # status_text | |
| None, # ocr_json_dl | |
| None, # repaired_json_dl | |
| None, # semantic_json_dl | |
| None, # merge_tex_dl | |
| None, # visualization_pdf_dl | |
| [], # semantic_state | |
| 0, # formula_idx | |
| {} # formula_details | |
| ] | |
| # Connect events | |
| outputs = [ | |
| ocr_log, | |
| repair_log, | |
| semantic_log, | |
| summary_log, | |
| status_text, | |
| ocr_json_dl, | |
| repaired_json_dl, | |
| semantic_json_dl, | |
| merge_tex_dl, | |
| visualization_pdf_dl, | |
| semantic_state | |
| ] | |
| run_btn.click( | |
| fn=process_with_streaming, | |
| inputs=[pdf_input], | |
| outputs=outputs | |
| ) | |
| preview_btn.click( | |
| fn=show_formula_details, | |
| inputs=[formula_idx, semantic_state], | |
| outputs=[formula_details] | |
| ) | |
| formula_idx.change( | |
| fn=show_formula_details, | |
| inputs=[formula_idx, semantic_state], | |
| outputs=[formula_details] | |
| ) | |
| clear_btn.click( | |
| fn=clear_all, | |
| outputs=outputs + [formula_idx, formula_details] | |
| ) | |
| # ============================================================ | |
| # Main Execution | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| print("=" * 60) | |
| print("PDF Formula Processing Pipeline - Visualization Version") | |
| print(f"Device: {DEVICE}") | |
| print("=" * 60) | |
| # 尝试导入fpdf，如果失败则使用文本报告 | |
| try: | |
| from fpdf import FPDF | |
| print("✓ fpdf is available for PDF generation") | |
| except ImportError: | |
| print("⚠️ fpdf not available, will use text report instead") | |
| print(" To install: pip install fpdf") | |
| demo.queue().launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| debug=False, | |
| show_error=True | |
| ) |