Update app.py

app.py

@@ -1,45 +1,113 @@
 import gradio as gr
+import torch
+import logging
+import os
+import json
 from PIL import Image
-import io
 from surya.ocr import run_ocr
+from surya.detection import batch_text_detection
+from surya.layout import batch_layout_detection
+from surya.ordering import batch_ordering
 from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
 from surya.model.recognition.model import load_model as load_rec_model
 from surya.model.recognition.processor import load_processor as load_rec_processor
+from surya.model.ordering.model import load_model as load_order_model
+from surya.model.ordering.processor import load_processor as load_order_processor
+from surya.settings import settings
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # Load models and processors
+logger.info("Loading models and processors...")
 det_processor, det_model = load_det_processor(), load_det_model()
 rec_model, rec_processor = load_rec_model(), load_rec_processor()
+layout_model = load_det_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+layout_processor = load_det_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+order_model = load_order_model()
+order_processor = load_order_processor()
 
-def perform_ocr(image, language):
-    # Convert gradio image to PIL Image
-    if image is not None:
-        image = Image.fromarray(image)
-    else:
-        return "No image uploaded"
-
-    # Perform OCR
-    langs = [language]  # You can expand this to support multiple languages
-    predictions = run_ocr([image], [langs], det_model, det_processor, rec_model, rec_processor)
+# Compile the OCR model for better performance
+logger.info("Compiling OCR model...")
+os.environ['RECOGNITION_STATIC_CACHE'] = 'true'
+rec_model.decoder.model = torch.compile(rec_model.decoder.model)
 
-    # Extract text from predictions
-    result = ""
-    for page in predictions[0]:  # Assuming single image input
-        for line in page['text_lines']:
-            result += line['text'] + "\n"
+def process_image(image_path, langs):
+    logger.info(f"Processing image: {image_path}")
+    image = Image.open(image_path)
+    
+    # OCR
+    logger.info("Performing OCR...")
+    ocr_predictions = run_ocr([image], [langs.split(',')], det_model, det_processor, rec_model, rec_processor)
+    
+    # Text line detection
+    logger.info("Detecting text lines...")
+    line_predictions = batch_text_detection([image], det_model, det_processor)
+    
+    # Layout analysis
+    logger.info("Analyzing layout...")
+    layout_predictions = batch_layout_detection([image], layout_model, layout_processor, line_predictions)
+    
+    # Reading order
+    logger.info("Determining reading order...")
+    bboxes = [bbox['bbox'] for bbox in layout_predictions[0]['bboxes']]
+    order_predictions = batch_ordering([image], [bboxes], order_model, order_processor)
+    
+    # Combine results
+    results = {
+        "ocr": ocr_predictions[0],
+        "text_lines": line_predictions[0],
+        "layout": layout_predictions[0],
+        "reading_order": order_predictions[0]
+    }
+    
+    logger.info("Processing complete.")
+    return json.dumps(results, indent=2)
 
-    return result
+def surya_ui(image, langs):
+    if image is None:
+        return "Please upload an image."
+    
+    try:
+        result = process_image(image, langs)
+        return result
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        return f"An error occurred: {str(e)}"
 
-# Define the Gradio interface
+# Create Gradio interface
 iface = gr.Interface(
-    fn=perform_ocr,
+    fn=surya_ui,
     inputs=[
-        gr.Image(type="numpy", label="Upload an image"),
-        gr.Dropdown(choices=["en", "fr", "de", "es", "it"], label="Select language", value="en")
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Textbox(label="Languages (comma-separated, e.g., 'en,fr')", value="en")
     ],
-    outputs=gr.Textbox(label="Extracted Text"),
-    title="OCR with Surya",
-    description="Upload an image to extract text using Optical Character Recognition."
+    outputs=gr.Textbox(label="Results"),
+    title="Surya Document Analysis",
+    description="Upload an image to perform OCR, text line detection, layout analysis, and reading order detection.",
+    theme="huggingface",
+    css="""
+    .gradio-container {
+        font-family: 'IBM Plex Sans', sans-serif;
+    }
+    .gr-button {
+        color: white;
+        border-radius: 8px;
+        background: linear-gradient(45deg, #ff9a9e 0%, #fad0c4 99%, #fad0c4 100%);
+    }
+    .gr-button:hover {
+        background: linear-gradient(45deg, #fad0c4 0%, #ff9a9e 99%, #ff9a9e 100%);
+    }
+    .gr-form {
+        border-radius: 12px;
+        background-color: #ffffff;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    }
+    """
 )
 
-# Launch the app
-iface.launch()
+# Launch the interface
+if __name__ == "__main__":
+    logger.info("Starting Gradio interface...")
+    iface.launch()