Update app.py

app.py

@@ -1,50 +1,30 @@
-import io
+import spaces
+import json
+import math
 import os
-import tempfile
-import time
-import uuid
-import cv2
+import traceback
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+import re
+
+import fitz
 import gradio as gr
-import pymupdf
-import spaces
+import requests
 import torch
-from PIL import Image, ImageDraw, ImageFont
-from transformers import AutoModelForCausalLM, AutoProcessor, VisionEncoderDecoderModel
 from huggingface_hub import snapshot_download
-from qwen_vl_utils import process_vision_info
-from utils.utils import prepare_image, parse_layout_string, process_coordinates, ImageDimensions
-from utils.markdown_utils import MarkdownConverter
-
-# Define device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load dot.ocr model
-dot_ocr_model_id = "rednote-hilab/dots.ocr"
-dot_ocr_model = AutoModelForCausalLM.from_pretrained(
-    dot_ocr_model_id,
-    attn_implementation="flash_attention_2",
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    trust_remote_code=True
-)
-dot_ocr_processor = AutoProcessor.from_pretrained(
-    dot_ocr_model_id,
-    trust_remote_code=True
+from PIL import Image, ImageDraw, ImageFont
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+    TextIteratorStreamer,
 )
-
-# Load Dolphin model
-dolphin_model_id = "ByteDance/Dolphin"
-dolphin_processor = AutoProcessor.from_pretrained(dolphin_model_id)
-dolphin_model = VisionEncoderDecoderModel.from_pretrained(dolphin_model_id)
-dolphin_model.eval()
-dolphin_model.to(device)
-dolphin_model = dolphin_model.half()
-dolphin_tokenizer = dolphin_processor.tokenizer
+from qwen_vl_utils import process_vision_info
 
 # Constants
 MIN_PIXELS = 3136
 MAX_PIXELS = 11289600
 IMAGE_FACTOR = 28
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Prompts
 prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
@@ -66,11 +46,55 @@ prompt = """Please output the layout information from the PDF image, including e
 5. Final Output: The entire output must be a single JSON object.
 """
 
+# Load Camel-Doc-OCR-062825
+MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825"
+processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
+model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_M,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# Load Megalodon-OCR-Sync-0713
+MODEL_ID_T = "prithivMLmods/Megalodon-OCR-Sync-0713"
+processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
+model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_T,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# Load Nanonets-OCR-s
+MODEL_ID_C = "nanonets/Nanonets-OCR-s"
+processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True)
+model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_C,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# Load MonkeyOCR
+MODEL_ID_G = "echo840/MonkeyOCR"
+SUBFOLDER = "Recognition"
+processor_g = AutoProcessor.from_pretrained(
+    MODEL_ID_G,
+    trust_remote_code=True,
+    subfolder=SUBFOLDER
+)
+model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_G,
+    trust_remote_code=True,
+    subfolder=SUBFOLDER,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+
 # Utility functions
 def round_by_factor(number: int, factor: int) -> int:
     """Returns the closest integer to 'number' that is divisible by 'factor'."""
     return round(number / factor) * factor
 
+
 def smart_resize(
     height: int,
     width: int,
@@ -100,6 +124,7 @@ def smart_resize(
         w_bar = round_by_factor(width * beta, factor)
     return h_bar, w_bar
 
+
 def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
     """Fetch and process an image"""
     if isinstance(image_input, str):
@@ -112,29 +137,31 @@ def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
         image = image_input.convert('RGB')
     else:
         raise ValueError(f"Invalid image input type: {type(image_input)}")
-    
+
     if min_pixels is not None or max_pixels is not None:
         min_pixels = min_pixels or MIN_PIXELS
         max_pixels = max_pixels or MAX_PIXELS
         height, width = smart_resize(
-            image.height, 
-            image.width, 
+            image.height,
+            image.width,
             factor=IMAGE_FACTOR,
             min_pixels=min_pixels,
             max_pixels=max_pixels
         )
         image = image.resize((width, height), Image.LANCZOS)
-    
+
     return image
 
+
 def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
     """Load images from PDF file"""
     images = []
     try:
-        pdf_document = pymupdf.open(pdf_path)
+        pdf_document = fitz.open(pdf_path)
         for page_num in range(len(pdf_document)):
             page = pdf_document.load_page(page_num)
-            mat = pymupdf.Matrix(2.0, 2.0)  # Increase resolution
+            # Convert page to image
+            mat = fitz.Matrix(2.0, 2.0)  # Increase resolution
             pix = page.get_pixmap(matrix=mat)
             img_data = pix.tobytes("ppm")
             image = Image.open(BytesIO(img_data)).convert('RGB')
@@ -145,14 +172,16 @@ def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
         return []
     return images
 
+
 def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.Image:
     """Draw layout bounding boxes on image"""
     img_copy = image.copy()
     draw = ImageDraw.Draw(img_copy)
-    
+
+    # Colors for different categories
     colors = {
         'Caption': '#FF6B6B',
-        'Footnote': '#4ECDC4', 
+        'Footnote': '#4ECDC4',
         'Formula': '#45B7D1',
         'List-item': '#96CEB4',
         'Page-footer': '#FFEAA7',
@@ -163,58 +192,134 @@ def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.I
         'Text': '#74B9FF',
         'Title': '#E17055'
     }
-    
+
     try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
-    except Exception:
-        font = ImageFont.load_default()
-    
-    for item in layout_data:
-        if 'bbox' in item and 'category' in item:
-            bbox = item['bbox']
-            category = item['category']
-            color = colors.get(category, '#000000')
-            draw.rectangle(bbox, outline=color, width=2)
-            label = category
-            label_bbox = draw.textbbox((0, 0), label, font=font)
-            label_width = label_bbox[2] - label_bbox[0]
-            label_height = label_bbox[3] - label_bbox[1]
-            label_x = bbox[0]
-            label_y = max(0, bbox[1] - label_height - 2)
-            draw.rectangle(
-                [label_x, label_y, label_x + label_width + 4, label_y + label_height + 2],
-                fill=color
-            )
-            draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
+        # Load a font
+        try:
+            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
+        except Exception:
+            font = ImageFont.load_default()
+
+        for item in layout_data:
+            if 'bbox' in item and 'category' in item:
+                bbox = item['bbox']
+                category = item['category']
+                color = colors.get(category, '#000000')
+
+                # Draw rectangle
+                draw.rectangle(bbox, outline=color, width=2)
+
+                # Draw label
+                label = category
+                label_bbox = draw.textbbox((0, 0), label, font=font)
+                label_width = label_bbox[2] - label_bbox[0]
+                label_height = label_bbox[3] - label_bbox[1]
+
+                # Position label above the box
+                label_x = bbox[0]
+                label_y = max(0, bbox[1] - label_height - 2)
+
+                # Draw background for label
+                draw.rectangle(
+                    [label_x, label_y, label_x + label_width + 4, label_y + label_height + 2],
+                    fill=color
+                )
+
+                # Draw text
+                draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
+
+    except Exception as e:
+        print(f"Error drawing layout: {e}")
+
     return img_copy
 
+
+def is_arabic_text(text: str) -> bool:
+    """Check if text in headers and paragraphs contains mostly Arabic characters"""
+    if not text:
+        return False
+
+    # Extract text from headers and paragraphs only
+    # Match markdown headers (# ## ###) and regular paragraph text
+    header_pattern = r'^#{1,6}\s+(.+)$'
+    paragraph_pattern = r'^(?!#{1,6}\s|!\[|```|\||\s*[-*+]\s|\s*\d+\.\s)(.+)$'
+
+    content_text = []
+
+    for line in text.split('\n'):
+        line = line.strip()
+        if not line:
+            continue
+
+        # Check for headers
+        header_match = re.match(header_pattern, line, re.MULTILINE)
+        if header_match:
+            content_text.append(header_match.group(1))
+            continue
+
+        # Check for paragraph text (exclude lists, tables, code blocks, images)
+        if re.match(paragraph_pattern, line, re.MULTILINE):
+            content_text.append(line)
+
+    if not content_text:
+        return False
+
+    # Join all content text and check for Arabic characters
+    combined_text = ' '.join(content_text)
+
+    # Arabic Unicode ranges
+    arabic_chars = 0
+    total_chars = 0
+
+    for char in combined_text:
+        if char.isalpha():
+            total_chars += 1
+            # Arabic script ranges
+            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
+                arabic_chars += 1
+
+    if total_chars == 0:
+        return False
+
+    # Consider text as Arabic if more than 50% of alphabetic characters are Arabic
+    return (arabic_chars / total_chars) > 0.5
+
+
 def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
     """Convert layout JSON to markdown format"""
     import base64
     from io import BytesIO
-    
+
     markdown_lines = []
-    
+
     try:
-        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
-        
+        # Sort items by reading order (top to bottom, left to right)
+        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox',), x.get('bbox',)))
+
         for item in sorted_items:
             category = item.get('category', '')
             text = item.get(text_key, '')
             bbox = item.get('bbox', [])
-            
+
             if category == 'Picture':
+                # Extract image region and embed it
                 if bbox and len(bbox) == 4:
                     try:
+                        # Extract the image region
                         x1, y1, x2, y2 = bbox
+                        # Ensure coordinates are within image bounds
                         x1, y1 = max(0, int(x1)), max(0, int(y1))
                         x2, y2 = min(image.width, int(x2)), min(image.height, int(y2))
-                        
+
                         if x2 > x1 and y2 > y1:
                             cropped_img = image.crop((x1, y1, x2, y2))
+
+                            # Convert to base64 for embedding
                             buffer = BytesIO()
                             cropped_img.save(buffer, format='PNG')
                             img_data = base64.b64encode(buffer.getvalue()).decode()
+
+                            # Add as markdown image
                             markdown_lines.append(f"![Image](data:image/png;base64,{img_data})\n")
                         else:
                             markdown_lines.append("![Image](Image region detected)\n")
@@ -234,11 +339,13 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
             elif category == 'List-item':
                 markdown_lines.append(f"- {text}\n")
             elif category == 'Table':
+                # If text is already HTML, keep it as is
                 if text.strip().startswith('<'):
                     markdown_lines.append(f"{text}\n")
                 else:
                     markdown_lines.append(f"**Table:** {text}\n")
             elif category == 'Formula':
+                # If text is LaTeX, format it properly
                 if text.strip().startswith('$') or '\\' in text:
                     markdown_lines.append(f"$$\n{text}\n$$\n")
                 else:
@@ -248,16 +355,21 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
             elif category == 'Footnote':
                 markdown_lines.append(f"^{text}^\n")
             elif category in ['Page-header', 'Page-footer']:
+                # Skip headers and footers in main content
                 continue
             else:
                 markdown_lines.append(f"{text}\n")
-            markdown_lines.append("")
+
+            markdown_lines.append("")  # Add spacing
+
     except Exception as e:
         print(f"Error converting to markdown: {e}")
         return str(layout_data)
+
     return "\n".join(markdown_lines)
 
-# Global state variables
+
+# PDF handling state
 pdf_cache = {
     "images": [],
     "current_page": 0,
@@ -266,60 +378,74 @@ pdf_cache = {
     "is_parsed": False,
     "results": []
 }
-
-@spaces.GPU()
-def dot_ocr_inference(image: Image.Image, prompt: str, max_new_tokens: int = 24000) -> str:
-    """Run inference on an image with the given prompt using dot.ocr model"""
+@spaces.GPU
+def inference(model_name: str, image: Image.Image, prompt: str, max_new_tokens: int = 1024) -> str:
+    """Run inference on an image with the given prompt using the selected model."""
     try:
+        if model_name == "Camel-Doc-OCR-062825":
+            processor = processor_m
+            model = model_m
+        elif model_name == "Megalodon-OCR-Sync-0713":
+            processor = processor_t
+            model = model_t
+        elif model_name == "Nanonets-OCR-s":
+            processor = processor_c
+            model = model_c
+        elif model_name == "MonkeyOCR-Recognition":
+            processor = processor_g
+            model = model_g
+        else:
+            raise ValueError(f"Invalid model selected: {model_name}")
+
         messages = [
             {
                 "role": "user",
                 "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": prompt}
+                    {"type": "text", "text": prompt},
+                    {"type": "image"}
                 ]
             }
         ]
-        text = dot_ocr_processor.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        image_inputs, video_inputs = process_vision_info(messages)
-        inputs = dot_ocr_processor(
-            text=[text],
-            images=image_inputs,
-            videos=video_inputs,
-            padding=True,
-            return_tensors="pt",
-        )
-        inputs = inputs.to(device)
+
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text, [image], return_tensors="pt").to(device)
+
         with torch.no_grad():
-            generated_ids = dot_ocr_model.generate(
+            generated_ids = model.generate(
                 **inputs,
                 max_new_tokens=max_new_tokens,
                 do_sample=False,
                 temperature=0.1
             )
-        generated_ids_trimmed = [
+
+        generated_ids = [
             out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
         ]
-        output_text = dot_ocr_processor.batch_decode(
-            generated_ids_trimmed,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False
-        )
-        return output_text[0] if output_text else ""
+        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
+        return output_text
+
     except Exception as e:
-        print(f"Error during dot.ocr inference: {e}")
+        print(f"Error during inference: {e}")
+        traceback.print_exc()
         return f"Error during inference: {str(e)}"
 
-def process_image_dot_ocr(image: Image.Image, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None) -> Dict[str, Any]:
-    """Process a single image with the dot.ocr model"""
+
+def process_image(
+    model_name: str,
+    image: Image.Image,
+    min_pixels: Optional[int] = None,
+    max_pixels: Optional[int] = None
+) -> Dict[str, Any]:
+    """Process a single image with the specified prompt mode"""
     try:
+        # Resize image if needed
         if min_pixels is not None or max_pixels is not None:
             image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
-        raw_output = dot_ocr_inference(image, prompt)
+
+        # Run inference with the default prompt
+        raw_output = inference(model_name, image, prompt)
+
+        # Process results based on prompt mode
         result = {
             'original_image': image,
             'raw_output': raw_output,
@@ -327,19 +453,45 @@ def process_image_dot_ocr(image: Image.Image, min_pixels: Optional[int] = None,
             'layout_result': None,
             'markdown_content': None
         }
+
+        # Try to parse JSON and create visualizations (since we're doing layout analysis)
         try:
-            layout_data = json.loads(raw_output)
+            # Clean the output to be valid JSON
+            # Models sometimes add ```json ... ``` markers
+            json_match = re.search(r'```json\s*([\s\S]+?)\s*```', raw_output)
+            if json_match:
+                json_str = json_match.group(1)
+            else:
+                json_str = raw_output
+            
+            layout_data = json.loads(json_str)
             result['layout_result'] = layout_data
-            processed_image = draw_layout_on_image(image, layout_data)
-            result['processed_image'] = processed_image
-            markdown_content = layoutjson2md(image, layout_data, text_key='text')
-            result['markdown_content'] = markdown_content
+
+            # Create visualization with bounding boxes
+            try:
+                processed_image = draw_layout_on_image(image, layout_data)
+                result['processed_image'] = processed_image
+            except Exception as e:
+                print(f"Error drawing layout: {e}")
+                result['processed_image'] = image
+
+            # Generate markdown from layout data
+            try:
+                markdown_content = layoutjson2md(image, layout_data, text_key='text')
+                result['markdown_content'] = markdown_content
+            except Exception as e:
+                print(f"Error generating markdown: {e}")
+                result['markdown_content'] = raw_output
+
         except json.JSONDecodeError:
             print("Failed to parse JSON output, using raw output")
             result['markdown_content'] = raw_output
+
         return result
+
     except Exception as e:
-        print(f"Error processing image with dot.ocr: {e}")
+        print(f"Error processing image: {e}")
+        traceback.print_exc()
         return {
             'original_image': image,
             'raw_output': f"Error processing image: {str(e)}",
@@ -348,279 +500,23 @@ def process_image_dot_ocr(image: Image.Image, min_pixels: Optional[int] = None,
             'markdown_content': f"Error processing image: {str(e)}"
         }
 
-def process_all_pages_dot_ocr(file_path, min_pixels, max_pixels):
-    """Process all pages of a document with dot.ocr model"""
-    if file_path.lower().endswith('.pdf'):
-        images = load_images_from_pdf(file_path)
-    else:
-        images = [Image.open(file_path).convert('RGB')]
-    results = []
-    for img in images:
-        result = process_image_dot_ocr(img, min_pixels, max_pixels)
-        results.append(result)
-    return results
-
-# Dolphin model functions
-@spaces.GPU()
-def dolphin_model_chat(prompt, image):
-    """Process an image or batch of images with the given prompt(s) using Dolphin model"""
-    is_batch = isinstance(image, list)
-    if not is_batch:
-        images = [image]
-        prompts = [prompt]
-    else:
-        images = image
-        prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
-    batch_inputs = dolphin_processor(images, return_tensors="pt", padding=True)
-    batch_pixel_values = batch_inputs.pixel_values.half().to(device)
-    prompts = [f"<s>{p} <Answer/>" for p in prompts]
-    batch_prompt_inputs = dolphin_tokenizer(
-        prompts,
-        add_special_tokens=False,
-        return_tensors="pt"
-    )
-    batch_prompt_ids = batch_prompt_inputs.input_ids.to(device)
-    batch_attention_mask = batch_prompt_inputs.attention_mask.to(device)
-    outputs = dolphin_model.generate(
-        pixel_values=batch_pixel_values,
-        decoder_input_ids=batch_prompt_ids,
-        decoder_attention_mask=batch_attention_mask,
-        min_length=1,
-        max_length=4096,
-        pad_token_id=dolphin_tokenizer.pad_token_id,
-        eos_token_id=dolphin_tokenizer.eos_token_id,
-        use_cache=True,
-        bad_words_ids=[[dolphin_tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
-        do_sample=False,
-        num_beams=1,
-        repetition_penalty=1.1
-    )
-    sequences = dolphin_tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
-    results = []
-    for i, sequence in enumerate(sequences):
-        cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
-        results.append(cleaned)
-    if not is_batch:
-        return results[0]
-    return results
-
-def process_element_batch_dolphin(elements, prompt, max_batch_size=16):
-    """Process elements of the same type in batches for Dolphin model"""
-    results = []
-    batch_size = min(len(elements), max_batch_size)
-    for i in range(0, len(elements), batch_size):
-        batch_elements = elements[i:i+batch_size]
-        crops_list = [elem["crop"] for elem in batch_elements]
-        prompts_list = [prompt] * len(crops_list)
-        batch_results = dolphin_model_chat(prompts_list, crops_list)
-        for j, result in enumerate(batch_results):
-            elem = batch_elements[j]
-            results.append({
-                "label": elem["label"],
-                "bbox": elem["bbox"],
-                "text": result.strip(),
-                "reading_order": elem["reading_order"],
-            })
-    return results
-
-def process_page_dolphin(image_path):
-    """Process a single page with Dolphin model"""
-    pil_image = Image.open(image_path).convert("RGB")
-    layout_output = dolphin_model_chat("Parse the reading order of this document.", pil_image)
-    padded_image, dims = prepare_image(pil_image)
-    recognition_results = process_elements_dolphin(layout_output, padded_image, dims)
-    return recognition_results
-
-def process_elements_dolphin(layout_results, padded_image, dims):
-    """Parse all document elements for Dolphin model"""
-    layout_results = parse_layout_string(layout_results)
-    text_elements = []
-    table_elements = []
-    figure_results = []
-    previous_box = None
-    reading_order = 0
-    for bbox, label in layout_results:
-        try:
-            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
-                bbox, padded_image, dims, previous_box
-            )
-            cropped = padded_image[y1:y2, x1:x2]
-            if cropped.size > 0 and (cropped.shape[0] > 3 and cropped.shape[1] > 3):
-                if label == "fig":
-                    try:
-                        pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                        buffered = io.BytesIO()
-                        pil_crop.save(buffered, format="PNG")
-                        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-                        figure_results.append(
-                            {
-                                "label": label,
-                                "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                                "text": img_base64,
-                                "reading_order": reading_order,
-                            }
-                        )
-                    except Exception as e:
-                        print(f"Error encoding figure to base64: {e}")
-                        figure_results.append(
-                            {
-                                "label": label,
-                                "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                                "text": "",
-                                "reading_order": reading_order,
-                            }
-                        )
-                else:
-                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                    element_info = {
-                        "crop": pil_crop,
-                        "label": label,
-                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                        "reading_order": reading_order,
-                    }
-                    if label == "tab":
-                        table_elements.append(element_info)
-                    else:
-                        text_elements.append(element_info)
-            reading_order += 1
-        except Exception as e:
-            print(f"Error processing bbox with label {label}: {str(e)}")
-            continue
-    recognition_results = figure_results.copy()
-    if text_elements:
-        text_results = process_element_batch_dolphin(text_elements, "Read text in the image.")
-        recognition_results.extend(text_results)
-    if table_elements:
-        table_results = process_element_batch_dolphin(table_elements, "Parse the table in the image.")
-        recognition_results.extend(table_results)
-    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
-    return recognition_results
-
-def generate_markdown(recognition_results):
-    """Generate markdown from recognition results for Dolphin model"""
-    converter = MarkdownConverter()
-    return converter.convert(recognition_results)
-
-def convert_all_pdf_pages_to_images(file_path, target_size=896):
-    """Convert all pages of a PDF to images for Dolphin model"""
-    if file_path is None:
-        return []
-    try:
-        file_ext = os.path.splitext(file_path)[1].lower()
-        if file_ext == '.pdf':
-            doc = pymupdf.open(file_path)
-            image_paths = []
-            for page_num in range(len(doc)):
-                page = doc[page_num]
-                rect = page.rect
-                scale = target_size / max(rect.width, rect.height)
-                mat = pymupdf.Matrix(scale, scale)
-                pix = page.get_pixmap(matrix=mat)
-                img_data = pix.tobytes("png")
-                pil_image = Image.open(io.BytesIO(img_data))
-                with tempfile.NamedTemporaryFile(suffix=f"_page_{page_num}.png", delete=False) as tmp_file:
-                    pil_image.save(tmp_file.name, "PNG")
-                    image_paths.append(tmp_file.name)
-            doc.close()
-            return image_paths
-        else:
-            converted_path = convert_to_image(file_path, target_size)
-            return [converted_path] if converted_path else []
-    except Exception as e:
-        print(f"Error converting PDF pages to images: {e}")
-        return []
-
-def convert_to_image(file_path, target_size=896, page_num=0):
-    """Convert input file to image format for Dolphin model"""
-    if file_path is None:
-        return None
-    try:
-        file_ext = os.path.splitext(file_path)[1].lower()
-        if file_ext == '.pdf':
-            doc = pymupdf.open(file_path)
-            if page_num >= len(doc):
-                page_num = 0
-            page = doc[page_num]
-            rect = page.rect
-            scale = target_size / max(rect.width, rect.height)
-            mat = pymupdf.Matrix(scale, scale)
-            pix = page.get_pixmap(matrix=mat)
-            img_data = pix.tobytes("png")
-            pil_image = Image.open(io.BytesIO(img_data))
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
-                pil_image.save(tmp_file.name, "PNG")
-                doc.close()
-                return tmp_file.name
-        else:
-            pil_image = Image.open(file_path).convert("RGB")
-            w, h = pil_image.size
-            if max(w, h) > target_size:
-                if w > h:
-                    new_w, new_h = target_size, int(h * target_size / w)
-                else:
-                    new_w, new_h = int(w * target_size / h), target_size
-                pil_image = pil_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
-                pil_image.save(tmp_file.name, "PNG")
-                return tmp_file.name
-    except Exception as e:
-        print(f"Error converting file to image: {e}")
-        return file_path
-
-def process_all_pages_dolphin(file_path):
-    """Process all pages of a document with Dolphin model"""
-    image_paths = convert_all_pdf_pages_to_images(file_path)
-    per_page_results = []
-    for image_path in image_paths:
-        try:
-            original_image = Image.open(image_path).convert('RGB')
-            recognition_results = process_page_dolphin(image_path)
-            markdown_content = generate_markdown(recognition_results)
-            placeholder_text = "Layout visualization not available for Dolphin model"
-            processed_image = create_placeholder_image(placeholder_text, size=(original_image.width, original_image.height))
-            per_page_results.append({
-                'original_image': original_image,
-                'processed_image': processed_image,
-                'markdown_content': markdown_content,
-                'layout_result': recognition_results
-            })
-        except Exception as e:
-            print(f"Error processing page: {e}")
-            per_page_results.append({
-                'original_image': Image.new('RGB', (100, 100), color='white'),
-                'processed_image': create_placeholder_image("Error processing page", size=(100, 100)),
-                'markdown_content': f"Error processing page: {str(e)}",
-                'layout_result': None
-            })
-        finally:
-            if os.path.exists(image_path):
-                os.remove(image_path)
-    return per_page_results
-
-def create_placeholder_image(text, size=(400, 200)):
-    """Create a placeholder image with text"""
-    img = Image.new('RGB', size, color='white')
-    draw = ImageDraw.Draw(img)
-    try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16)
-    except Exception:
-        font = ImageFont.load_default()
-    draw.text((10, 10), text, fill='black', font=font)
-    return img
 
-# Gradio interface functions
 def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
     """Load file for preview (supports PDF and images)"""
     global pdf_cache
+
     if not file_path or not os.path.exists(file_path):
         return None, "No file selected"
-    file_ext = os.path.splitext(file_path)[1].lower()
+
+    file_ext = os.path.splitext(file_path).lower()
+
     try:
         if file_ext == '.pdf':
+            # Load PDF pages
             images = load_images_from_pdf(file_path)
             if not images:
                 return None, "Failed to load PDF"
+
             pdf_cache.update({
                 "images": images,
                 "current_page": 0,
@@ -629,9 +525,13 @@ def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
                 "is_parsed": False,
                 "results": []
             })
-            return images[0], f"Page 1 / {len(images)}"
+
+            return images, f"Page 1 / {len(images)}"
+
         elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
+            # Load single image
             image = Image.open(file_path).convert('RGB')
+
             pdf_cache.update({
                 "images": [image],
                 "current_page": 0,
@@ -640,78 +540,73 @@ def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
                 "is_parsed": False,
                 "results": []
             })
+
             return image, "Page 1 / 1"
         else:
             return None, f"Unsupported file format: {file_ext}"
+
     except Exception as e:
         print(f"Error loading file: {e}")
         return None, f"Error loading file: {str(e)}"
 
-def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, str, Optional[Image.Image], Optional[Dict]]:
+
+def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, Any, Optional[Image.Image], Optional[Dict]]:
     """Navigate through PDF pages and update all relevant outputs."""
     global pdf_cache
+
     if not pdf_cache["images"]:
-        return None, "No file loaded", "No results yet", None, None
+        return None, '<div class="page-info">No file loaded</div>', "No results yet", None, None
+
     if direction == "prev":
         pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1)
     elif direction == "next":
-        pdf_cache["current_page"] = min(pdf_cache["total_pages"] - 1, pdf_cache["current_page"] + 1)
+        pdf_cache["current_page"] = min(
+            pdf_cache["total_pages"] - 1,
+            pdf_cache["current_page"] + 1
+        )
+
     index = pdf_cache["current_page"]
     current_image_preview = pdf_cache["images"][index]
-    page_info_html = f"Page {index + 1} / {pdf_cache['total_pages']}"
-    if pdf_cache["is_parsed"] and index < len(pdf_cache["results"]):
+    page_info_html = f'<div class="page-info">Page {index + 1} / {pdf_cache["total_pages"]}</div>'
+
+    # Initialize default result values
+    markdown_content = "Page not processed yet"
+    processed_img = None
+    layout_json = None
+
+    # Get results for current page if available
+    if (pdf_cache["is_parsed"] and
+        index < len(pdf_cache["results"]) and
+        pdf_cache["results"][index]):
+
         result = pdf_cache["results"][index]
-        processed_img = result['processed_image']
-        markdown_content = result['markdown_content'] or "No content available"
-        layout_json = result['layout_result']
+        markdown_content = result.get('markdown_content') or result.get('raw_output', 'No content available')
+        processed_img = result.get('processed_image', None) # Get the processed image
+        layout_json = result.get('layout_result', None) # Get the layout JSON
+
+    # Check for Arabic text to set RTL property
+    if is_arabic_text(markdown_content):
+        markdown_update = gr.update(value=markdown_content, rtl=True)
     else:
-        processed_img = None
-        markdown_content = "Page not processed yet"
-        layout_json = None
-    return current_image_preview, page_info_html, markdown_content, processed_img, layout_json
+        markdown_update = markdown_content
+
+    return current_image_preview, page_info_html, markdown_update, processed_img, layout_json
 
-def process_document(model_choice, file_path, max_tokens, min_pix, max_pix):
-    """Process the uploaded document with the selected model"""
-    global pdf_cache
-    try:
-        if not file_path:
-            return None, "Please upload a file first.", None
-        if model_choice == "dot.ocr":
-            results = process_all_pages_dot_ocr(file_path, min_pix, max_pix)
-        elif model_choice == "Dolphin":
-            results = process_all_pages_dolphin(file_path)
-        else:
-            raise ValueError("Invalid model choice")
-        pdf_cache["results"] = results
-        pdf_cache["is_parsed"] = True
-        first_result = results[0]
-        if model_choice == "dot.ocr":
-            processed_img = first_result['processed_image']
-            markdown_content = first_result['markdown_content']
-            layout_json = first_result['layout_result']
-        else:
-            processed_img = first_result['processed_image']
-            markdown_content = first_result['markdown_content']
-            layout_json = first_result['layout_result']
-        return processed_img, markdown_content, layout_json
-    except Exception as e:
-        error_msg = f"Error processing document: {str(e)}"
-        print(error_msg)
-        return None, error_msg, None
 
 def create_gradio_interface():
     """Create the Gradio interface"""
+
     css = """
     .main-container { max-width: 1400px; margin: 0 auto; }
     .header-text { text-align: center; color: #2c3e50; margin-bottom: 20px; }
-    .process-button { 
-        border: none !important; 
-        color: white !important; 
-        font-weight: bold !important; 
-        background-color: blue !important;} 
-    .process-button:hover { 
+    .process-button {
+        border: none !important;
+        color: white !important;
+        font-weight: bold !important;
+        background-color: blue !important;}
+    .process-button:hover {
         background-color: darkblue !important;
-        transform: translateY(-2px) !important; 
+        transform: translateY(-2px) !important;
         box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
     .info-box { border: 1px solid #dee2e6; border-radius: 8px; padding: 15px; margin: 10px 0; }
     .page-info { text-align: center; padding: 8px 16px; border-radius: 20px; font-weight: bold; margin: 10px 0; }
@@ -727,34 +622,41 @@ def create_gradio_interface():
             </p>
         </div>
         """)
+
+        # Main interface
         with gr.Row():
+            # Left column - Input and controls
             with gr.Column(scale=1):
+
+                # Model selection
                 model_choice = gr.Radio(
-                    choices=["dot.ocr", "Dolphin"],
+                    choices=["Camel-Doc-OCR-062825", "MonkeyOCR-Recognition", "Nanonets-OCR-s", "Megalodon-OCR-Sync-0713"],
                     label="Select Model",
-                    value="dot.ocr"
+                    value="Camel-Doc-OCR-062825"
                 )
+
+                # File input
                 file_input = gr.File(
                     label="Upload Image or PDF",
                     file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"],
                     type="filepath"
                 )
-                with gr.Row():
-                    examples = gr.Examples(
-                        examples=["examples/sample_image1.png", "examples/sample_image2.png", "examples/sample_pdf.pdf"],
-                        inputs=file_input,
-                        label="Example Documents"
-                    )
+
+                # Image preview
                 image_preview = gr.Image(
                     label="Preview",
                     type="pil",
                     interactive=False,
                     height=300
                 )
+
+                # Page navigation for PDFs
                 with gr.Row():
                     prev_page_btn = gr.Button("◀ Previous", size="md")
-                    page_info = gr.HTML("No file loaded")
+                    page_info = gr.HTML('<div class="page-info">No file loaded</div>')
                     next_page_btn = gr.Button("Next ▶", size="md")
+
+                # Advanced settings
                 with gr.Accordion("Advanced Settings", open=False):
                     max_new_tokens = gr.Slider(
                         minimum=1000,
@@ -764,25 +666,36 @@ def create_gradio_interface():
                         label="Max New Tokens",
                         info="Maximum number of tokens to generate"
                     )
+
                     min_pixels = gr.Number(
                         value=MIN_PIXELS,
                         label="Min Pixels",
                         info="Minimum image resolution"
                     )
+
                     max_pixels = gr.Number(
                         value=MAX_PIXELS,
-                        label="Max Pixels", 
+                        label="Max Pixels",
                         info="Maximum image resolution"
                     )
+
+                # Process button
                 process_btn = gr.Button(
                     "🚀 Process Document",
                     variant="primary",
                     elem_classes=["process-button"],
                     size="lg"
                 )
+
+                # Clear button
                 clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
+
+            # Right column - Results
             with gr.Column(scale=2):
+
+                # Results tabs
                 with gr.Tabs():
+                    # Processed image tab
                     with gr.Tab("🖼️ Processed Image"):
                         processed_image = gr.Image(
                             label="Image with Layout Detection",
@@ -790,11 +703,13 @@ def create_gradio_interface():
                             interactive=False,
                             height=500
                         )
+                    # Markdown output tab
                     with gr.Tab("📝 Extracted Content"):
                         markdown_output = gr.Markdown(
                             value="Click 'Process Document' to see extracted content...",
                             height=500
                         )
+                    # JSON layout tab
                     with gr.Tab("📋 Layout JSON"):
                         json_output = gr.JSON(
                             label="Layout Analysis Results",
@@ -802,8 +717,114 @@ def create_gradio_interface():
                         )
 
         # Event handlers
+        def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
+            """Process the uploaded document"""
+            global pdf_cache
+
+            try:
+                if not file_path:
+                    return None, "Please upload a file first.", None
+
+                # Load and preview file
+                image, page_info = load_file_for_preview(file_path)
+                if image is None:
+                    return None, page_info, None
+
+                # Process the image(s)
+                if pdf_cache["file_type"] == "pdf":
+                    # Process all pages for PDF
+                    all_results = []
+                    all_markdown = []
+
+                    for i, img in enumerate(pdf_cache["images"]):
+                        result = process_image(
+                            model_name,
+                            img,
+                            min_pixels=int(min_pix) if min_pix else None,
+                            max_pixels=int(max_pix) if max_pix else None
+                        )
+                        all_results.append(result)
+                        if result.get('markdown_content'):
+                            all_markdown.append(f"## Page {i+1}\n\n{result['markdown_content']}")
+
+                    pdf_cache["results"] = all_results
+                    pdf_cache["is_parsed"] = True
+
+                    # Show results for first page
+                    first_result = all_results
+                    combined_markdown = "\n\n---\n\n".join(all_markdown)
+
+                    # Check if the combined markdown contains mostly Arabic text
+                    if is_arabic_text(combined_markdown):
+                        markdown_update = gr.update(value=combined_markdown, rtl=True)
+                    else:
+                        markdown_update = combined_markdown
+
+                    return (
+                        first_result['processed_image'],
+                        markdown_update,
+                        first_result['layout_result']
+                    )
+                else:
+                    # Process single image
+                    result = process_image(
+                        model_name,
+                        image,
+                        min_pixels=int(min_pix) if min_pix else None,
+                        max_pixels=int(max_pix) if max_pix else None
+                    )
+
+                    pdf_cache["results"] = [result]
+                    pdf_cache["is_parsed"] = True
+
+                    # Check if the content contains mostly Arabic text
+                    content = result['markdown_content'] or "No content extracted"
+                    if is_arabic_text(content):
+                        markdown_update = gr.update(value=content, rtl=True)
+                    else:
+                        markdown_update = content
+
+                    return (
+                        result['processed_image'],
+                        markdown_update,
+                        result['layout_result']
+                    )
+
+            except Exception as e:
+                error_msg = f"Error processing document: {str(e)}"
+                print(error_msg)
+                traceback.print_exc()
+                return None, error_msg, None
+
+        def handle_file_upload(file_path):
+            """Handle file upload and show preview"""
+            if not file_path:
+                return None, "No file loaded"
+
+            image, page_info = load_file_for_preview(file_path)
+            return image, page_info
+
+        def clear_all():
+            """Clear all data and reset interface"""
+            global pdf_cache
+
+            pdf_cache = {
+                "images": [], "current_page": 0, "total_pages": 0,
+                "file_type": None, "is_parsed": False, "results": []
+            }
+
+            return (
+                None,  # file_input
+                None,  # image_preview
+                '<div class="page-info">No file loaded</div>',  # page_info
+                None,  # processed_image
+                "Click 'Process Document' to see extracted content...",  # markdown_output
+                None,  # json_output
+            )
+
+        # Wire up event handlers
         file_input.change(
-            lambda file_path: load_file_for_preview(file_path),
+            handle_file_upload,
             inputs=[file_input],
             outputs=[image_preview, page_info]
         )
@@ -825,12 +846,23 @@ def create_gradio_interface():
         )
 
         clear_btn.click(
-            lambda: (None, None, "No file loaded", None, "Click 'Process Document' to see extracted content...", None),
-            outputs=[file_input, image_preview, page_info, processed_image, markdown_output, json_output]
+            clear_all,
+            outputs=[
+                file_input, image_preview, page_info, processed_image,
+                markdown_output, json_output
+            ]
         )
 
     return demo
 
+
 if __name__ == "__main__":
+    # Create and launch the interface
     demo = create_gradio_interface()
-    demo.queue(max_size=10).launch(share=False, debug=True, show_error=True)
+    demo.queue(max_size=10).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True,
+        show_error=True
+    )