Update app.py

app.py

@@ -18,387 +18,208 @@ from transformers import (
     AutoProcessor,
     TextIteratorStreamer,
 )
-from qwen_vl_utils import process_vision_info
 
-# Constants
-MIN_PIXELS = 3136
-MAX_PIXELS = 11289600
-IMAGE_FACTOR = 28
-MAX_INPUT_TOKEN_LENGTH = 2048
+# --- Constants and Model Setup ---
+MAX_INPUT_TOKEN_LENGTH = 4096
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 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.
+# --- Prompts for Different Tasks ---
+layout_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: Format its text as HTML.
-    - All Others (Text, Title, etc.): Format their text as Markdown.
-
+    - For tables, provide the content in a structured JSON format.
+    - For all other elements, provide the plain text.
 4. Constraints:
-    - The output text must be the original text from the image, with no translation.
+    - The output must be the original text from the image.
     - All layout elements must be sorted according to human reading order.
-
-5. Final Output: The entire output must be a single JSON object.
+5. Final Output: The entire output must be a single JSON object wrapped in ```json ... ```.
 """
 
-# Load models
+ocr_prompt = "Perform precise OCR on the image. Extract all text content, maintaining the original structure, paragraphs, and tables as formatted markdown."
+
+# --- Model Loading ---
 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
-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(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
-    import base64
-    from io import BytesIO
+# --- Utility Functions ---
+def layoutjson2md(layout_data: List[Dict]) -> str:
+    """Converts the structured JSON from Layout Analysis into formatted Markdown."""
     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', [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_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
-            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")
+            text = item.get('text', '')
+            if not text: continue
+
+            if category == 'Title': markdown_lines.append(f"# {text}\n")
+            elif category == 'Section-header': markdown_lines.append(f"## {text}\n")
             elif category == 'Table':
-                if text.strip().startswith('<'):
+                 # Handle structured table JSON
+                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.extend([header, separator] + rows)
+                    markdown_lines.append("\n")
+                else: # Fallback for simple text
                     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("")
+                 markdown_lines.append(f"{text}\n")
     except Exception as e:
         print(f"Error converting to markdown: {e}")
-        return str(layout_data)
+        return "### Error converting JSON to Markdown."
     return "\n".join(markdown_lines)
 
+# --- Core Application Logic ---
 @spaces.GPU
-def inference(model_name: str, image: Image.Image, text: str, max_new_tokens: int = 1024) -> str:
-    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}")
-
-        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"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
+def process_document_stream(model_name: str, task_choice: str, image: Image.Image, max_new_tokens: int):
+    """
+    Main generator function that handles both OCR and Layout Analysis tasks.
+    """
+    if image is None:
+        yield "Please upload an image.", "Please upload an image.", None
+        return
+
+    # 1. Select prompt based on user's task choice
+    text_prompt = ocr_prompt if task_choice == "Content Extraction" else layout_prompt
+
+    # 2. Select 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
+
+    # 3. Prepare model inputs and streamer
+    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}
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # 4. Stream raw output to the UI in real-time
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "")
+        time.sleep(0.01)
+        yield buffer, "⏳ Processing...", {"status": "streaming"}
+
+    # 5. Post-process the final buffer based on the selected task
+    if task_choice == "Content Extraction":
+        # For OCR, the buffer is the final result.
+        yield buffer, buffer, None
+    else: # Layout Analysis
         try:
-            json_match = re.search(r'```json
-            json_str = json_match.group(1) if json_match else buffer
+            json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
+            if not json_match:
+                raise json.JSONDecodeError("JSON object not found in output.", buffer, 0)
+            
+            json_str = json_match.group(1)
             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
-
-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)}"
-
+            markdown_content = layoutjson2md(layout_data)
+            
+            yield buffer, markdown_content, layout_data
+        except Exception as e:
+            error_md = f"❌ **Error:** Failed to parse Layout JSON.\n\n**Details:**\n`{str(e)}`"
+            error_json = {"error": "ProcessingError", "details": str(e), "raw_output": buffer}
+            yield buffer, error_md, error_json
+
+# --- Gradio UI Definition ---
 def create_gradio_interface():
+    """Builds and returns the Gradio web 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 { 
-        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; }
+    .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; }
     """
     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 to markdown document processing
+                Advanced Vision-Language Model for Image Content and Layout Extraction
             </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"
+                    label="Select Model", value="Camel-Doc-OCR-062825"
                 )
-                file_input = gr.File(
-                    label="Upload Image",
-                    file_types  =[".jpg", ".jpeg", ".png", ".bmp", ".tiff"],
-                    type="filepath"
+                task_choice = gr.Radio(
+                    choices=["Content Extraction", "Layout Analysis"],
+                    label="Select Task", value="Content Extraction"
                 )
-                image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
+                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
                 with gr.Accordion("Advanced Settings", open=False):
-                    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")
+                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens")
+                
                 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"):
-                        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="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])
+                with gr.Tabs() as tabs:
+                    with gr.Tab("📝 Extracted Content", id=0):
+                        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")
+                    
+                    with gr.Tab("📋 Layout Analysis Results", id=1):
+                        json_output = gr.JSON(label="Structured Layout Data (JSON)")
+        
+        # Event Handlers
+        def clear_all_outputs():
+            return None, "Raw output will appear here.", "Formatted results will appear here.", None
+
         process_btn.click(
-            process_document,
-            inputs=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
-            outputs=[output, json_output]
+            fn=process_document_stream,
+            inputs=[model_choice, task_choice, image_input, max_new_tokens],
+            outputs=[raw_output_stream, markdown_output, json_output]
         )
         clear_btn.click(
-            clear_all,
-            outputs=[file_input, image_preview, output, json_output]
+            clear_all_outputs,
+            outputs=[image_input, raw_output_stream, markdown_output, json_output]
         )
     return demo
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    demo.queue(max_size=10).launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True, show_error=True)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)