Update app.py

app.py

@@ -1,325 +1,82 @@
-# import gradio as gr
-# print("GRADIO VERSION:", gr.__version__)
-# import json
-# import os
-# import tempfile
-# from pathlib import Path
-
-# # NOTE: You must ensure that 'working_yolo_pipeline.py' exists 
-# # and defines the following items correctly:
-# from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-# # Since I don't have this file, I am assuming the imports are correct.
-
-# # Define placeholders for assumed constants if the pipeline file isn't present
-# # You should replace these with your actual definitions if they are missing
-# 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.")
-#     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_pdf(pdf_file, layoutlmv3_model_path=None):
-#     """
-#     Wrapper function for Gradio interface.
-
-#     Args:
-#         pdf_file: Gradio UploadButton file object
-#         layoutlmv3_model_path: Optional custom model path
-
-#     Returns:
-#         Tuple of (JSON string, download file path)
-#     """
-#     if pdf_file is None:
-#         return "❌ Error: No PDF file uploaded.", None
-
-#     # Use default model path if not provided
-#     if not layoutlmv3_model_path:
-#         layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
-
-#     # Verify model and weights exist
-#     if not os.path.exists(layoutlmv3_model_path):
-#         return f"❌ Error: LayoutLMv3 model not found at {layoutlmv3_model_path}", None
-
-#     if not os.path.exists(WEIGHTS_PATH):
-#         return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
-
-#     try:
-#         # Get the uploaded PDF path
-#         pdf_path = pdf_file.name
-
-#         # Run the pipeline
-#         result = run_document_pipeline(pdf_path, layoutlmv3_model_path, 'label_studio_import.json')
-
-#         if result is None:
-#             return "❌ Error: Pipeline failed to process the PDF. Check console for details.", None
-
-#         # Create a temporary file for download
-#         output_filename = f"{Path(pdf_path).stem}_analysis.json"
-#         temp_output = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='analysis_')
-
-#         # Dump results to the temporary file
-#         with open(temp_output.name, 'w', encoding='utf-8') as f:
-#             json.dump(result, f, indent=2, ensure_ascii=False)
-
-#         # Format JSON for display
-#         json_display = json.dumps(result, indent=2, ensure_ascii=False)
-
-#         return json_display, temp_output.name
-
-#     except Exception as e:
-#         return f"❌ Error during processing: {str(e)}", None
-
-
-# # Create Gradio interface
-# # FIX APPLIED: Removed 'theme=gr.themes.Soft()' which caused the TypeError
-# with gr.Blocks(title="Document Analysis Pipeline") as demo:
-#     gr.Markdown("""
-#     # 📄 Document Analysis Pipeline
-
-#     Upload a PDF document to extract structured data including questions, options, answers, passages, and embedded images.
-
-#     **Pipeline Steps:**
-#     1. 🔍 YOLO/OCR Preprocessing (word extraction + figure/equation detection)
-#     2. 🤖 LayoutLMv3 Inference (BIO tagging)
-#     3. 📊 Structured JSON Decoding
-#     4. 🖼️ Base64 Image Embedding
-#     """)
-
-#     with gr.Row():
-#         with gr.Column(scale=1):
-#             pdf_input = gr.File(
-#                 label="Upload PDF Document",
-#                 file_types=[".pdf"],
-#                 type="filepath"
-#             )
-
-#             model_path_input = gr.Textbox(
-#                 label="LayoutLMv3 Model Path (optional)",
-#                 placeholder=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-#                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-#                 interactive=True
-#             )
-
-#             process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
-
-#             gr.Markdown("""
-#             ### ℹ️ Notes:
-#             - Processing may take several minutes depending on PDF size
-#             - Figures and equations will be extracted and embedded as Base64
-#             - The output JSON includes structured questions, options, and answers
-#             """)
-
-#         with gr.Column(scale=2):
-#             json_output = gr.Code(
-#                 label="Structured JSON Output",
-#                 language="json",
-#                 lines=25
-#             )
-
-#             download_output = gr.File(
-#                 label="Download Full JSON",
-#                 interactive=False
-#             )
-
-#     # Status/Examples section
-#     with gr.Row():
-#         gr.Markdown("""
-#         ### 📋 Output Format
-#         The pipeline generates JSON with the following structure:
-#         - **Questions**: Extracted question text
-#         - **Options**: Multiple choice options (A, B, C, D, etc.)
-#         - **Answers**: Correct answer(s)
-#         - **Passages**: Associated reading passages
-#         - **Images**: Base64-encoded figures and equations (embedded with keys like `figure1`, `equation2`)
-#         """)
-
-#     # Connect the button to the processing function
-#     process_btn.click(
-#         fn=process_pdf,
-#         inputs=[pdf_input, model_path_input],
-#         outputs=[json_output, download_output],
-#         api_name="process_document"
-#     )
-
-#     # Example section (optional - add example PDFs if available)
-#     # gr.Examples(
-#     #     examples=[
-#     #         ["examples/sample1.pdf"],
-#     #         ["examples/sample2.pdf"],
-#     #     ],
-#     #     inputs=pdf_input,
-#     # )
-
-# # Launch the app
-# if __name__ == "__main__":
-#     demo.launch(
-#         server_name="0.0.0.0",
-#         server_port=7860,
-#         share=False,
-#         show_error=True
-#     )
-
-
-
-
-
 import gradio as gr
-print("GRADIO VERSION:", gr.__version__)
-import json
+import cv2
+import numpy as np
+from paddleocr import PPStructure
+from huggingface_hub import snapshot_download
 import os
-import tempfile
-from pathlib import Path
-
-# ==============================
-# WRITE CUSTOM CSS FOR FONTS
-# ==============================
-
-# CUSTOM_CSS = """
-# @font-face {
-#     font-family: 'NotoSansMath';
-#     src: url('./NotoSansMath-Regular.ttf') format('truetype');
-#     font-weight: normal;
-#     font-style: normal;
-# }
-
-# html, body, * {
-#     font-family: 'NotoSansMath', sans-serif !important;
-# }
-# """
-
-# # Optionally write the CSS file if needed (not required for inline css)
-# if not os.path.exists("custom.css"):
-#     with open("custom.css", "w") as f:
-#         f.write(CUSTOM_CSS)
-# ==============================
-
-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.")
-    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_pdf(pdf_file, layoutlmv3_model_path=None):
-    if pdf_file is None:
-        return "❌ Error: No PDF file uploaded.", None
-
-    if not layoutlmv3_model_path:
-        layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
-
-    if not os.path.exists(layoutlmv3_model_path):
-        return f"❌ Error: LayoutLMv3 model not found at {layoutlmv3_model_path}", None
-
-    if not os.path.exists(WEIGHTS_PATH):
-        return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
-
-    try:
-        pdf_path = pdf_file.name
-
-        result = run_document_pipeline(pdf_path, layoutlmv3_model_path, 'label_studio_import.json')
-
-        if result is None:
-            return "❌ Error: Pipeline failed to process the PDF. Check console for details.", None
-
-        output_filename = f"{Path(pdf_path).stem}_analysis.json"
-        temp_output = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='analysis_')
-
-        with open(temp_output.name, 'w', encoding='utf-8') as f:
-            json.dump(result, f, indent=2, ensure_ascii=False)
-
-        json_display = json.dumps(result, indent=2, ensure_ascii=False)
-
-        return json_display, temp_output.name
-
-    except Exception as e:
-        return f"❌ Error during processing: {str(e)}", None
-
-
-with gr.Blocks(
-    title="Document Analysis Pipeline"
-) as demo:
-
-
-    gr.HTML()
-
-    gr.Markdown("""
-    # 📄 Document Analysis Pipeline
-
-    Upload a PDF document to extract structured data including questions, options, answers, passages, and embedded images.
-
-    **Pipeline Steps:**
-    1. 🔍 YOLO/OCR Preprocessing (word extraction + figure/equation detection)
-    2. 🤖 LayoutLMv3 Inference (BIO tagging)
-    3. 📊 Structured JSON Decoding
-    4. 🖼️ Base64 Image Embedding
-    """)
 
+# --- STEP 1: Download the Model from Hugging Face ---
+# We download the 'main' branch which contains the Paddle inference weights
+print("Downloading PP-DocLayoutV3 from Hugging Face...")
+model_path = snapshot_download(repo_id="PaddlePaddle/PP-DocLayoutV3", allow_patterns=["*.pdiparams", "*.pdmodel", "*.yml", "*.json"])
+print(f"Model downloaded to: {model_path}")
+
+# --- STEP 2: Initialize the Layout Engine ---
+# We use PPStructure, which is PaddleOCR's layout analysis module.
+# We point it to the downloaded model folder.
+layout_engine = PPStructure(
+    layout_model_dir=model_path,
+    table=False,          # Disable table structure recognition for speed
+    ocr=False,            # Disable OCR for now (we just want to see layout)
+    show_log=True,
+    use_angle_cls=True,   # Helps with orientation
+    enable_mkldnn=False   # CRITICAL: Fixes the CPU crash
+)
+
+def analyze_layout(input_image):
+    if input_image is None:
+        return None, "No image uploaded"
+
+    image_np = np.array(input_image)
+    
+    # Run Inference
+    # result is a list of dictionaries, one per detected region
+    result = layout_engine(image_np)
+
+    viz_image = image_np.copy()
+    detections_text = []
+
+    # --- STEP 3: Visualize Results ---
+    for region in result:
+        # Extract Box (4 points)
+        box = region['layout_bbox']
+        label = region['label']
+        
+        # Convert to numpy format for drawing
+        # layout_bbox is usually [x1, y1, x2, y2]
+        x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
+        
+        # Color coding based on type
+        color = (0, 255, 0) # Green for Text
+        if label == 'title': color = (0, 0, 255) # Red for Title
+        elif label == 'figure': color = (255, 0, 0) # Blue for Figures
+        elif label == 'table': color = (255, 255, 0) # Cyan for Tables
+
+        # Draw Rectangle
+        cv2.rectangle(viz_image, (x1, y1), (x2, y2), color, 3)
+        
+        # Draw Label
+        cv2.putText(viz_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
+        
+        detections_text.append(f"Found {label} at {box}")
+
+    return viz_image, "\n".join(detections_text)
+
+# --- Gradio UI ---
+with gr.Blocks(title="PP-DocLayoutV3 Demo") as demo:
+    gr.Markdown("## 📄 PP-DocLayoutV3 Explorer")
+    gr.Markdown("This model detects **layout regions** (Text, Tables, Titles) instead of reading characters. It is excellent for de-warping and segmenting messy documents.")
+    
     with gr.Row():
-        with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="Upload PDF Document",
-                file_types=[".pdf"],
-                type="filepath"
-            )
-
-            model_path_input = gr.Textbox(
-                label="LayoutLMv3 Model Path (optional)",
-                placeholder=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-                value=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-                interactive=True
-            )
-
-            process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
-
-            gr.Markdown("""
-            ### ℹ️ Notes:
-            - Processing may take several minutes depending on PDF size
-            - Figures and equations will be extracted and embedded as Base64
-            - The output JSON includes structured questions, options, and answers
-            """)
-
-        with gr.Column(scale=2):
-            json_output = gr.Code(
-                label="Structured JSON Output",
-                language="json",
-                lines=25
-            )
-
-            download_output = gr.File(
-                label="Download Full JSON",
-                interactive=False
-            )
-
-    with gr.Row():
-        gr.Markdown("""
-        ### 📋 Output Format
-        The pipeline generates JSON with the following structure:
-        - **Questions**: Extracted question text
-        - **Options**: Multiple choice options
-        - **Answers**: Correct answer(s)
-        - **Passages**: Associated reading passages
-        - **Images**: Base64-encoded figures and equations
-        """)
-
-    process_btn.click(
-        fn=process_pdf,
-        inputs=[pdf_input, model_path_input],
-        outputs=[json_output, download_output],
-        api_name="process_document"
-    )
+        with gr.Column():
+            input_img = gr.Image(type="pil", label="Input Document")
+            submit_btn = gr.Button("Analyze Layout", variant="primary")
+        
+        with gr.Column():
+            output_img = gr.Image(label="Layout Visualization")
+            output_log = gr.Textbox(label="Detected Regions", lines=10)
 
+    submit_btn.click(fn=analyze_layout, inputs=input_img, outputs=[output_img, output_log])
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860)