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=""" ... ... """, 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(""" ...your full XML... """) except ImportError: print("Pillow not installed, can't create dummy image.") except Exception as e: print(f"Error creating dummy image: {e}") demo.launch()