Upload 2 files

app.py

@@ -0,0 +1,527 @@
+import gradio as gr
+from PIL import Image
+import xml.etree.ElementTree as ET
+import os
+import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText, pipeline
+import spaces
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"  # turn on HF_TRANSFER
+# --- Global Model and Processor ---
+MODELS = {}
+PROCESSORS = {}
+PIPELINES = {}
+MODEL_LOAD_ERROR_MSG = {}
+
+# Available models
+AVAILABLE_MODELS = ["RolmOCR", "Nanonets-OCR-s"]
+
+# Load RolmOCR
+try:
+    PROCESSORS["RolmOCR"] = AutoProcessor.from_pretrained("reducto/RolmOCR")
+    MODELS["RolmOCR"] = AutoModelForImageTextToText.from_pretrained(
+        "reducto/RolmOCR", torch_dtype=torch.bfloat16, device_map="auto"
+    )
+    PIPELINES["RolmOCR"] = pipeline(
+        "image-text-to-text", model=MODELS["RolmOCR"], processor=PROCESSORS["RolmOCR"]
+    )
+except Exception as e:
+    MODEL_LOAD_ERROR_MSG["RolmOCR"] = f"Failed to load RolmOCR: {str(e)}"
+    print(f"Error loading RolmOCR: {e}")
+
+# Load Nanonets-OCR-s
+try:
+    PROCESSORS["Nanonets-OCR-s"] = AutoProcessor.from_pretrained(
+        "nanonets/Nanonets-OCR-s"
+    )
+    MODELS["Nanonets-OCR-s"] = AutoModelForImageTextToText.from_pretrained(
+        "nanonets/Nanonets-OCR-s", torch_dtype=torch.bfloat16, device_map="auto"
+    )
+    PIPELINES["Nanonets-OCR-s"] = pipeline(
+        "image-text-to-text",
+        model=MODELS["Nanonets-OCR-s"],
+        processor=PROCESSORS["Nanonets-OCR-s"],
+    )
+except Exception as e:
+    MODEL_LOAD_ERROR_MSG["Nanonets-OCR-s"] = f"Failed to load Nanonets-OCR-s: {str(e)}"
+    print(f"Error loading Nanonets-OCR-s: {e}")
+
+
+# --- Helper Functions ---
+
+
+def get_xml_namespace(xml_file_path):
+    """
+    Dynamically gets the namespace from the XML file.
+    Returns both the namespace and the format type (ALTO or PAGE).
+    """
+    try:
+        tree = ET.parse(xml_file_path)
+        root = tree.getroot()
+        if "}" in root.tag:
+            ns = root.tag.split("}")[0] + "}"
+            # Determine format based on root element
+            if "PcGts" in root.tag:
+                return ns, "PAGE"
+            elif "alto" in root.tag.lower():
+                return ns, "ALTO"
+    except ET.ParseError:
+        print(f"Error parsing XML to find namespace: {xml_file_path}")
+    return "", "UNKNOWN"
+
+
+def parse_page_xml_for_text(xml_file_path):
+    """
+    Parses a PAGE XML file to extract text content.
+    Returns:
+        - full_text (str): All extracted text concatenated.
+    """
+    full_text_lines = []
+
+    if not xml_file_path or not os.path.exists(xml_file_path):
+        return "Error: XML file not provided or does not exist."
+
+    try:
+        ns_prefix, _ = get_xml_namespace(xml_file_path)
+        tree = ET.parse(xml_file_path)
+        root = tree.getroot()
+
+        # Find all TextLine elements
+        for text_line in root.findall(f".//{ns_prefix}TextLine"):
+            # First try to get text from TextEquiv/Unicode
+            text_equiv = text_line.find(f"{ns_prefix}TextEquiv/{ns_prefix}Unicode")
+            if text_equiv is not None and text_equiv.text:
+                full_text_lines.append(text_equiv.text)
+                continue
+
+            # If no TextEquiv, try to get text from Word elements
+            line_text_parts = []
+            for word in text_line.findall(f"{ns_prefix}Word"):
+                word_text = word.find(f"{ns_prefix}TextEquiv/{ns_prefix}Unicode")
+                if word_text is not None and word_text.text:
+                    line_text_parts.append(word_text.text)
+
+            if line_text_parts:
+                full_text_lines.append(" ".join(line_text_parts))
+
+        return "\n".join(full_text_lines)
+
+    except ET.ParseError as e:
+        return f"Error parsing XML: {e}"
+    except Exception as e:
+        return f"An unexpected error occurred during XML parsing: {e}"
+
+
+def parse_alto_xml_for_text(xml_file_path):
+    """
+    Parses an ALTO XML file to extract text content.
+    Returns:
+        - full_text (str): All extracted text concatenated.
+    """
+    full_text_lines = []
+
+    if not xml_file_path or not os.path.exists(xml_file_path):
+        return "Error: XML file not provided or does not exist."
+
+    try:
+        ns_prefix, _ = get_xml_namespace(xml_file_path)
+        tree = ET.parse(xml_file_path)
+        root = tree.getroot()
+
+        for text_line in root.findall(f".//{ns_prefix}TextLine"):
+            line_text_parts = []
+            for string_element in text_line.findall(f"{ns_prefix}String"):
+                if text := string_element.get("CONTENT"):
+                    line_text_parts.append(text)
+            if line_text_parts:
+                full_text_lines.append(" ".join(line_text_parts))
+
+        return "\n".join(full_text_lines)
+
+    except ET.ParseError as e:
+        return f"Error parsing XML: {e}"
+    except Exception as e:
+        return f"An unexpected error occurred during XML parsing: {e}"
+
+
+def parse_xml_for_text(xml_file_path):
+    """
+    Main function to parse XML files, automatically detecting the format.
+    """
+    if not xml_file_path or not os.path.exists(xml_file_path):
+        return "Error: XML file not provided or does not exist."
+
+    try:
+        _, xml_format = get_xml_namespace(xml_file_path)
+
+        if xml_format == "PAGE":
+            return parse_page_xml_for_text(xml_file_path)
+        elif xml_format == "ALTO":
+            return parse_alto_xml_for_text(xml_file_path)
+        else:
+            return "Error: Unsupported XML format. Expected ALTO or PAGE XML."
+
+    except Exception as e:
+        return f"Error determining XML format: {str(e)}"
+
+
+@spaces.GPU
+def predict(pil_image, model_name="RolmOCR"):
+    """Performs OCR prediction using the selected Hugging Face model."""
+    global PIPELINES, MODEL_LOAD_ERROR_MSG
+
+    if model_name not in PIPELINES:
+        error_to_report = MODEL_LOAD_ERROR_MSG.get(
+            model_name,
+            f"Model {model_name} could not be initialized or is not available.",
+        )
+        raise RuntimeError(error_to_report)
+
+    selected_pipe = PIPELINES[model_name]
+
+    # Format the message based on the model
+    if model_name == "RolmOCR":
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": pil_image},
+                    {
+                        "type": "text",
+                        "text": "Return the plain text representation of this document as if you were reading it naturally.\n",
+                    },
+                ],
+            }
+        ]
+    else:  # Nanonets-OCR-s
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": pil_image},
+                    {
+                        "type": "text",
+                        "text": "Extract and return all the text from this image. Include all text elements and maintain the reading order. If there are tables, convert them to markdown format. If there are mathematical equations, convert them to LaTeX format.",
+                    },
+                ],
+            }
+        ]
+    max_tokens = 8096
+    # Use the pipeline with the properly formatted messages
+    return selected_pipe(messages, max_new_tokens=max_tokens)
+
+
+def run_hf_ocr(image_path, model_name="RolmOCR"):
+    """
+    Runs OCR on the provided image using the selected Hugging Face model (via predict function).
+    """
+    if image_path is None:
+        return "No image provided for OCR."
+
+    try:
+        pil_image = Image.open(image_path).convert("RGB")
+        ocr_results = predict(
+            pil_image, model_name
+        )  # predict handles model loading and inference
+
+        # Parse the output based on the user's example structure
+        if (
+            isinstance(ocr_results, list)
+            and ocr_results
+            and "generated_text" in ocr_results[0]
+        ):
+            generated_content = ocr_results[0]["generated_text"]
+
+            if isinstance(generated_content, str):
+                return generated_content
+
+            if isinstance(generated_content, list) and generated_content:
+                if assistant_message := next(
+                    (
+                        msg["content"]
+                        for msg in reversed(generated_content)
+                        if isinstance(msg, dict)
+                        and msg.get("role") == "assistant"
+                        and "content" in msg
+                    ),
+                    None,
+                ):
+                    return assistant_message
+
+                # Fallback if the specific assistant message structure isn't found but there's content
+                if (
+                    isinstance(generated_content[0], dict)
+                    and "content" in generated_content[0]
+                ):
+                    if (
+                        len(generated_content) > 1
+                        and isinstance(generated_content[1], dict)
+                        and "content" in generated_content[1]
+                    ):
+                        return generated_content[1][
+                            "content"
+                        ]  # Assuming second part is assistant
+                    else:
+                        return generated_content[0]["content"]
+
+            print(f"Unexpected OCR output structure from HF model: {ocr_results}")
+            return "Error: Could not parse OCR model output. Check console."
+
+        else:
+            print(f"Unexpected OCR output structure from HF model: {ocr_results}")
+            return "Error: OCR model did not return expected output. Check console."
+
+    except RuntimeError as e:  # Catch model loading/initialization errors from predict
+        return str(e)
+    except Exception as e:
+        print(f"Error during Hugging Face OCR processing: {e}")
+        return f"Error during Hugging Face OCR: {str(e)}"
+
+
+# --- Gradio Interface Function ---
+
+
+def process_files(image_path, xml_path, model_name):
+    """
+    Main function for the Gradio interface.
+    Processes the image for display, runs OCR with selected model,
+    and parses XML if provided.
+    """
+    img_to_display = None
+    xml_text_output = "XML not provided or not processed."
+    hf_ocr_text_output = "Image not provided or OCR not run."
+    ocr_download = gr.DownloadButton(visible=False)
+    xml_download = gr.DownloadButton(visible=False)
+
+    if image_path:
+        try:
+            img_to_display = Image.open(image_path).convert("RGB")
+            hf_ocr_text_output = run_hf_ocr(image_path, model_name)
+
+            # Create download file for OCR output
+            if hf_ocr_text_output and not hf_ocr_text_output.startswith("Error"):
+                ocr_filename = f"vlm_ocr_output_{model_name}.txt"
+                with open(ocr_filename, "w", encoding="utf-8") as f:
+                    f.write(hf_ocr_text_output)
+                ocr_download = gr.DownloadButton(
+                    label="Download VLM OCR", value=ocr_filename, visible=True
+                )
+        except Exception as e:
+            img_to_display = None  # Clear image if it failed to load
+            hf_ocr_text_output = f"Error loading image or running {model_name} OCR: {e}"
+    else:
+        hf_ocr_text_output = "Please upload an image to perform OCR."
+
+    if xml_path:
+        xml_text_output = parse_xml_for_text(xml_path)
+
+        # Create download file for XML text
+        if xml_text_output and not xml_text_output.startswith("Error"):
+            xml_filename = "traditional_ocr_output.txt"
+            with open(xml_filename, "w", encoding="utf-8") as f:
+                f.write(xml_text_output)
+            xml_download = gr.DownloadButton(
+                label="Download XML Text", value=xml_filename, visible=True
+            )
+    else:
+        xml_text_output = "No XML file uploaded."
+
+    # If only XML is provided without an image
+    if not image_path and xml_path:
+        img_to_display = None  # No image to display
+        hf_ocr_text_output = "Upload an image to perform OCR."
+
+    return (
+        img_to_display,
+        xml_text_output,
+        hf_ocr_text_output,
+        ocr_download,
+        xml_download,
+    )
+
+
+# --- Create Gradio App ---
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🕰️ OCR Time Machine")
+    gr.Markdown(
+        "Travel through time to see how OCR technology has evolved! \n\n "
+        "For decades, galleries, libraries, archives, and museums (GLAMs) have used Optical Character Recognition "
+        "to transform digitized books, newspapers, and manuscripts into machine-readable text. Traditional OCR "
+        "produces complex XML formats like ALTO, packed with layout details but difficult to use. "
+        "Now, cutting-edge Vision-Language Models (VLMs) are revolutionizing OCR with simpler, cleaner Markdown output. "
+        "This Space makes it easy to compare these two approaches and see which works best for your historical documents. "
+        "Upload a historical document image and its XML file to compare these approaches side-by-side. "
+        "We'll extract the reading order from your XML for an apples-to-apples comparison of the actual text content.\n\n"
+        "**Available models:** [RolmOCR](https://huggingface.co/reducto/RolmOCR) | "
+        "[Nanonets-OCR-s](https://huggingface.co/nanonets/Nanonets-OCR-s)"
+    )
+
+    gr.Markdown("---")
+
+    # How it works section
+    gr.Markdown("## 🚀 How it works")
+    gr.Markdown(
+        "1. 📤 **Upload Image**: Select a historical document image (JPG, PNG, JP2)\n"
+        "2. 📄 **Upload XML** (Optional): Add the corresponding ALTO or PAGE XML file for comparison\n"
+        "3. 🤖 **Choose Model**: Select between RolmOCR (new) or Nanonets-OCR-s (even newer!)\n"
+        "4. 🔍 **Compare**: Click 'Compare OCR Methods' to process\n"
+        "5. 💾 **Download**: Save the results for further analysis"
+    )
+
+    gr.Markdown("---")
+
+    # Input section
+    gr.Markdown("## 📥 Upload Files")
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("### 📤 Step 1: Upload your document")
+                image_input = gr.File(
+                    label="Historical Document Image",
+                    type="filepath",
+                    file_types=["image"],
+                )
+                xml_input = gr.File(
+                    label="XML File (Optional - ALTO or PAGE format)",
+                    type="filepath",
+                    file_types=[".xml"],
+                )
+
+            with gr.Group():
+                gr.Markdown("### 🤖 Step 2: Select OCR Model")
+                model_selector = gr.Radio(
+                    choices=AVAILABLE_MODELS,
+                    value="RolmOCR",
+                    label="Choose Model",
+                    info="RolmOCR: Fast & general-purpose | Nanonets: Advanced with table/math support",
+                )
+
+            submit_button = gr.Button(
+                "🔍 Compare OCR Methods", variant="primary", size="lg"
+            )
+
+    # Results section
+    gr.Markdown("## 📊 Results")
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("### 🖼️ Document Image")
+                output_image_display = gr.Image(
+                    label="Uploaded Document", type="pil", interactive=False
+                )
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("### 🤖 Modern VLM OCR Output")
+                hf_ocr_output_textbox = gr.Markdown(
+                    label="Markdown Format",
+                    show_copy_button=True,
+                )
+                ocr_download_btn = gr.DownloadButton(
+                    label="💾 Download VLM OCR", visible=False, size="sm"
+                )
+            with gr.Group():
+                gr.Markdown("### 📜 Traditional OCR Output")
+                xml_output_textbox = gr.Textbox(
+                    label="XML Reading Order",
+                    lines=10,
+                    interactive=False,
+                    show_copy_button=True,
+                )
+                xml_download_btn = gr.DownloadButton(
+                    label="💾 Download XML Text", visible=False, size="sm"
+                )
+
+    submit_button.click(
+        fn=process_files,
+        inputs=[image_input, xml_input, model_selector],
+        outputs=[
+            output_image_display,
+            xml_output_textbox,
+            hf_ocr_output_textbox,
+            ocr_download_btn,
+            xml_download_btn,
+        ],
+    )
+
+    gr.Markdown("---")
+
+    # Examples section
+    with gr.Group():
+        gr.Markdown("## 🎯 Try an Example")
+        gr.Examples(
+            examples=[
+                [
+                    "examples/one/74442232.3.jpg",
+                    "examples/one/74442232.34.xml",
+                    "RolmOCR",
+                ],
+                [
+                    "examples/one/74442232.3.jpg",
+                    "examples/one/74442232.34.xml",
+                    "Nanonets-OCR-s",
+                ],
+            ],
+            inputs=[image_input, xml_input, model_selector],
+            outputs=[
+                output_image_display,
+                xml_output_textbox,
+                hf_ocr_output_textbox,
+                ocr_download_btn,
+                xml_download_btn,
+            ],
+            fn=process_files,
+            cache_examples=False,
+        )
+        gr.Markdown(
+            "*Example from ['A Medical History of British India'](https://data.nls.uk/data/digitised-collections/a-medical-history-of-british-india/) "
+            "collection, National Library of Scotland*"
+        )
+
+    gr.Markdown("---")
+
+    # Tips section
+    with gr.Accordion("💡 Tips & Information", open=False):
+        gr.Markdown(
+            "### 📚 About ALTO/PAGE XML\n"
+            "- **ALTO** (Analyzed Layout and Text Object) and **PAGE** are XML formats that store OCR results with detailed layout information\n"
+            "- These files are typically generated by traditional OCR software and include position data for each text element\n"
+            "- This tool extracts just the reading order text for easier comparison\n\n"
+            "### 🎯 Best Practices\n"
+            "- Use high-resolution scans (300+ DPI) for best results\n"
+            "- Historical documents with clear text work best\n"
+            "- The VLM models can handle complex layouts, tables, and mathematical notation\n\n"
+        )
+        gr.Code(
+            value=(
+                """<alto xmlns="http://www.loc.gov/standards/alto/v3/alto.xsd">
+  <Layout>
+    <Page>
+      <PrintSpace>
+        <TextLine>
+          <String CONTENT="Hello World"/>
+        </TextLine>
+      </PrintSpace>
+    </Page>
+  </Layout>
+</alto>"""
+            ),
+            interactive=False,
+        )
+
+    # Footer
+    gr.Markdown("---")
+    gr.Markdown(
+        "<center>\n\n"
+        "Built with ❤️ for the GLAM community | "
+        "[Learn more about OCR formats](https://www.loc.gov/standards/alto/) | "
+        "Questions? [Open an issue](https://github.com/davanstrien/ocr-playground/issues)\n\n"
+        "</center>"
+    )
+
+if __name__ == "__main__":
+    print("Attempting to launch Gradio demo...")
+    print(
+        "If the Hugging Face model is large, initial startup might take some time due to model download/loading (on first OCR attempt)."
+    )
+    demo.launch()

requirements.txt

@@ -0,0 +1,11 @@
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+gradio 
+Pillow
+lxml
+transformers
+spaces
+torchvision
+accelerate
+hf-transfer
+hf-xet