update app [main]

app.py

@@ -0,0 +1,213 @@
+import gradio as gr
+import torch
+import os
+import sys
+from PIL import Image, ImageDraw
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, GenerationConfig
+from huggingface_hub import snapshot_download
+import spaces
+from typing import Optional, Tuple, Dict, Any, Iterable
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+
+print("Downloading model snapshot to ensure all scripts are present...")
+model_dir = snapshot_download(repo_id="nvidia/NVIDIA-Nemotron-Parse-v1.1")
+print(f"Model downloaded to: {model_dir}")
+
+sys.path.append(model_dir)
+
+try:
+    from postprocessing import extract_classes_bboxes, transform_bbox_to_original, postprocess_text
+    print("Successfully imported postprocessing functions.")
+except ImportError as e:
+    print(f" Error importing postprocessing: {e}")
+    raise e
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+colors.steel_blue = colors.Color(
+    name="steel_blue",
+    c50="#EBF3F8",
+    c100="#D3E5F0",
+    c200="#A8CCE1",
+    c300="#7DB3D2",
+    c400="#529AC3",
+    c500="#4682B4",
+    c600="#3E72A0",
+    c700="#36638C",
+    c800="#2E5378",
+    c900="#264364",
+    c950="#1E3450",
+)
+
+class SteelBlueTheme(Soft):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.gray,
+        secondary_hue: colors.Color | str = colors.steel_blue,
+        neutral_hue: colors.Color | str = colors.slate,
+        text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary="*primary_50",
+            background_fill_primary_dark="*primary_900",
+            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
+            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
+            button_primary_text_color="white",
+            button_primary_text_color_hover="white",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            slider_color="*secondary_500",
+            slider_color_dark="*secondary_600",
+            block_title_text_weight="600",
+            block_border_width="3px",
+            block_shadow="*shadow_drop_lg",
+            button_primary_shadow="*shadow_drop_lg",
+            button_large_padding="11px",
+            color_accent_soft="*primary_100",
+            block_label_background_fill="*primary_200",
+        )
+
+steel_blue_theme = SteelBlueTheme()
+css = """
+#main-title h1 { font-size: 2.3em !important; }
+#output-title h2 { font-size: 2.1em !important; }
+"""
+
+print("Loading Model components...")
+
+processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_dir,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+).to(device).eval()
+
+try:
+    generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)
+except Exception as e:
+    print(f"Warning: Could not load GenerationConfig: {e}. Using default.")
+    generation_config = GenerationConfig(max_new_tokens=4096)
+
+print("Model loaded successfully.")
+
+@spaces.GPU
+def process_ocr_task(image):
+    """
+    Processes an image with NVIDIA-Nemotron-Parse-v1.1.
+    """
+    if image is None:
+        return "Please upload an image first.", None
+        
+    task_prompt = "</s><s><predict_bbox><predict_classes><output_markdown>"
+    
+    inputs = processor(images=[image], text=task_prompt, return_tensors="pt").to(device)
+    
+    if device.type == 'cuda':
+        inputs = {k: v.to(torch.bfloat16) if v.dtype == torch.float32 else v for k, v in inputs.items()}
+
+    print("Running inference...")
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs, 
+            generation_config=generation_config
+        )
+
+    generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    
+    try:
+        classes, bboxes, texts = extract_classes_bboxes(generated_text)
+    except Exception as e:
+        print(f"Error extracting boxes: {e}")
+        return generated_text, image
+
+    bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
+
+    table_format = 'latex'  
+    text_format = 'markdown'
+    blank_text_in_figures = False 
+
+    processed_texts = [
+        postprocess_text(
+            text, 
+            cls=cls, 
+            table_format=table_format, 
+            text_format=text_format, 
+            blank_text_in_figures=blank_text_in_figures
+        ) 
+        for text, cls in zip(texts, classes)
+    ]
+
+    result_image = image.copy()
+    draw = ImageDraw.Draw(result_image)
+    
+    color_map = {
+        "Table": "red",
+        "Figure": "blue",
+        "Text": "green",
+        "Title": "purple"
+    }
+
+    final_output_text = ""
+    
+    for cls, bbox, txt in zip(classes, bboxes, processed_texts):
+        color = color_map.get(cls, "red")
+        draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], outline=color, width=3)
+        
+        if cls == "Table":
+            final_output_text += f"\n\n--- [Table] ---\n{txt}\n-----------------\n"
+        elif cls == "Figure":
+            final_output_text += f"\n\n--- [Figure] ---\n(Figure Detected)\n-----------------\n"
+        else:
+            final_output_text += f"{txt}\n"
+
+    if not final_output_text.strip() and generated_text:
+        final_output_text = generated_text
+
+    return final_output_text, result_image
+
+with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
+    gr.Markdown("# **NVIDIA Nemotron Parse v1.1 [OCR/Parsing]**", elem_id="main-title")
+    gr.Markdown("Upload a document image to extract text, tables, and layout structures using NVIDIA's state-of-the-art Parse model.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard"])
+            submit_btn = gr.Button("Process Document", variant="primary")
+
+            examples = gr.Examples(
+                examples=["examples/1.jpg"],
+                inputs=image_input, 
+                label="Examples"
+            )
+
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(label="Parsed Content (Markdown/LaTeX)", lines=8, show_copy_button=True)
+            output_image = gr.Image(label="Detected Layout & Bounding Boxes", type="pil")
+            
+    submit_btn.click(
+        fn=process_ocr_task, 
+        inputs=[image_input], 
+        outputs=[output_text, output_image]
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(share=True, mcp_server=True, ssr_mode=False)

requirements.txt

@@ -0,0 +1,19 @@
+opencv_python_headless==4.9.0.80
+opencv_python==4.8.0.74
+transformers==4.51.3
+huggingface_hub
+open-clip-torch
+beautifulsoup4
+albumentations
+sentencepiece 
+numpy==1.26.4
+torchmetrics
+torchvision
+mdtex2html
+html2text
+spaces
+einops
+gradio
+pillow
+torch
+timm