Update app.py

app.py

@@ -1,172 +1,3 @@
-# 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
@@ -175,28 +6,8 @@ import tempfile
 from pathlib import Path
 
 # ==============================
-# WRITE CUSTOM CSS FOR FONTS
+# PIPELINE IMPORT
 # ==============================
-
-# 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:
@@ -207,9 +18,12 @@ except ImportError:
     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
+def process_file(uploaded_file, layoutlmv3_model_path=None):
+    """
+    Handles both PDF and Image uploads and routes them to the YOLO/OCR pipeline.
+    """
+    if uploaded_file is None:
+        return "❌ Error: No file uploaded.", None
 
     if not layoutlmv3_model_path:
         layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
@@ -221,14 +35,22 @@ def process_pdf(pdf_file, layoutlmv3_model_path=None):
         return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
 
     try:
-        pdf_path = pdf_file.name
+        file_path = uploaded_file.name
+        
+        # Determine file type for logging
+        ext = Path(file_path).suffix.lower()
+        file_type = "Image" if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'] else "PDF"
+        print(f"🚀 Starting pipeline for {file_type}: {file_path}")
 
-        result = run_document_pipeline(pdf_path, layoutlmv3_model_path, 'label_studio_import.json')
+        # Call the pipeline exactly as before. 
+        # Our modified working_yolo_pipeline now handles the branching internally.
+        result = run_document_pipeline(file_path, layoutlmv3_model_path)
 
         if result is None:
-            return "❌ Error: Pipeline failed to process the PDF. Check console for details.", None
+            return "❌ Error: Pipeline failed to process the document. Check console for details.", None
 
-        output_filename = f"{Path(pdf_path).stem}_analysis.json"
+        # Prepare output file for download
+        output_filename = f"{Path(file_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:
@@ -242,30 +64,30 @@ def process_pdf(pdf_file, layoutlmv3_model_path=None):
         return f"❌ Error during processing: {str(e)}", None
 
 
-with gr.Blocks(
-    title="Document Analysis Pipeline"
-) as demo:
-
-
-    gr.HTML()
+# ==============================
+# GRADIO INTERFACE
+# ==============================
+with gr.Blocks(title="Document Analysis Pipeline") as demo:
 
     gr.Markdown("""
-    # 📄 Document Analysis Pipeline
+    # 📄 Document & Image Analysis Pipeline
 
-    Upload a PDF document to extract structured data including questions, options, answers, passages, and embedded images.
+    Upload a **PDF document** or an **Image (JPG/PNG)** to extract structured data.
+    
+    **Supported Formats:** `.pdf`, `.jpg`, `.jpeg`, `.png`, `.bmp`, `.webp`
 
     **Pipeline Steps:**
-    1. 🔍 YOLO/OCR Preprocessing (word extraction + figure/equation detection)
-    2. 🤖 LayoutLMv3 Inference (BIO tagging)
-    3. 📊 Structured JSON Decoding
-    4. 🖼️ Base64 Image Embedding
+    1. 🔍 **YOLO/OCR**: Word extraction + Figure/Equation detection
+    2. 🤖 **LayoutLMv3**: BIO tagging and structural analysis
+    3. 📊 **Decoding**: Conversion to hierarchical JSON
+    4. 🖼️ **Extraction**: Base64 embedding of detected visual elements
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="Upload PDF Document",
-                file_types=[".pdf"],
+            file_input = gr.File(
+                label="Upload PDF or Image",
+                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
                 type="filepath"
             )
 
@@ -276,13 +98,13 @@ with gr.Blocks(
                 interactive=True
             )
 
-            process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
+            process_btn = gr.Button("🚀 Process File", 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
+            - **Images** are treated as single-page documents.
+            - **PDFs** are processed page-by-page.
+            - High-resolution Tesseract OCR is used for all image content.
             """)
 
         with gr.Column(scale=2):
@@ -297,25 +119,14 @@ with gr.Blocks(
                 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
-        """)
-
+    # UI Logic
     process_btn.click(
-        fn=process_pdf,
-        inputs=[pdf_input, model_path_input],
+        fn=process_file,
+        inputs=[file_input, model_path_input],
         outputs=[json_output, download_output],
         api_name="process_document"
     )
 
-
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",