app.py

import gradio as gr
from PIL import Image # ImageDraw, ImageFont are no longer needed for overlay
import xml.etree.ElementTree as ET
import os

# --- Helper Functions ---

def get_alto_namespace(xml_file_path):
    """
    Dynamically gets the ALTO namespace from the XML file.
    """
    try:
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        if '}' in root.tag:
            return root.tag.split('}')[0] + '}'
    except ET.ParseError:
        print(f"Error parsing XML to find namespace: {xml_file_path}")
    return ''

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_alto_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'):
            line_text_parts = []
            for string_element in text_line.findall(f'{ns_prefix}String'):
                text = string_element.get('CONTENT')
                if text: # Ensure text is not None
                    line_text_parts.append(text)
            # Also consider <SP/> (Space) elements if they contribute to word separation
            # and are not implicitly handled by joining CONTENT attributes.
            # For now, just joining CONTENT attributes.
            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}"

# The draw_ocr_on_image function is no longer needed.

# --- Gradio Interface Function ---

def process_image_and_xml(image_path, xml_path):
    """
    Main function for the Gradio interface.
    Processes the image and XML to return the image and extracted text.
    """
    if image_path is None: # If no image is uploaded at all
        return None, "Please upload an image."
    
    try:
        img_pil = Image.open(image_path).convert("RGB")
    except Exception as e:
        return None, f"Error loading image: {e}"

    if xml_path is None: # If XML is missing, but image is present
        return img_pil, "Please upload an OCR XML file."

    # Both image and XML are presumably present
    extracted_text = parse_alto_xml_for_text(xml_path)
    
    return img_pil, extracted_text


# --- Create Gradio App ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# OCR Viewer (ALTO XML) - Text Extractor")
    gr.Markdown(
        "Upload an image and its corresponding ALTO OCR XML file. "
        "The app will display the image and extract/show the plain text."
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.File(label="Upload Image (PNG, JPG, etc.)", type="filepath")
            xml_input = gr.File(label="Upload ALTO XML File (.xml)", type="filepath")
            # show_overlay_checkbox has been removed
            submit_button = gr.Button("Process Files", variant="primary")

    with gr.Row():
        with gr.Column(scale=1):
            output_image_orig = gr.Image(label="Uploaded Image", type="pil", interactive=False)
        with gr.Column(scale=1):
            output_text = gr.Textbox(label="Extracted Plain Text", lines=15, interactive=False)
    
    # output_image_overlay has been removed

    def update_interface(image_filepath, xml_filepath):
        # image_filepath and xml_filepath are now strings (paths) or None
        
        if image_filepath is None and xml_filepath is None:
            return None, "Please upload an image and an XML file."
        # process_image_and_xml handles cases where one is None
        
        img, text = process_image_and_xml(image_filepath, xml_filepath)
        
        return img, text

    submit_button.click(
        fn=update_interface,
        inputs=[image_input, xml_input], # show_overlay_checkbox removed
        outputs=[output_image_orig, output_text] # output_image_overlay removed
    )
    
    # The .change event for show_overlay_checkbox has been removed

    gr.Markdown("---")
    gr.Markdown("### Example ALTO XML Snippet (for `String` element extraction):")
    gr.Code(
        value="""
<alto xmlns="http://www.loc.gov/standards/alto/v3/alto.xsd">
  <Description>...</Description>
  <Styles>...</Styles>
  <Layout>
    <Page ID="Page13" PHYSICAL_IMG_NR="13" WIDTH="2394" HEIGHT="3612">
      <PrintSpace>
        <TextLine WIDTH="684" HEIGHT="108" ID="p13_t1" HPOS="465" VPOS="196">
          <String ID="p13_w1" CONTENT="Introduction" HPOS="465" VPOS="196" WIDTH="684" HEIGHT="108" STYLEREFS="font0"/>
        </TextLine>
        <TextLine WIDTH="1798" HEIGHT="51" ID="p13_t2" HPOS="492" VPOS="523">
          <String ID="p13_w2" CONTENT="Britain" HPOS="492" VPOS="523" WIDTH="166" HEIGHT="51" STYLEREFS="font1"/>
          <SP WIDTH="24" VPOS="523" HPOS="658"/>
          <String ID="p13_w3" CONTENT="1981" HPOS="682" VPOS="523" WIDTH="117" HEIGHT="51" STYLEREFS="font1"/>
          <!-- ... more String and SP elements ... -->
        </TextLine>
        <!-- ... more TextLine elements ... -->
      </PrintSpace>
    </Page>
  </Layout>
</alto>
        """,
        interactive=False
    )


if __name__ == "__main__":
    try:
        # Create a dummy image for testing
        img_test = Image.new('RGB', (2394, 3612), color = 'lightgray') # Dimensions from example XML
        img_test.save("dummy_image.png")
        print("Created dummy_image.png for testing.")
        
        # Ensure the example XML file (189819724.34.xml) exists in the same directory
        # or provide the correct path if it's elsewhere.
        example_xml_filename = "189819724.34.xml" 
        if not os.path.exists(example_xml_filename):
            print(f"WARNING: Example XML '{example_xml_filename}' not found. Please create it (using the content from the prompt) or upload your own.")

    except ImportError:
        print("Pillow not installed, can't create dummy image.")
    except Exception as e:
        print(f"Error during setup: {e}")

    demo.launch()