Update app.py

app.py

@@ -18,227 +18,387 @@ from transformers import (
     AutoProcessor,
     TextIteratorStreamer,
 )
+from qwen_vl_utils import process_vision_info
 
-# --- Constants and Model Setup ---
-MAX_INPUT_TOKEN_LENGTH = 4096
+# Constants
+MIN_PIXELS = 3136
+MAX_PIXELS = 11289600
+IMAGE_FACTOR = 28
+MAX_INPUT_TOKEN_LENGTH = 2048
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# The detailed prompt to instruct the model to generate structured JSON
+# Prompts
 prompt = """Please output the layout information from the image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
 
 1. Bbox format: [x1, y1, x2, y2]
+
 2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
+
 3. Text Extraction & Formatting Rules:
     - Picture: For the 'Picture' category, the text field should be omitted.
     - Formula: Format its text as LaTeX.
-    - Table: For tables, provide the content in a structured format within the JSON.
+    - Table: Format its text as HTML.
     - All Others (Text, Title, etc.): Format their text as Markdown.
+
 4. Constraints:
     - The output text must be the original text from the image, with no translation.
     - All layout elements must be sorted according to human reading order.
-5. Final Output: The entire output must be a single JSON object wrapped in ```json ... ```.
+
+5. Final Output: The entire output must be a single JSON object.
 """
 
 # Load models
 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
+    MODEL_ID_M,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
 ).to(device).eval()
 
 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
+    MODEL_ID_T,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
 ).to(device).eval()
 
 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
+    MODEL_ID_C,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
 ).to(device).eval()
 
 MODEL_ID_G = "echo840/MonkeyOCR"
 SUBFOLDER = "Recognition"
 processor_g = AutoProcessor.from_pretrained(
-    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER
+    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
+    MODEL_ID_G,
+    trust_remote_code=True,
+    subfolder=SUBFOLDER,
+    torch_dtype=torch.float16
 ).to(device).eval()
 
-# --- Utility Functions ---
+# Utility functions
+def round_by_factor(number: int, factor: int) -> int:
+    return round(number / factor) * factor
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = 28,
+    min_pixels: int = 3136,
+    max_pixels: int = 11289600,
+):
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(f"Aspect ratio too extreme: {max(height, width) / min(height, width)}")
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = round_by_factor(height / beta, factor)
+        w_bar = round_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = round_by_factor(height * beta, factor)
+        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):
+    if isinstance(image_input, str):
+        if image_input.startswith(("http://", "https://")):
+            response = requests.get(image_input)
+            image = Image.open(BytesIO(response.content)).convert('RGB')
+        else:
+            image = Image.open(image_input).convert('RGB')
+    elif isinstance(image_input, Image.Image):
+        image = image_input.convert('RGB')
+    else:
+        raise ValueError(f"Invalid image input type: {type.image_input)}")
+    if min_pixels or max_pixels:
+        min_pixels = min_pixels or MIN_PIXELS
+        max_pixels = max_pixels or MAX_PIXELS
+        height, width = smart_resize(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 is_arabic_text(text: str) -> bool:
+    if not text:
+        return False
+    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
+        header_match = re.match(header_pattern, line, re.MULTILINE)
+        if header_match:
+            content_text.append(header_match.group(1))
+            continue
+        if re.match(paragraph_pattern, line, re.MULTILINE):
+            content_text.append(line)
+    if not content_text:
+        return False
+    combined_text = ' '.join(content_text)
+    arabic_chars = 0
+    total_chars = 0
+    for char in combined_text:
+        if char.isalpha():
+            total_chars += 1
+            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
+                arabic_chars += 1
+    return total_chars > 0 and (arabic_chars / total_chars) > 0.5
 
-def layoutjson2md(layout_data: List[Dict]) -> str:
-    """Converts the structured JSON layout data into formatted Markdown."""
+def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
+    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]))
         for item in sorted_items:
             category = item.get('category', '')
-            text = item.get('text', '')
-
-            if not text:
+            text = item.get(text_key, '')
+            bbox = item.get('bbox', [])
+            if category == 'Picture':
+                if bbox and len(bbox) == 4:
+                    try:
+                        x1, y1, x2, y2 = bbox
+                        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))
+                            buffer = BytesIO()
+                            cropped_img.save(buffer, format='PNG')
+                            img_data = base64.b64encode(buffer.getvalue()).decode()
+                            markdown_lines.append(f"<image-card alt="Image" src="data:image/png;base64,{img_data}" ></image-card>\n")
+                        else:
+                            markdown_lines.append("<image-card alt="Image" src="Image region detected" ></image-card>\n")
+                    except Exception as e:
+                        print(f"Error processing image region: {e}")
+                        markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\n")
+                else:
+                    markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\n")
+            elif not text:
                 continue
-            
-            if category == 'Title':
+            elif category == 'Title':
                 markdown_lines.append(f"# {text}\n")
             elif category == 'Section-header':
                 markdown_lines.append(f"## {text}\n")
+            elif category == 'Text':
+                markdown_lines.append(f"{text}\n")
+            elif category == 'List-item':
+                markdown_lines.append(f"- {text}\n")
             elif category == 'Table':
-                # Check if the text is a dictionary representing a structured table
-                if isinstance(text, dict) and 'header' in text and 'rows' in text:
-                    header = '| ' + ' | '.join(map(str, text['header'])) + ' |'
-                    separator = '| ' + ' | '.join(['---'] * len(text['header'])) + ' |'
-                    rows = ['| ' + ' | '.join(map(str, row)) + ' |' for row in text['rows']]
-                    markdown_lines.append(header)
-                    markdown_lines.append(separator)
-                    markdown_lines.extend(rows)
-                    markdown_lines.append("\n")
-                else:
-                    # Fallback for unstructured table text
+                if text.strip().startswith('<'):
                     markdown_lines.append(f"{text}\n")
+                else:
+                    markdown_lines.append(f"**Table:** {text}\n")
+            elif category == 'Formula':
+                if text.strip().startswith('$') or '\\' in text:
+                    markdown_lines.append(f"$$  \n{text}\n  $$\n")
+                else:
+                    markdown_lines.append(f"**Formula:** {text}\n")
+            elif category == 'Caption':
+                markdown_lines.append(f"*{text}*\n")
+            elif category == 'Footnote':
+                markdown_lines.append(f"^{text}^\n")
+            elif category in ['Page-header', 'Page-footer']:
+                continue
             else:
                 markdown_lines.append(f"{text}\n")
-                
+            markdown_lines.append("")
     except Exception as e:
         print(f"Error converting to markdown: {e}")
-        return "### Error converting JSON to Markdown."
+        return str(layout_data)
     return "\n".join(markdown_lines)
 
-# --- Core Application Logic ---
-
 @spaces.GPU
-def process_document_stream(model_name: str, image: Image.Image, text_prompt: str, max_new_tokens: int):
-    """
-    Main generator function that streams raw model output and then processes it into
-    formatted Markdown and structured JSON for the UI.
-    """
-    if image is None:
-        yield "Please upload an image.", "Please upload an image.", None
-        return
-
-    # Select the model and processor
-    if model_name == "Camel-Doc-OCR-062825": processor, model = processor_m, model_m
-    elif model_name == "Megalodon-OCR-Sync-0713": processor, model = processor_t, model_t
-    elif model_name == "Nanonets-OCR-s": processor, model = processor_c, model_c
-    elif model_name == "MonkeyOCR-Recognition": processor, model = processor_g, model_g
-    else:
-        yield "Invalid model selected.", "Invalid model selected.", None
-        return
-
-    # Prepare model inputs
-    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text_prompt}]}]
-    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)
-    
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-    
-    # Start generation in a separate thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    # Stream raw output to the UI
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        buffer = buffer.replace("<|im_end|>", "")
-        time.sleep(0.01)
-        # Yield the raw stream and placeholders for the final results
-        yield buffer, "⏳ Formatting Markdown...", {"status": "processing"}
-
-    # After streaming is complete, process the final buffer
+def inference(model_name: str, image: Image.Image, text: str, max_new_tokens: int = 1024) -> str:
     try:
-        # Extract the JSON object from the buffer
-        json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
-        if not json_match:
-            raise json.JSONDecodeError("JSON object not found in the model's output.", buffer, 0)
-        
-        json_str = json_match.group(1)
-        layout_data = json.loads(json_str)
-        
-        # Convert the parsed JSON to formatted markdown
-        markdown_content = layoutjson2md(layout_data)
-        
-        # Yield the final, complete results
-        yield buffer, markdown_content, layout_data
-
-    except json.JSONDecodeError as e:
-        print(f"JSON parsing failed: {e}")
-        error_md = f"❌ **Error:** Failed to parse JSON from the model's output.\n\nSee the raw output stream for details."
-        error_json = {"error": "JSONDecodeError", "details": str(e), "raw_output": buffer}
-        yield buffer, error_md, error_json
+        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}")
+
+        if image is None:
+            yield "Please upload an image.", "Please upload an image."
+            return
+
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": text},
+            ]
+        }]
+        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(
+            text=[prompt_full],
+            images=[image],
+            return_tensors="pt",
+            padding=True,
+            truncation=False,
+            max_length=MAX_INPUT_TOKEN_LENGTH
+        ).to(device)
+        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer, buffer
     except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        yield buffer, f"❌ An unexpected error occurred: {e}", None
+        print(f"Error during inference: {e}")
+        traceback.print_exc()
+        yield f"Error during inference: {str(e)}", f"Error during inference: {str(e)}"
 
+def process_image(
+    model_name: str,
+    image: Image.Image,
+    min_pixels: Optional[int] = None,
+    max_pixels: Optional[int] = None,
+    max_new_tokens: int = 1024
+):
+    try:
+        if min_pixels or max_pixels:
+            image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
+        buffer = ""
+        for raw_output, _ in inference(model_name, image, prompt, max_new_tokens):
+            buffer = raw_output
+            yield buffer, None  # Yield raw OCR stream and None for JSON during processing
+        try:
+            json_match = re.search(r'```json
+            json_str = json_match.group(1) if json_match else buffer
+            layout_data = json.loads(json_str)
+            yield buffer, layout_data  # Final yield with raw OCR and parsed JSON
+        except json.JSONDecodeError:
+            print("Failed to parse JSON output, using raw output")
+            yield buffer, None  # If JSON parsing fails, yield raw OCR with no JSON
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        traceback.print_exc()
+        yield f"Error processing image: {str(e)}", None
 
-# --- Gradio UI Definition ---
+def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
+    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()
+    try:
+        if file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
+            image = Image.open(file_path).convert('RGB')
+            return image, "Image loaded"
+        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 create_gradio_interface():
-    """Builds and returns the Gradio web interface."""
     css = """
     .main-container { max-width: 1400px; margin: 0 auto; }
-    .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; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
+    .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 { 
+        background-color: darkblue !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; }
+    .model-status { padding: 10px; border-radius: 8px; margin: 10px 0; text-align: center; font-weight: bold; }
+    .status-ready { background: #d1edff; color: #0c5460; border: 1px solid #b8daff; }
     """
     with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         gr.HTML("""
         <div class="title" style="text-align: center">
             <h1>Dot<span style="color: red;">●</span><strong></strong>OCR Comparator</h1>
             <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
-                Advanced Vision-Language Model for Image Layout Analysis
+                Advanced vision-language model for image to markdown document processing
             </p>
         </div>
         """)
-        
         with gr.Row():
-            # --- Left Column (Inputs) ---
             with gr.Column(scale=1):
                 model_choice = gr.Radio(
                     choices=["Camel-Doc-OCR-062825", "MonkeyOCR-Recognition", "Nanonets-OCR-s", "Megalodon-OCR-Sync-0713"],
                     label="Select Model",
                     value="Camel-Doc-OCR-062825"
                 )
-                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
+                file_input = gr.File(
+                    label="Upload Image",
+                    file_types  =[".jpg", ".jpeg", ".png", ".bmp", ".tiff"],
+                    type="filepath"
+                )
+                image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
                 with gr.Accordion("Advanced Settings", open=False):
-                    max_new_tokens = gr.Slider(minimum=1000, maximum=8192, value=4096, step=256, label="Max New Tokens")
-                
+                    max_new_tokens = gr.Slider(minimum=1000, maximum=32000, value=24000, step=1000, label="Max New Tokens")
+                    min_pixels = gr.Number(value=MIN_PIXELS, label="Min Pixels")
+                    max_pixels = gr.Number(value=MAX_PIXELS, label="Max Pixels")
                 process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
                 clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-
-            # --- Right Column (Outputs) ---
             with gr.Column(scale=2):
                 with gr.Tabs():
                     with gr.Tab("📝 Extracted Content"):
-                        raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=15, show_copy_button=True)
-                        with gr.Accordion("(Formatted Result)", open=True):
-                            markdown_output = gr.Markdown(label="Formatted Markdown (from JSON)")
-                    
+                        output = gr.Textbox(label="Raw OCR Stream", interactive=False, lines=10, show_copy_button=True)
                     with gr.Tab("📋 Layout Analysis Results"):
-                        json_output = gr.JSON(label="Structured Layout Data (JSON)", value=None)
-        
-        # --- Event Handlers ---
-        def clear_all_outputs():
-            """Resets all input and output fields to their default state."""
-            return None, "Raw output will appear here.", "Formatted results will appear here.", None
-
-        # Connect the process button to the main generator function
+                        json_output = gr.JSON(label="Layout Analysis Results", value=None)
+        def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
+            try:
+                if not file_path:
+                    return "Please upload an image.", None
+                image, status = load_file_for_preview(file_path)
+                if image is None:
+                    return status, None
+                for raw_output, layout_result in 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, max_new_tokens=max_tokens):
+                    yield raw_output, layout_result
+            except Exception as e:
+                error_msg = f"Error processing document: {str(e)}"
+                print(error_msg)
+                traceback.print_exc()
+                yield error_msg, None
+        def handle_file_upload(file_path):
+            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():
+            return None, None, "No file loaded", None
+        file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview, output])
         process_btn.click(
-            fn=process_document_stream,
-            inputs=[model_choice, image_input, gr.Textbox(value=prompt, visible=False), max_new_tokens],
-            outputs=[raw_output_stream, markdown_output, json_output]
+            process_document,
+            inputs=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
+            outputs=[output, json_output]
         )
-
-        # Connect the clear button
         clear_btn.click(
-            clear_all_outputs,
-            outputs=[image_input, raw_output_stream, markdown_output, json_output]
+            clear_all,
+            outputs=[file_input, image_preview, output, json_output]
         )
-
     return demo
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, 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)