davanstrien HF Staff commited on
Commit
5f3165f
·
verified ·
1 Parent(s): 898d181

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +81 -160
app.py CHANGED
@@ -1,11 +1,14 @@
1
  import gradio as gr
2
- from PIL import Image, ImageDraw, ImageFont
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 parse_alto_xml(xml_file_path):
 
 
 
 
 
19
  full_text_lines = []
20
- ocr_data = []
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
- try:
34
- hpos = int(float(string_element.get('HPOS')))
35
- vpos = int(float(string_element.get('VPOS')))
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
- return "\n".join(full_text_lines), ocr_data
 
 
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
- final_text_top_y = tentative_text_top_y
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, show_overlay):
158
- if image_path is None:
159
- return None, "Please upload an image.", None
 
 
 
 
 
160
  try:
161
  img_pil = Image.open(image_path).convert("RGB")
162
  except Exception as e:
163
- return None, f"Error loading image: {e}", None
164
-
165
- if xml_path is None:
166
- return img_pil, "Please upload an OCR XML file.", None
167
-
168
- extracted_text, ocr_box_data = parse_alto_xml(xml_path)
169
- overlay_image_pil = None
170
- if show_overlay:
171
- if ocr_box_data:
172
- img_for_overlay = img_pil.copy()
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, extract and show the plain text, "
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 = gr.Checkbox(label="Show OCR Overlay on Image", value=False)
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
- output_image_overlay = gr.Image(label="Image with OCR Overlay", type="pil", interactive=False, visible=True)
 
201
 
202
- def update_interface(image_filepath, xml_filepath, show_overlay_val):
 
 
203
  if image_filepath is None and xml_filepath is None:
204
- return None, "Please upload an image and an XML file.", None
205
- img, text, overlay_img = process_image_and_xml(image_filepath, xml_filepath, show_overlay_val)
206
- return img, text, overlay_img
 
 
 
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, output_image_overlay]
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
- <!-- ... more TextLine and String elements ... -->
 
 
 
 
 
 
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
- img_test = Image.new('RGB', (2394, 3612), color = 'lightgray')
 
243
  img_test.save("dummy_image.png")
244
  print("Created dummy_image.png for testing.")
245
- example_xml_filename = "189819724.34.xml"
 
 
 
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()