app.py

import gradio as gr
from PIL import Image, ImageDraw, ImageFont
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(xml_file_path):
    """
    Parses an ALTO XML file to extract text content and bounding box info.
    Returns:
        - full_text (str): All extracted text concatenated.
        - ocr_data (list): A list of dictionaries, each with
                           {'text': str, 'x': int, 'y': int, 'w': int, 'h': int}
    """
    full_text_lines = []
    ocr_data = []
    
    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()

        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:
                    line_text_parts.append(text)
                    try:
                        hpos = int(float(string_element.get('HPOS')))
                        vpos = int(float(string_element.get('VPOS')))
                        width = int(float(string_element.get('WIDTH')))
                        height = int(float(string_element.get('HEIGHT')))
                        ocr_data.append({
                            'text': text,
                            'x': hpos,
                            'y': vpos,
                            'w': width,
                            'h': height
                        })
                    except (ValueError, TypeError) as e:
                        print(f"Warning: Could not parse coordinates for '{text}': {e}")
                        ocr_data.append({
                            'text': text, 'x': 0, 'y': 0, 'w': 10, 'h': 10
                        })
            if line_text_parts:
                full_text_lines.append(" ".join(line_text_parts))
        
        return "\n".join(full_text_lines), ocr_data

    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 draw_ocr_on_image(image_pil, ocr_data):
    """
    Draws bounding boxes and text from ocr_data onto the image.
    """
    if not image_pil or not ocr_data:
        return image_pil

    draw = ImageDraw.Draw(image_pil)
    
    try:
        avg_height = sum(d['h'] for d in ocr_data if d['h'] > 0) / len(ocr_data) if ocr_data and any(d['h'] > 0 for d in ocr_data) else 10
        font_size = max(8, int(avg_height * 0.6)) 
        font = ImageFont.truetype("arial.ttf", font_size)
    except (IOError, ZeroDivisionError): # Added ZeroDivisionError for safety
        font = ImageFont.load_default()
        font_size = 10 
        print("Arial font not found or issue with height calculation, using default font.")

    for item in ocr_data:
        x, y, w, h = item['x'], item['y'], item['w'], item['h']
        text = item['text']
        draw.rectangle([(x, y), (x + w, y + h)], outline="red", width=2)
        text_position = (x + 2, y - font_size - 2 if y - font_size - 2 > 0 else y + 2)
        draw.text(text_position, text, fill="green", font=font)
        
    return image_pil

# --- Gradio Interface Function ---

def process_image_and_xml(image_path, xml_path, show_overlay):
    """
    Main function for the Gradio interface.
    image_path and xml_path are now file paths (strings).
    """
    if image_path is None:
        return None, "Please upload an image.", None
    if xml_path is None:
        # If image_path is not None, we can still show the image
        try:
            img_pil_orig = Image.open(image_path).convert("RGB")
        except Exception as e:
            return None, f"Error loading image: {e}. Also, please upload an OCR XML file.", None
        return img_pil_orig, "Please upload an OCR XML file.", None

    try:
        img_pil = Image.open(image_path).convert("RGB")
    except Exception as e:
        return None, f"Error loading image: {e}", None

    extracted_text, ocr_box_data = parse_alto_xml(xml_path)

    overlay_image_pil = None
    if show_overlay and ocr_box_data:
        img_for_overlay = img_pil.copy()
        overlay_image_pil = draw_ocr_on_image(img_for_overlay, ocr_box_data)
    elif show_overlay and not ocr_box_data and not (isinstance(extracted_text, str) and extracted_text.startswith("Error")):
        extracted_text += "\n(No bounding box data found or parsed for overlay)"

    return img_pil, extracted_text, overlay_image_pil


# --- Create Gradio App ---

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

    with gr.Row():
        with gr.Column(scale=1):
            # Corrected: type="filepath"
            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 = gr.Checkbox(label="Show OCR Overlay on Image", value=False)
            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 = gr.Image(label="Image with OCR Overlay", type="pil", interactive=False, visible=True)

    def update_interface(image_filepath, xml_filepath, show_overlay_val):
        # image_filepath and xml_filepath are now strings (paths) or None
        
        # Initial check for None inputs to provide guidance
        if image_filepath is None and xml_filepath is None:
            return None, "Please upload an image and an XML file.", None
        if image_filepath is None:
            # xml_filepath might be present, but we need the image first
            return None, "Please upload an image file.", None
        # xml_filepath can be None here, process_image_and_xml handles it

        # Call the main processing function
        img, text, overlay_img = process_image_and_xml(image_filepath, xml_filepath, show_overlay_val)
        
        return img, text, overlay_img

    submit_button.click(
        fn=update_interface,
        inputs=[image_input, xml_input, show_overlay_checkbox],
        outputs=[output_image_orig, output_text, output_image_overlay]
    )
    
    show_overlay_checkbox.change(
        fn=update_interface,
        inputs=[image_input, xml_input, show_overlay_checkbox],
        outputs=[output_image_orig, output_text, output_image_overlay]
    )

    gr.Markdown("---")
    gr.Markdown("### Example ALTO XML Snippet (for `String` element extraction):")
    gr.Code(
        language="xml",
        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:
        # from PIL import Image as PImage # Already imported as Image
        img = Image.new('RGB', (2394, 3612), color = 'lightgray')
        img.save("dummy_image.png")
        print("Created dummy_image.png for testing.")
        
        # Make sure your example XML is named 189819724.34.xml and is in the same directory
        # Or, create it if it doesn't exist with the content you provided
        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 or upload your own.")
            # You could also write the example XML content here if needed for testing
            # with open(example_xml_filename, "w") as f:
            #     f.write("""<?xml version="1.0" encoding="UTF-8"?> ...your full XML... </alto>""")


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

    demo.launch()