Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -1,11 +1,14 @@
|
|
1 |
import gradio as gr
|
2 |
-
from PIL import Image
|
3 |
import xml.etree.ElementTree as ET
|
4 |
import os
|
5 |
|
6 |
# --- Helper Functions ---
|
7 |
|
8 |
def get_alto_namespace(xml_file_path):
|
|
|
|
|
|
|
9 |
try:
|
10 |
tree = ET.parse(xml_file_path)
|
11 |
root = tree.getroot()
|
@@ -15,206 +18,111 @@ def get_alto_namespace(xml_file_path):
|
|
15 |
print(f"Error parsing XML to find namespace: {xml_file_path}")
|
16 |
return ''
|
17 |
|
18 |
-
def
|
|
|
|
|
|
|
|
|
|
|
19 |
full_text_lines = []
|
20 |
-
|
21 |
if not xml_file_path or not os.path.exists(xml_file_path):
|
22 |
-
return "Error: XML file not provided or does not exist."
|
|
|
23 |
try:
|
24 |
ns_prefix = get_alto_namespace(xml_file_path)
|
25 |
tree = ET.parse(xml_file_path)
|
26 |
root = tree.getroot()
|
|
|
|
|
27 |
for text_line in root.findall(f'.//{ns_prefix}TextLine'):
|
28 |
line_text_parts = []
|
29 |
for string_element in text_line.findall(f'{ns_prefix}String'):
|
30 |
text = string_element.get('CONTENT')
|
31 |
-
if text:
|
32 |
line_text_parts.append(text)
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
width = int(float(string_element.get('WIDTH')))
|
37 |
-
height = int(float(string_element.get('HEIGHT')))
|
38 |
-
ocr_data.append({
|
39 |
-
'text': text, 'x': hpos, 'y': vpos, 'w': width, 'h': height
|
40 |
-
})
|
41 |
-
except (ValueError, TypeError) as e:
|
42 |
-
print(f"Warning: Could not parse coordinates for '{text}': {e}")
|
43 |
-
ocr_data.append({
|
44 |
-
'text': text, 'x': 0, 'y': 0, 'w': 10, 'h': 10
|
45 |
-
})
|
46 |
if line_text_parts:
|
47 |
full_text_lines.append(" ".join(line_text_parts))
|
48 |
-
|
|
|
|
|
49 |
except ET.ParseError as e:
|
50 |
-
return f"Error parsing XML: {e}"
|
51 |
except Exception as e:
|
52 |
-
return f"An unexpected error occurred during XML parsing: {e}"
|
53 |
-
|
54 |
-
def draw_ocr_on_image(image_pil, ocr_data):
|
55 |
-
"""
|
56 |
-
Draws styled bounding boxes and text from ocr_data onto the image.
|
57 |
-
"""
|
58 |
-
if not image_pil or not ocr_data:
|
59 |
-
return image_pil
|
60 |
-
|
61 |
-
draw = ImageDraw.Draw(image_pil)
|
62 |
-
|
63 |
-
for item in ocr_data:
|
64 |
-
x, y, w, h = item['x'], item['y'], item['w'], item['h']
|
65 |
-
text = item['text']
|
66 |
-
|
67 |
-
# 1. Draw bounding box for the original OCR segment
|
68 |
-
draw.rectangle([(x, y), (x + w, y + h)], outline="rgba(255,0,0,190)", width=1) # Red, bit transparent
|
69 |
-
|
70 |
-
# 2. Determine font for this specific item
|
71 |
-
# Adaptive font size based on item height, with min/max caps
|
72 |
-
# Making it slightly smaller relative to box height for overlay text
|
73 |
-
current_font_size = max(8, min(16, int(item['h'] * 0.60)))
|
74 |
-
is_truetype = False
|
75 |
-
try:
|
76 |
-
font = ImageFont.truetype("arial.ttf", current_font_size)
|
77 |
-
is_truetype = True
|
78 |
-
except IOError:
|
79 |
-
try:
|
80 |
-
font = ImageFont.truetype("DejaVuSans.ttf", current_font_size)
|
81 |
-
is_truetype = True
|
82 |
-
except IOError:
|
83 |
-
font = ImageFont.load_default() # Small bitmap font, size not controllable
|
84 |
-
# print("Arial and DejaVuSans fonts not found, using default PIL font for some items.")
|
85 |
-
|
86 |
-
# 3. Get actual text dimensions for precise placement
|
87 |
-
text_actual_width = 0
|
88 |
-
text_actual_height = 0
|
89 |
-
text_draw_origin_is_top_left = False # Flag to know how to use coordinates with draw.text
|
90 |
-
|
91 |
-
try: # Modern Pillow: use textbbox with anchor for precise bounds
|
92 |
-
# Using 'lt' (left-top) anchor means (0,0) in textbbox is relative to top-left of text
|
93 |
-
bbox = font.getbbox(text, anchor='lt')
|
94 |
-
text_actual_width = bbox[2] - bbox[0]
|
95 |
-
text_actual_height = bbox[3] - bbox[1]
|
96 |
-
text_draw_origin_is_top_left = True # draw.text with anchor='lt' will use xy as top-left
|
97 |
-
except (AttributeError, TypeError): # Fallback for older Pillow or if anchor not supported by getbbox
|
98 |
-
try: # Try font.getsize (deprecated but widely available)
|
99 |
-
size = font.getsize(text)
|
100 |
-
text_actual_width = size[0]
|
101 |
-
if is_truetype: # For TrueType, getsize height is usually baseline to top
|
102 |
-
ascent, descent = font.getmetrics()
|
103 |
-
text_actual_height = ascent + descent # Full line height
|
104 |
-
else: # For default font, getsize height is total height
|
105 |
-
text_actual_height = size[1]
|
106 |
-
except AttributeError: # Ultimate fallback if getsize also fails (unlikely for loaded font)
|
107 |
-
text_actual_width = len(text) * current_font_size // 2 # Very rough estimate
|
108 |
-
text_actual_height = current_font_size # Very rough estimate
|
109 |
-
|
110 |
-
# 4. Calculate position for the TOP-LEFT of the overlay text
|
111 |
-
buffer = 3 # Buffer in pixels between OCR box and text overlay
|
112 |
-
text_draw_x = x # Align left edge of text with left edge of box
|
113 |
-
|
114 |
-
# Try to place text above the box
|
115 |
-
tentative_text_top_y = y - text_actual_height - buffer
|
116 |
-
if tentative_text_top_y < buffer: # If it goes off (or too close to) the top of the image
|
117 |
-
tentative_text_top_y = y + h + buffer # Place it below the box
|
118 |
-
|
119 |
-
# Ensure it doesn't go off the bottom either
|
120 |
-
if tentative_text_top_y + text_actual_height > image_pil.height - buffer:
|
121 |
-
tentative_text_top_y = image_pil.height - text_actual_height - buffer # Pin to bottom
|
122 |
-
if tentative_text_top_y < buffer: # If image is too short for text, pin to top
|
123 |
-
tentative_text_top_y = buffer
|
124 |
|
125 |
-
|
126 |
-
|
127 |
-
# 5. Draw background for the overlay text
|
128 |
-
bg_padding = 2
|
129 |
-
bg_x0 = text_draw_x - bg_padding
|
130 |
-
bg_y0 = final_text_top_y - bg_padding
|
131 |
-
bg_x1 = text_draw_x + text_actual_width + bg_padding
|
132 |
-
bg_y1 = final_text_top_y + text_actual_height + bg_padding
|
133 |
-
|
134 |
-
draw.rectangle([(bg_x0, bg_y0), (bg_x1, bg_y1)], fill="rgba(255, 255, 220, 235)") # Light yellow, fairly opaque
|
135 |
-
|
136 |
-
# 6. Draw the overlay text
|
137 |
-
draw_coords = (text_draw_x, final_text_top_y)
|
138 |
-
if text_draw_origin_is_top_left:
|
139 |
-
try:
|
140 |
-
draw.text(draw_coords, text, fill="black", font=font, anchor='lt')
|
141 |
-
except (TypeError, AttributeError): # Fallback if anchor='lt' fails at draw time
|
142 |
-
# This fallback means background might be slightly off if is_truetype and text_draw_origin_is_top_left was True due to getbbox working
|
143 |
-
# but draw.text doesn't support anchor. For simplicity, draw at top_y assuming it's baseline.
|
144 |
-
ascent = font.getmetrics()[0] if is_truetype else text_actual_height
|
145 |
-
draw.text((text_draw_x, final_text_top_y + ascent if is_truetype else final_text_top_y), text, fill="black", font=font)
|
146 |
-
else: # Older Pillow, (x,y) is baseline for TrueType, top-left for default
|
147 |
-
if is_truetype:
|
148 |
-
ascent, _ = font.getmetrics()
|
149 |
-
draw.text((text_draw_x, final_text_top_y + ascent), text, fill="black", font=font)
|
150 |
-
else: # Default font, (x,y) is top-left
|
151 |
-
draw.text((text_draw_x, final_text_top_y), text, fill="black", font=font)
|
152 |
-
|
153 |
-
return image_pil
|
154 |
|
155 |
# --- Gradio Interface Function ---
|
156 |
|
157 |
-
def process_image_and_xml(image_path, xml_path
|
158 |
-
|
159 |
-
|
|
|
|
|
|
|
|
|
|
|
160 |
try:
|
161 |
img_pil = Image.open(image_path).convert("RGB")
|
162 |
except Exception as e:
|
163 |
-
return None, f"Error loading image: {e}"
|
164 |
-
|
165 |
-
if xml_path is None:
|
166 |
-
return img_pil, "Please upload an OCR XML file."
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
overlay_image_pil = draw_ocr_on_image(img_for_overlay, ocr_box_data)
|
174 |
-
elif not (isinstance(extracted_text, str) and extracted_text.startswith("Error")):
|
175 |
-
if isinstance(extracted_text, str):
|
176 |
-
extracted_text += "\n(No bounding box data for overlay)"
|
177 |
-
else:
|
178 |
-
extracted_text = "(No bounding box data for overlay)"
|
179 |
-
return img_pil, extracted_text, overlay_image_pil
|
180 |
|
181 |
# --- Create Gradio App ---
|
|
|
182 |
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
183 |
-
gr.Markdown("# OCR Viewer (ALTO XML)")
|
184 |
gr.Markdown(
|
185 |
"Upload an image and its corresponding ALTO OCR XML file. "
|
186 |
-
"The app will display the image
|
187 |
-
"and optionally overlay the OCR predictions on the image."
|
188 |
)
|
|
|
189 |
with gr.Row():
|
190 |
with gr.Column(scale=1):
|
191 |
image_input = gr.File(label="Upload Image (PNG, JPG, etc.)", type="filepath")
|
192 |
xml_input = gr.File(label="Upload ALTO XML File (.xml)", type="filepath")
|
193 |
-
show_overlay_checkbox
|
194 |
submit_button = gr.Button("Process Files", variant="primary")
|
|
|
195 |
with gr.Row():
|
196 |
with gr.Column(scale=1):
|
197 |
output_image_orig = gr.Image(label="Uploaded Image", type="pil", interactive=False)
|
198 |
with gr.Column(scale=1):
|
199 |
output_text = gr.Textbox(label="Extracted Plain Text", lines=15, interactive=False)
|
200 |
-
|
|
|
201 |
|
202 |
-
def update_interface(image_filepath, xml_filepath
|
|
|
|
|
203 |
if image_filepath is None and xml_filepath is None:
|
204 |
-
return None, "Please upload an image and an XML file."
|
205 |
-
|
206 |
-
|
|
|
|
|
|
|
207 |
|
208 |
submit_button.click(
|
209 |
fn=update_interface,
|
210 |
-
inputs=[image_input, xml_input, show_overlay_checkbox
|
211 |
-
outputs=[output_image_orig, output_text
|
212 |
-
)
|
213 |
-
show_overlay_checkbox.change(
|
214 |
-
fn=update_interface,
|
215 |
-
inputs=[image_input, xml_input, show_overlay_checkbox],
|
216 |
-
outputs=[output_image_orig, output_text, output_image_overlay]
|
217 |
)
|
|
|
|
|
|
|
218 |
gr.Markdown("---")
|
219 |
gr.Markdown("### Example ALTO XML Snippet (for `String` element extraction):")
|
220 |
gr.Code(
|
@@ -228,7 +136,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
228 |
<TextLine WIDTH="684" HEIGHT="108" ID="p13_t1" HPOS="465" VPOS="196">
|
229 |
<String ID="p13_w1" CONTENT="Introduction" HPOS="465" VPOS="196" WIDTH="684" HEIGHT="108" STYLEREFS="font0"/>
|
230 |
</TextLine>
|
231 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
232 |
</PrintSpace>
|
233 |
</Page>
|
234 |
</Layout>
|
@@ -237,16 +151,23 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
237 |
interactive=False
|
238 |
)
|
239 |
|
|
|
240 |
if __name__ == "__main__":
|
241 |
try:
|
242 |
-
|
|
|
243 |
img_test.save("dummy_image.png")
|
244 |
print("Created dummy_image.png for testing.")
|
245 |
-
|
|
|
|
|
|
|
246 |
if not os.path.exists(example_xml_filename):
|
247 |
print(f"WARNING: Example XML '{example_xml_filename}' not found. Please create it (using the content from the prompt) or upload your own.")
|
|
|
248 |
except ImportError:
|
249 |
print("Pillow not installed, can't create dummy image.")
|
250 |
except Exception as e:
|
251 |
print(f"Error during setup: {e}")
|
|
|
252 |
demo.launch()
|
|
|
1 |
import gradio as gr
|
2 |
+
from PIL import Image # ImageDraw, ImageFont are no longer needed for overlay
|
3 |
import xml.etree.ElementTree as ET
|
4 |
import os
|
5 |
|
6 |
# --- Helper Functions ---
|
7 |
|
8 |
def get_alto_namespace(xml_file_path):
|
9 |
+
"""
|
10 |
+
Dynamically gets the ALTO namespace from the XML file.
|
11 |
+
"""
|
12 |
try:
|
13 |
tree = ET.parse(xml_file_path)
|
14 |
root = tree.getroot()
|
|
|
18 |
print(f"Error parsing XML to find namespace: {xml_file_path}")
|
19 |
return ''
|
20 |
|
21 |
+
def parse_alto_xml_for_text(xml_file_path):
|
22 |
+
"""
|
23 |
+
Parses an ALTO XML file to extract text content.
|
24 |
+
Returns:
|
25 |
+
- full_text (str): All extracted text concatenated.
|
26 |
+
"""
|
27 |
full_text_lines = []
|
28 |
+
|
29 |
if not xml_file_path or not os.path.exists(xml_file_path):
|
30 |
+
return "Error: XML file not provided or does not exist."
|
31 |
+
|
32 |
try:
|
33 |
ns_prefix = get_alto_namespace(xml_file_path)
|
34 |
tree = ET.parse(xml_file_path)
|
35 |
root = tree.getroot()
|
36 |
+
|
37 |
+
# Find all TextLine elements
|
38 |
for text_line in root.findall(f'.//{ns_prefix}TextLine'):
|
39 |
line_text_parts = []
|
40 |
for string_element in text_line.findall(f'{ns_prefix}String'):
|
41 |
text = string_element.get('CONTENT')
|
42 |
+
if text: # Ensure text is not None
|
43 |
line_text_parts.append(text)
|
44 |
+
# Also consider <SP/> (Space) elements if they contribute to word separation
|
45 |
+
# and are not implicitly handled by joining CONTENT attributes.
|
46 |
+
# For now, just joining CONTENT attributes.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
47 |
if line_text_parts:
|
48 |
full_text_lines.append(" ".join(line_text_parts))
|
49 |
+
|
50 |
+
return "\n".join(full_text_lines)
|
51 |
+
|
52 |
except ET.ParseError as e:
|
53 |
+
return f"Error parsing XML: {e}"
|
54 |
except Exception as e:
|
55 |
+
return f"An unexpected error occurred during XML parsing: {e}"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
56 |
|
57 |
+
# The draw_ocr_on_image function is no longer needed.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
|
59 |
# --- Gradio Interface Function ---
|
60 |
|
61 |
+
def process_image_and_xml(image_path, xml_path):
|
62 |
+
"""
|
63 |
+
Main function for the Gradio interface.
|
64 |
+
Processes the image and XML to return the image and extracted text.
|
65 |
+
"""
|
66 |
+
if image_path is None: # If no image is uploaded at all
|
67 |
+
return None, "Please upload an image."
|
68 |
+
|
69 |
try:
|
70 |
img_pil = Image.open(image_path).convert("RGB")
|
71 |
except Exception as e:
|
72 |
+
return None, f"Error loading image: {e}"
|
73 |
+
|
74 |
+
if xml_path is None: # If XML is missing, but image is present
|
75 |
+
return img_pil, "Please upload an OCR XML file."
|
76 |
+
|
77 |
+
# Both image and XML are presumably present
|
78 |
+
extracted_text = parse_alto_xml_for_text(xml_path)
|
79 |
+
|
80 |
+
return img_pil, extracted_text
|
81 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
|
83 |
# --- Create Gradio App ---
|
84 |
+
|
85 |
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
86 |
+
gr.Markdown("# OCR Viewer (ALTO XML) - Text Extractor")
|
87 |
gr.Markdown(
|
88 |
"Upload an image and its corresponding ALTO OCR XML file. "
|
89 |
+
"The app will display the image and extract/show the plain text."
|
|
|
90 |
)
|
91 |
+
|
92 |
with gr.Row():
|
93 |
with gr.Column(scale=1):
|
94 |
image_input = gr.File(label="Upload Image (PNG, JPG, etc.)", type="filepath")
|
95 |
xml_input = gr.File(label="Upload ALTO XML File (.xml)", type="filepath")
|
96 |
+
# show_overlay_checkbox has been removed
|
97 |
submit_button = gr.Button("Process Files", variant="primary")
|
98 |
+
|
99 |
with gr.Row():
|
100 |
with gr.Column(scale=1):
|
101 |
output_image_orig = gr.Image(label="Uploaded Image", type="pil", interactive=False)
|
102 |
with gr.Column(scale=1):
|
103 |
output_text = gr.Textbox(label="Extracted Plain Text", lines=15, interactive=False)
|
104 |
+
|
105 |
+
# output_image_overlay has been removed
|
106 |
|
107 |
+
def update_interface(image_filepath, xml_filepath):
|
108 |
+
# image_filepath and xml_filepath are now strings (paths) or None
|
109 |
+
|
110 |
if image_filepath is None and xml_filepath is None:
|
111 |
+
return None, "Please upload an image and an XML file."
|
112 |
+
# process_image_and_xml handles cases where one is None
|
113 |
+
|
114 |
+
img, text = process_image_and_xml(image_filepath, xml_filepath)
|
115 |
+
|
116 |
+
return img, text
|
117 |
|
118 |
submit_button.click(
|
119 |
fn=update_interface,
|
120 |
+
inputs=[image_input, xml_input], # show_overlay_checkbox removed
|
121 |
+
outputs=[output_image_orig, output_text] # output_image_overlay removed
|
|
|
|
|
|
|
|
|
|
|
122 |
)
|
123 |
+
|
124 |
+
# The .change event for show_overlay_checkbox has been removed
|
125 |
+
|
126 |
gr.Markdown("---")
|
127 |
gr.Markdown("### Example ALTO XML Snippet (for `String` element extraction):")
|
128 |
gr.Code(
|
|
|
136 |
<TextLine WIDTH="684" HEIGHT="108" ID="p13_t1" HPOS="465" VPOS="196">
|
137 |
<String ID="p13_w1" CONTENT="Introduction" HPOS="465" VPOS="196" WIDTH="684" HEIGHT="108" STYLEREFS="font0"/>
|
138 |
</TextLine>
|
139 |
+
<TextLine WIDTH="1798" HEIGHT="51" ID="p13_t2" HPOS="492" VPOS="523">
|
140 |
+
<String ID="p13_w2" CONTENT="Britain" HPOS="492" VPOS="523" WIDTH="166" HEIGHT="51" STYLEREFS="font1"/>
|
141 |
+
<SP WIDTH="24" VPOS="523" HPOS="658"/>
|
142 |
+
<String ID="p13_w3" CONTENT="1981" HPOS="682" VPOS="523" WIDTH="117" HEIGHT="51" STYLEREFS="font1"/>
|
143 |
+
<!-- ... more String and SP elements ... -->
|
144 |
+
</TextLine>
|
145 |
+
<!-- ... more TextLine elements ... -->
|
146 |
</PrintSpace>
|
147 |
</Page>
|
148 |
</Layout>
|
|
|
151 |
interactive=False
|
152 |
)
|
153 |
|
154 |
+
|
155 |
if __name__ == "__main__":
|
156 |
try:
|
157 |
+
# Create a dummy image for testing
|
158 |
+
img_test = Image.new('RGB', (2394, 3612), color = 'lightgray') # Dimensions from example XML
|
159 |
img_test.save("dummy_image.png")
|
160 |
print("Created dummy_image.png for testing.")
|
161 |
+
|
162 |
+
# Ensure the example XML file (189819724.34.xml) exists in the same directory
|
163 |
+
# or provide the correct path if it's elsewhere.
|
164 |
+
example_xml_filename = "189819724.34.xml"
|
165 |
if not os.path.exists(example_xml_filename):
|
166 |
print(f"WARNING: Example XML '{example_xml_filename}' not found. Please create it (using the content from the prompt) or upload your own.")
|
167 |
+
|
168 |
except ImportError:
|
169 |
print("Pillow not installed, can't create dummy image.")
|
170 |
except Exception as e:
|
171 |
print(f"Error during setup: {e}")
|
172 |
+
|
173 |
demo.launch()
|