# import gradio as gr # import json # import os # import tempfile # import img2pdf # import glob # import shutil # from img2pdf import Rotation # from pathlib import Path # print("--- DEBUG: Current Working Directory ---") # print(os.getcwd()) # print("--- DEBUG: Files in Root ---") # print(os.listdir('.')) # # ============================== # # PIPELINE IMPORT # # ============================== # # try: # # from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH # # except ImportError: # # print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.") # try: # from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH # except Exception as e: # Catch ALL exceptions # print(f"Warning: Failed to import pipeline: {e}") # import traceback # traceback.print_exc() # Show the actual error # def run_document_pipeline(*args): # return {"error": "Placeholder pipeline function called."} # DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model" # WEIGHTS_PATH = "./weights/yolo_weights.pt" # def process_file(uploaded_files, layoutlmv3_model_path=None): # """ # Robust handler for multiple or single file uploads. # Returns the final JSON and a LIST of all intermediate JSON files (OCR, Predictions, BIO). # """ # if uploaded_files is None: # return "❌ Error: No files uploaded.", None # if not isinstance(uploaded_files, list): # file_list = [uploaded_files] # else: # file_list = uploaded_files # if len(file_list) == 0: # return "❌ Error: Empty file list.", None # # 1. Resolve all file paths safely # resolved_paths = [] # for f in file_list: # try: # if isinstance(f, dict) and "path" in f: # resolved_paths.append(f["path"]) # elif hasattr(f, 'path'): # resolved_paths.append(f.path) # else: # resolved_paths.append(str(f)) # except Exception as e: # print(f"Error resolving path for {f}: {e}") # if not resolved_paths: # return "❌ Error: Could not resolve file paths.", None # # 2. Determine if we should merge into a single PDF # first_file = Path(resolved_paths[0]) # is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff'] # try: # if len(resolved_paths) > 1 or is_image: # print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF...") # temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") # with open(temp_pdf.name, "wb") as f_out: # f_out.write(img2pdf.convert(resolved_paths, rotation=Rotation.ifvalid)) # processing_path = temp_pdf.name # else: # processing_path = resolved_paths[0] # # 3. Standard Pipeline Checks # final_model_path = layoutlmv3_model_path or DEFAULT_LAYOUTLMV3_MODEL_PATH # if not os.path.exists(final_model_path): # return f"❌ Error: Model not found at {final_model_path}", None # # 4. Call the pipeline # print(f"🚀 Starting pipeline for: {processing_path}") # result = run_document_pipeline(processing_path, final_model_path) # # 5. SCRAPE FOR INTERMEDIATE FILES # # We look for all .json files in /tmp/ created during this run # base_name = Path(processing_path).stem # # This matches common patterns like /tmp/pipeline_run_... or filenames in /tmp/ # search_patterns = [ # f"/tmp/pipeline_run_{base_name}*/*.json", # f"/tmp/*{base_name}*.json" # ] # all_intermediate_jsons = [] # for pattern in search_patterns: # all_intermediate_jsons.extend(glob.glob(pattern)) # # Remove duplicates while preserving order # all_intermediate_jsons = list(dict.fromkeys(all_intermediate_jsons)) # # 6. Prepare Final Output for Display # if result is None or (isinstance(result, list) and len(result) == 0): # display_text = "⚠️ Pipeline failed at Step 3 (BIO Decoding).\nDownload the intermediate JSONs below to inspect OCR and Model Predictions." # else: # display_text = json.dumps(result, indent=2, ensure_ascii=False) # # If the final result succeeded, save it to a temp file so it can be downloaded too # temp_final = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='final_result_') # json.dump(result, temp_final, indent=2, ensure_ascii=False) # temp_final.close() # all_intermediate_jsons.append(temp_final.name) # return display_text, all_intermediate_jsons # except Exception as e: # import traceback # traceback.print_exc() # return f"❌ Error: {str(e)}", None # # def visualize_detections(uploaded_files): # # """Shows the first uploaded image with YOLO bounding boxes""" # # if not uploaded_files: # # return None # # # Get first file path # # file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files # # if isinstance(file_path, dict): # # file_path = file_path["path"] # # import cv2 # # from ultralytics import YOLO # # # Load image # # img = cv2.imread(str(file_path)) # # if img is None: # # return None # # # Run YOLO # # model = YOLO(WEIGHTS_PATH) # # results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False) # # # Draw boxes # # for box in results[0].boxes: # # class_id = int(box.cls[0]) # # class_name = model.names[class_id] # # if class_name in ['figure', 'equation']: # # x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy()) # # color = (0, 255, 0) if class_name == 'figure' else (255, 0, 0) # # cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) # # cv2.putText(img, f"{class_name} {box.conf[0]:.2f}", # # (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) # # # Save and return # # temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name # # cv2.imwrite(temp_path, img) # # return temp_path # # # ============================== # # # GRADIO INTERFACE # # # ============================== # # with gr.Blocks(title="Document Analysis Pipeline") as demo: # # gr.Markdown("# 📄 Full Pipeline Analysis") # # gr.Markdown("### 🔍 Intermediate File Recovery Active") # # gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)") # # with gr.Row(): # # with gr.Column(scale=1): # # file_input = gr.File( # # label="Upload PDFs or Images", # # file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"], # # file_count="multiple", # # type="filepath" # # ) # # model_path_input = gr.Textbox( # # label="Model Path", # # value=DEFAULT_LAYOUTLMV3_MODEL_PATH # # ) # # process_btn = gr.Button("🚀 Run Pipeline", variant="primary") # # with gr.Column(scale=2): # # json_output = gr.Code(label="Final Structured Output", language="json", lines=20) # # # IMPORTANT: file_count="multiple" allows returning the list of all stage files # # download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple") # # process_btn.click( # # fn=process_file, # # inputs=[file_input, model_path_input], # # outputs=[json_output, download_output] # # ) # # if __name__ == "__main__": # # demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) # # # ============================== # # # VISUAL DEBUG FUNCTION # # # ============================== # # def visualize_detections(uploaded_files): # # """Shows the first uploaded image with YOLO bounding boxes""" # # if not uploaded_files: # # return None # # try: # # # Get first file path # # file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files # # if isinstance(file_path, dict): # # file_path = file_path["path"] # # elif hasattr(file_path, 'path'): # # file_path = file_path.path # # import cv2 # # import numpy as np # # from ultralytics import YOLO # # import fitz # # # Handle PDF conversion to image # # if str(file_path).lower().endswith('.pdf'): # # doc = fitz.open(file_path) # # page_idx = int(page_num) - 1 # # page = doc.load_page(page_idx) # # pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) # # img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) # # if pix.n == 3: # # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # # elif pix.n == 4: # # img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) # # doc.close() # # else: # # img = cv2.imread(str(file_path)) # # if img is None: # # return None # # # Run YOLO detection # # model = YOLO(WEIGHTS_PATH) # # results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False) # # # Draw bounding boxes # # detection_count = {'figure': 0, 'equation': 0} # # for box in results[0].boxes: # # class_id = int(box.cls[0]) # # class_name = model.names[class_id] # # if class_name in ['figure', 'equation']: # # x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy()) # # conf = float(box.conf[0]) # # # Green for figures, Red for equations # # color = (0, 255, 0) if class_name == 'figure' else (0, 0, 255) # # cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) # # # Add label with confidence # # label = f"{class_name.upper()} {conf:.2f}" # # (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) # # cv2.rectangle(img, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1) # # cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) # # detection_count[class_name] += 1 # # # Add summary text at top # # summary = f"Detected: {detection_count['figure']} Figures (GREEN), {detection_count['equation']} Equations (RED)" # # cv2.rectangle(img, (10, 10), (10 + len(summary) * 10, 40), (0, 0, 0), -1) # # cv2.putText(img, summary, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # # # Save to temp file # # temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name # # cv2.imwrite(temp_path, img) # # return temp_path # # except Exception as e: # # print(f"Error in visualize_detections: {e}") # # import traceback # # traceback.print_exc() # # return None # # # ============================== # # # GRADIO INTERFACE # # # ============================== # # with gr.Blocks(title="Document Analysis Pipeline") as demo: # # gr.Markdown("# 📄 Full Pipeline Analysis") # # gr.Markdown("### 🔍 Intermediate File Recovery Active") # # gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)") # # with gr.Row(): # # with gr.Column(scale=1): # # file_input = gr.File( # # label="Upload PDFs or Images", # # file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"], # # file_count="multiple", # # type="filepath" # # ) # # page_selector = gr.Slider( # # minimum=1, # # maximum=100, # # value=1, # # step=1, # # label="PDF Page Number (for preview)", # # visible=True # # ) # # model_path_input = gr.Textbox( # # label="Model Path", # # value=DEFAULT_LAYOUTLMV3_MODEL_PATH # # ) # # # Debug button for visual inspection # # debug_btn = gr.Button("🔍 Show YOLO Detections (First Page)", variant="secondary") # # # Main processing button # # process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary") # # with gr.Column(scale=2): # # # Visual debug output # # detection_preview = gr.Image(label="YOLO Detection Preview (Green=Figure, Red=Equation)", type="filepath") # # # Final JSON output # # json_output = gr.Code(label="Final Structured Output", language="json", lines=20) # # # Download all intermediate files # # download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple") # # # Wire up the debug button # # debug_btn.click( # # fn=visualize_detections, # # inputs=[file_input, page_selector], # # outputs=[detection_preview] # # ) # # # Wire up the main processing button # # process_btn.click( # # fn=process_file, # # inputs=[file_input, model_path_input], # # outputs=[json_output, download_output] # # ) # # if __name__ == "__main__": # # demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) # # ============================== # # VISUAL DEBUG FUNCTION # # ============================== # def visualize_detections(uploaded_files, page_num): # """Shows the selected PDF page or image with YOLO bounding boxes""" # if not uploaded_files: # return None # try: # import cv2 # import numpy as np # import tempfile # from ultralytics import YOLO # import fitz # # Get first file path # file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files # if isinstance(file_path, dict): # file_path = file_path["path"] # elif hasattr(file_path, 'path'): # file_path = file_path.path # # Handle PDF conversion to image # if str(file_path).lower().endswith('.pdf'): # doc = fitz.open(file_path) # # Ensure the selected page exists in the document # page_idx = min(max(int(page_num) - 1, 0), len(doc) - 1) # page = doc.load_page(page_idx) # pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) # img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) # if pix.n == 3: # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # elif pix.n == 4: # img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) # doc.close() # else: # img = cv2.imread(str(file_path)) # if img is None: # return None # # Run YOLO detection # model = YOLO(WEIGHTS_PATH) # results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False) # # Draw bounding boxes # detection_count = {'figure': 0, 'equation': 0} # for box in results[0].boxes: # class_id = int(box.cls[0]) # class_name = model.names[class_id] # if class_name in ['figure', 'equation']: # x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy()) # conf = float(box.conf[0]) # # Green for figures, Red for equations # color = (0, 255, 0) if class_name == 'figure' else (0, 0, 255) # cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) # # Add label with confidence # label = f"{class_name.upper()} {conf:.2f}" # (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) # cv2.rectangle(img, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1) # cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) # detection_count[class_name] += 1 # # Add summary text at top # summary = f"Page {page_num} | Detected: {detection_count['figure']} Figures, {detection_count['equation']} Equations" # cv2.rectangle(img, (10, 10), (10 + len(summary) * 11, 40), (0, 0, 0), -1) # cv2.putText(img, summary, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # # Save to temp file # temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name # cv2.imwrite(temp_path, img) # return temp_path # except Exception as e: # print(f"Error in visualize_detections: {e}") # import traceback # traceback.print_exc() # return None # # ============================== # # GRADIO INTERFACE # # ============================== # with gr.Blocks(title="Document Analysis Pipeline") as demo: # gr.Markdown("# 📄 Full Pipeline Analysis") # gr.Markdown("### 🔍 Intermediate File Recovery Active") # gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)") # with gr.Row(): # with gr.Column(scale=1): # file_input = gr.File( # label="Upload PDFs or Images", # file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"], # file_count="multiple", # type="filepath" # ) # page_selector = gr.Slider( # minimum=1, # maximum=100, # value=1, # step=1, # label="PDF Page Number (for preview)", # visible=True # ) # model_path_input = gr.Textbox( # label="Model Path", # value=DEFAULT_LAYOUTLMV3_MODEL_PATH # ) # # Debug button for visual inspection # debug_btn = gr.Button("🔍 Show YOLO Detections", variant="secondary") # # Main processing button # process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary") # with gr.Column(scale=2): # # Visual debug output # detection_preview = gr.Image(label="YOLO Detection Preview (Green=Figure, Red=Equation)", type="filepath") # # Final JSON output # json_output = gr.Code(label="Final Structured Output", language="json", lines=20) # # Download all intermediate files # download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple") # # Wire up the debug button # debug_btn.click( # fn=visualize_detections, # inputs=[file_input, page_selector], # outputs=[detection_preview] # ) # # Wire up the main processing button # process_btn.click( # fn=process_file, # inputs=[file_input, model_path_input], # outputs=[json_output, download_output] # ) # if __name__ == "__main__": # demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) import gradio as gr import json import os import tempfile import img2pdf import glob import shutil from img2pdf import Rotation from pathlib import Path print("--- DEBUG: Current Working Directory ---") print(os.getcwd()) print("--- DEBUG: Files in Root ---") print(os.listdir('.')) # ============================== # PIPELINE IMPORT # ============================== try: from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH except Exception as e: # Catch ALL exceptions print(f"Warning: Failed to import pipeline: {e}") import traceback traceback.print_exc() # Show the actual error def run_document_pipeline(*args): yield {"status": "error", "message": "Placeholder pipeline function called."} DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model" WEIGHTS_PATH = "./weights/yolo_weights.pt" # ============================== # MAIN PROCESSING GENERATOR # ============================== def process_file(uploaded_files, layoutlmv3_model_path=None): """ Robust handler for multiple or single file uploads. Streams the estimation first, then yields the final JSON and intermediate files. """ if uploaded_files is None: yield "❌ Error: No files uploaded.", None return if not isinstance(uploaded_files, list): file_list = [uploaded_files] else: file_list = uploaded_files if len(file_list) == 0: yield "❌ Error: Empty file list.", None return # 1. Resolve all file paths safely resolved_paths = [] for f in file_list: try: if isinstance(f, dict) and "path" in f: resolved_paths.append(f["path"]) elif hasattr(f, 'path'): resolved_paths.append(f.path) else: resolved_paths.append(str(f)) except Exception as e: print(f"Error resolving path for {f}: {e}") if not resolved_paths: yield "❌ Error: Could not resolve file paths.", None return # 2. Determine if we should merge into a single PDF first_file = Path(resolved_paths[0]) is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff'] try: if len(resolved_paths) > 1 or is_image: print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF...") temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") with open(temp_pdf.name, "wb") as f_out: f_out.write(img2pdf.convert(resolved_paths, rotation=Rotation.ifvalid)) processing_path = temp_pdf.name else: processing_path = resolved_paths[0] # 3. Standard Pipeline Checks final_model_path = layoutlmv3_model_path or DEFAULT_LAYOUTLMV3_MODEL_PATH if not os.path.exists(final_model_path): yield f"❌ Error: Model not found at {final_model_path}", None return # 4. Call the pipeline generator print(f"🚀 Starting pipeline for: {processing_path}") # Iterate through the yields from run_document_pipeline for pipeline_update in run_document_pipeline(processing_path, final_model_path): # --- Handle Estimation Yield --- if pipeline_update.get("status") == "estimating": display_text = "⏱️ ESTIMATING PROCESSING TIME...\n\n" + json.dumps(pipeline_update, indent=2) yield display_text, None # --- Handle Final Complete Yield --- elif pipeline_update.get("status") == "complete": final_result = pipeline_update.get("result") # SCRAPE FOR INTERMEDIATE FILES base_name = Path(processing_path).stem search_patterns = [ f"/tmp/pipeline_run_{base_name}*/*.json", f"/tmp/*{base_name}*.json" ] all_intermediate_jsons = [] for pattern in search_patterns: all_intermediate_jsons.extend(glob.glob(pattern)) all_intermediate_jsons = list(dict.fromkeys(all_intermediate_jsons)) # Prepare Final Output for Display if final_result is None or (isinstance(final_result, list) and len(final_result) == 0): display_text = "⚠️ Pipeline failed at Step 3 (BIO Decoding).\nDownload the intermediate JSONs below to inspect OCR and Model Predictions." else: display_text = json.dumps(final_result, indent=2, ensure_ascii=False) # Save it to a temp file so it can be downloaded too temp_final = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='final_result_') json.dump(final_result, temp_final, indent=2, ensure_ascii=False) temp_final.close() all_intermediate_jsons.append(temp_final.name) yield display_text, all_intermediate_jsons # --- Handle Error Yield --- elif pipeline_update.get("status") == "error": yield f"❌ Error: {pipeline_update.get('message')}", None except Exception as e: import traceback traceback.print_exc() yield f"❌ Error: {str(e)}", None # ============================== # VISUAL DEBUG FUNCTION # ============================== def visualize_detections(uploaded_files, page_num): """Shows the selected PDF page or image with YOLO bounding boxes""" if not uploaded_files: return None try: import cv2 import numpy as np import tempfile from ultralytics import YOLO import fitz # Get first file path file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files if isinstance(file_path, dict): file_path = file_path["path"] elif hasattr(file_path, 'path'): file_path = file_path.path # Handle PDF conversion to image if str(file_path).lower().endswith('.pdf'): doc = fitz.open(file_path) # Ensure the selected page exists in the document page_idx = min(max(int(page_num) - 1, 0), len(doc) - 1) page = doc.load_page(page_idx) pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)) img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) if pix.n == 3: img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) elif pix.n == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) doc.close() else: img = cv2.imread(str(file_path)) if img is None: return None # Run YOLO detection model = YOLO(WEIGHTS_PATH) results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False) # Draw bounding boxes detection_count = {'figure': 0, 'equation': 0} for box in results[0].boxes: class_id = int(box.cls[0]) class_name = model.names[class_id] if class_name in ['figure', 'equation']: x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy()) conf = float(box.conf[0]) # Green for figures, Red for equations color = (0, 255, 0) if class_name == 'figure' else (0, 0, 255) cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) # Add label with confidence label = f"{class_name.upper()} {conf:.2f}" (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) cv2.rectangle(img, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1) cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) detection_count[class_name] += 1 # Add summary text at top summary = f"Page {page_num} | Detected: {detection_count['figure']} Figures, {detection_count['equation']} Equations" cv2.rectangle(img, (10, 10), (10 + len(summary) * 11, 40), (0, 0, 0), -1) cv2.putText(img, summary, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # Save to temp file temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name cv2.imwrite(temp_path, img) return temp_path except Exception as e: print(f"Error in visualize_detections: {e}") import traceback traceback.print_exc() return None # ============================== # GRADIO INTERFACE # ============================== with gr.Blocks(title="Document Analysis Pipeline") as demo: gr.Markdown("# 📄 Full Pipeline Analysis") gr.Markdown("### 🔍 Intermediate File Recovery Active") gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)") with gr.Row(): with gr.Column(scale=1): file_input = gr.File( label="Upload PDFs or Images", file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"], file_count="multiple", type="filepath" ) page_selector = gr.Slider( minimum=1, maximum=100, value=1, step=1, label="PDF Page Number (for preview)", visible=True ) model_path_input = gr.Textbox( label="Model Path", value=DEFAULT_LAYOUTLMV3_MODEL_PATH ) # Debug button for visual inspection debug_btn = gr.Button("🔍 Show YOLO Detections", variant="secondary") # Main processing button process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary") with gr.Column(scale=2): # Visual debug output detection_preview = gr.Image(label="YOLO Detection Preview (Green=Figure, Red=Equation)", type="filepath") # Final JSON output (Will update with estimation, then final result) json_output = gr.Code(label="Pipeline Output", language="json", lines=20) # Download all intermediate files download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple") # Wire up the debug button debug_btn.click( fn=visualize_detections, inputs=[file_input, page_selector], outputs=[detection_preview] ) # Wire up the main processing button process_btn.click( fn=process_file, inputs=[file_input, model_path_input], outputs=[json_output, download_output], api_name="process" # This enables the streaming endpoint /api/process ) if __name__ == "__main__": # IMPORTANT: .queue() is required for streaming generators to work! demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)