Update app.py

app.py

@@ -1,149 +1,8 @@
-# import gradio as gr
-# print("GRADIO VERSION:", gr.__version__)
-# import json
-# import os
-# import tempfile
-# from pathlib import Path
-
-# # ==============================
-# # 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.")
-#     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_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
-
-#     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:
-#         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}")
-
-#         # 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 document. Check console for details.", None
-
-#         # 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:
-#             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
-
-
-# # ==============================
-# # GRADIO INTERFACE
-# # ==============================
-# with gr.Blocks(title="Document Analysis Pipeline") as demo:
-
-#     gr.Markdown("""
-#     # 📄 Document & Image Analysis Pipeline
-
-#     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**: 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):
-#             file_input = gr.File(
-#                 label="Upload PDF or Image",
-#                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
-#                 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 File", variant="primary", size="lg")
-
-#             gr.Markdown("""
-#             ### ℹ️ Notes:
-#             - **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):
-#             json_output = gr.Code(
-#                 label="Structured JSON Output",
-#                 language="json",
-#                 lines=25
-#             )
-
-#             download_output = gr.File(
-#                 label="Download Full JSON",
-#                 interactive=False
-#             )
-
-#     # UI Logic
-#     process_btn.click(
-#         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",
-#         server_port=7860,
-#         share=False,
-#         show_error=True
-#     )
-
-
-
-
-
 import gradio as gr
-print("GRADIO VERSION:", gr.__version__)
 import json
 import os
 import tempfile
+import img2pdf
 from pathlib import Path
 
 # ==============================
@@ -158,71 +17,78 @@ except ImportError:
     DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
     WEIGHTS_PATH = "./weights/yolo_weights.pt"
 
-
-def process_file(uploaded_file, layoutlmv3_model_path=None):
+def process_file(uploaded_files, layoutlmv3_model_path=None):
     """
-    Handles both PDF and Image uploads and routes them to the YOLO/OCR pipeline.
+    Converts multiple images into a single PDF (if necessary) and routes 
+    the result to the YOLO/OCR pipeline as a single entity.
     """
-    if uploaded_file is None:
-        return "❌ Error: No file uploaded.", None
+    if not uploaded_files:
+        return "❌ Error: No files uploaded.", None
+
+    # Ensure we are working with a list of files (Gradio file_count="multiple" returns a list)
+    if not isinstance(uploaded_files, list):
+        uploaded_files = [uploaded_files]
+
+    # 1. Resolve all file paths
+    resolved_paths = []
+    for f in uploaded_files:
+        if hasattr(f, 'path'):
+            resolved_paths.append(f.path)
+        elif isinstance(f, dict):
+            resolved_paths.append(f.get("path"))
+        else:
+            resolved_paths.append(str(f))
 
-    # --- FIX FOR GRADIO 6.x FILE HANDLING ---
-    # If multiple files were somehow uploaded or Gradio returned a list
-    if isinstance(uploaded_file, list):
-        uploaded_file = uploaded_file[0]
+    # 2. Determine if we should merge into a single PDF
+    # We merge if there are multiple files OR if the single file is an image
+    first_file = Path(resolved_paths[0])
+    is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
+    
+    processing_path = None
     
-    # Extract the actual file path string. 
-    # Gradio File objects have a '.path' attribute for the temporary local location.
     try:
-        if hasattr(uploaded_file, 'path'):
-            file_path = uploaded_file.path
-        elif isinstance(uploaded_file, dict):
-            file_path = uploaded_file.get("path")
+        if len(resolved_paths) > 1 or (len(resolved_paths) == 1 and is_image):
+            print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF entity...")
+            temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+            # img2pdf.convert converts a list of image paths into PDF bytes
+            with open(temp_pdf.name, "wb") as f:
+                f.write(img2pdf.convert(resolved_paths))
+            processing_path = temp_pdf.name
         else:
-            file_path = str(uploaded_file)
-    except Exception as e:
-        return f"❌ Error resolving file path: {str(e)}", None
-    # ---------------------------------------
+            # It's a single PDF, process directly
+            processing_path = resolved_paths[0]
 
-    if not layoutlmv3_model_path:
-        layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
+        # 3. Standard Pipeline Checks
+        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(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
+        if not os.path.exists(WEIGHTS_PATH):
+            return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
 
-    try:
-        # Determine file type for logging safely
-        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}")
+        print(f"🚀 Starting pipeline for merged entity: {processing_path}")
 
-        # Call the pipeline
-        result = run_document_pipeline(file_path, layoutlmv3_model_path)
+        # 4. Call the pipeline
+        result = run_document_pipeline(processing_path, layoutlmv3_model_path)
 
         if result is None:
-            return "❌ Error: Pipeline failed to process the document. Check console for details.", None
+            return "❌ Error: Pipeline failed to process the document.", None
 
-        # Prepare output file for download
-        output_filename = f"{Path(file_path).stem}_analysis.json"
+        # 5. Prepare output
         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:
-        # This is where your previous error message was being caught and returned
         import traceback
-        traceback.print_exc() # This prints the full error to your terminal for debugging
+        traceback.print_exc()
         return f"❌ Error during processing: {str(e)}", None
 
-
 # ==============================
 # GRADIO INTERFACE
 # ==============================
@@ -230,16 +96,16 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
 
     gr.Markdown("""
     # 📄 Document & Image Analysis Pipeline
-    Upload a **PDF document** or an **Image (JPG/PNG)** to extract structured data.
+    Upload **multiple images** or a **PDF**. Multiple images will be processed together as a single continuous document.
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
             file_input = gr.File(
-                label="Upload PDF or Image",
+                label="Upload PDFs or Images",
                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
                 type="filepath",
-                file_count="single"  # Force single file to avoid list/tuple issues
+                file_count="multiple"  # ALLOWS MULTIPLE FILES
             )
 
             model_path_input = gr.Textbox(
@@ -249,11 +115,11 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
                 interactive=True
             )
 
-            process_btn = gr.Button("🚀 Process File", variant="primary", size="lg")
+            process_btn = gr.Button("🚀 Process Files", variant="primary", size="lg")
 
         with gr.Column(scale=2):
             json_output = gr.Code(
-                label="Structured JSON Output",
+                label="Combined Structured JSON Output",
                 language="json",
                 lines=25
             )