Spaces:
Running
on
Zero
Running
on
Zero
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): | |
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): | |
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 styled 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) | |
for item in ocr_data: | |
x, y, w, h = item['x'], item['y'], item['w'], item['h'] | |
text = item['text'] | |
# 1. Draw bounding box for the original OCR segment | |
draw.rectangle([(x, y), (x + w, y + h)], outline="rgba(255,0,0,190)", width=1) # Red, bit transparent | |
# 2. Determine font for this specific item | |
# Adaptive font size based on item height, with min/max caps | |
# Making it slightly smaller relative to box height for overlay text | |
current_font_size = max(8, min(16, int(item['h'] * 0.60))) | |
is_truetype = False | |
try: | |
font = ImageFont.truetype("arial.ttf", current_font_size) | |
is_truetype = True | |
except IOError: | |
try: | |
font = ImageFont.truetype("DejaVuSans.ttf", current_font_size) | |
is_truetype = True | |
except IOError: | |
font = ImageFont.load_default() # Small bitmap font, size not controllable | |
# print("Arial and DejaVuSans fonts not found, using default PIL font for some items.") | |
# 3. Get actual text dimensions for precise placement | |
text_actual_width = 0 | |
text_actual_height = 0 | |
text_draw_origin_is_top_left = False # Flag to know how to use coordinates with draw.text | |
try: # Modern Pillow: use textbbox with anchor for precise bounds | |
# Using 'lt' (left-top) anchor means (0,0) in textbbox is relative to top-left of text | |
bbox = font.getbbox(text, anchor='lt') | |
text_actual_width = bbox[2] - bbox[0] | |
text_actual_height = bbox[3] - bbox[1] | |
text_draw_origin_is_top_left = True # draw.text with anchor='lt' will use xy as top-left | |
except (AttributeError, TypeError): # Fallback for older Pillow or if anchor not supported by getbbox | |
try: # Try font.getsize (deprecated but widely available) | |
size = font.getsize(text) | |
text_actual_width = size[0] | |
if is_truetype: # For TrueType, getsize height is usually baseline to top | |
ascent, descent = font.getmetrics() | |
text_actual_height = ascent + descent # Full line height | |
else: # For default font, getsize height is total height | |
text_actual_height = size[1] | |
except AttributeError: # Ultimate fallback if getsize also fails (unlikely for loaded font) | |
text_actual_width = len(text) * current_font_size // 2 # Very rough estimate | |
text_actual_height = current_font_size # Very rough estimate | |
# 4. Calculate position for the TOP-LEFT of the overlay text | |
buffer = 3 # Buffer in pixels between OCR box and text overlay | |
text_draw_x = x # Align left edge of text with left edge of box | |
# Try to place text above the box | |
tentative_text_top_y = y - text_actual_height - buffer | |
if tentative_text_top_y < buffer: # If it goes off (or too close to) the top of the image | |
tentative_text_top_y = y + h + buffer # Place it below the box | |
# Ensure it doesn't go off the bottom either | |
if tentative_text_top_y + text_actual_height > image_pil.height - buffer: | |
tentative_text_top_y = image_pil.height - text_actual_height - buffer # Pin to bottom | |
if tentative_text_top_y < buffer: # If image is too short for text, pin to top | |
tentative_text_top_y = buffer | |
final_text_top_y = tentative_text_top_y | |
# 5. Draw background for the overlay text | |
bg_padding = 2 | |
bg_x0 = text_draw_x - bg_padding | |
bg_y0 = final_text_top_y - bg_padding | |
bg_x1 = text_draw_x + text_actual_width + bg_padding | |
bg_y1 = final_text_top_y + text_actual_height + bg_padding | |
draw.rectangle([(bg_x0, bg_y0), (bg_x1, bg_y1)], fill="rgba(255, 255, 220, 235)") # Light yellow, fairly opaque | |
# 6. Draw the overlay text | |
draw_coords = (text_draw_x, final_text_top_y) | |
if text_draw_origin_is_top_left: | |
try: | |
draw.text(draw_coords, text, fill="black", font=font, anchor='lt') | |
except (TypeError, AttributeError): # Fallback if anchor='lt' fails at draw time | |
# This fallback means background might be slightly off if is_truetype and text_draw_origin_is_top_left was True due to getbbox working | |
# but draw.text doesn't support anchor. For simplicity, draw at top_y assuming it's baseline. | |
ascent = font.getmetrics()[0] if is_truetype else text_actual_height | |
draw.text((text_draw_x, final_text_top_y + ascent if is_truetype else final_text_top_y), text, fill="black", font=font) | |
else: # Older Pillow, (x,y) is baseline for TrueType, top-left for default | |
if is_truetype: | |
ascent, _ = font.getmetrics() | |
draw.text((text_draw_x, final_text_top_y + ascent), text, fill="black", font=font) | |
else: # Default font, (x,y) is top-left | |
draw.text((text_draw_x, final_text_top_y), text, fill="black", font=font) | |
return image_pil | |
# --- Gradio Interface Function --- | |
def process_image_and_xml(image_path, xml_path, show_overlay): | |
if image_path is None: | |
return None, "Please upload an image.", None | |
try: | |
img_pil = Image.open(image_path).convert("RGB") | |
except Exception as e: | |
return None, f"Error loading image: {e}", None | |
if xml_path is None: | |
return img_pil, "Please upload an OCR XML file.", None | |
extracted_text, ocr_box_data = parse_alto_xml(xml_path) | |
overlay_image_pil = None | |
if show_overlay: | |
if ocr_box_data: | |
img_for_overlay = img_pil.copy() | |
overlay_image_pil = draw_ocr_on_image(img_for_overlay, ocr_box_data) | |
elif not (isinstance(extracted_text, str) and extracted_text.startswith("Error")): | |
if isinstance(extracted_text, str): | |
extracted_text += "\n(No bounding box data for overlay)" | |
else: | |
extracted_text = "(No bounding box data 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): | |
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): | |
if image_filepath is None and xml_filepath is None: | |
return None, "Please upload an image and an XML file.", None | |
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( | |
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> | |
<!-- ... more TextLine and String elements ... --> | |
</PrintSpace> | |
</Page> | |
</Layout> | |
</alto> | |
""", | |
interactive=False | |
) | |
if __name__ == "__main__": | |
try: | |
img_test = Image.new('RGB', (2394, 3612), color = 'lightgray') | |
img_test.save("dummy_image.png") | |
print("Created dummy_image.png for testing.") | |
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() |