| import xml.etree.ElementTree as ET | |
| from PIL import Image, ImageDraw, ImageFont | |
| import tempfile | |
| import os | |
| from typing import Optional, Tuple, List | |
| import math | |
| def htrflow_visualizer(image_path: str, htr_document_path: str) -> Optional[str]: | |
| """ | |
| Visualize HTR results by overlaying text regions and polygons on the original image. | |
| Args: | |
| image_path (str): Path to the original document image file | |
| htr_document_path (str): Path to the HTR XML file (ALTO or PAGE format) | |
| Returns: | |
| Optional[str]: Path to the generated visualization image, or None if failed | |
| """ | |
| try: | |
| if not image_path or not htr_document_path: | |
| return None | |
| image = Image.open(image_path) | |
| draw = ImageDraw.Draw(image) | |
| tree = ET.parse(htr_document_path) | |
| root = tree.getroot() | |
| if "alto" in root.tag.lower() or root.find(".//TextBlock") is not None: | |
| _visualize_alto_xml(draw, root, image.size) | |
| elif "PAGE" in root.tag or "PcGts" in root.tag: | |
| _visualize_page_xml(draw, root, image.size) | |
| else: | |
| if root.find(".//*[@points]") is not None: | |
| _visualize_page_xml(draw, root, image.size) | |
| else: | |
| _visualize_alto_xml(draw, root, image.size) | |
| temp_dir = tempfile.mkdtemp() | |
| output_path = os.path.join(temp_dir, "htr_visualization.png") | |
| image.save(output_path) | |
| return output_path | |
| except Exception: | |
| return None | |
| def _parse_points(points_str: str) -> List[Tuple[int, int]]: | |
| if not points_str: | |
| return [] | |
| points = [] | |
| for coord in points_str.strip().split(): | |
| if "," in coord: | |
| try: | |
| x, y = coord.split(",") | |
| points.append((int(float(x)), int(float(y)))) | |
| except ValueError: | |
| continue | |
| return points | |
| def _calculate_polygon_area(points: List[Tuple[int, int]]) -> float: | |
| if len(points) < 3: | |
| return 0 | |
| area = 0 | |
| n = len(points) | |
| for i in range(n): | |
| j = (i + 1) % n | |
| area += points[i][0] * points[j][1] | |
| area -= points[j][0] * points[i][1] | |
| return abs(area) / 2 | |
| def _get_dynamic_font_size( | |
| polygons: List[List[Tuple[int, int]]], image_size: Tuple[int, int] | |
| ) -> int: | |
| if not polygons: | |
| return 16 | |
| total_area = 0 | |
| valid_count = 0 | |
| for points in polygons: | |
| area = _calculate_polygon_area(points) | |
| if area > 0: | |
| total_area += area | |
| valid_count += 1 | |
| if valid_count == 0: | |
| return 16 | |
| avg_area = total_area / valid_count | |
| font_size = int(math.sqrt(avg_area) * 0.2) | |
| return max(12, min(72, font_size)) | |
| def _get_font(size: int) -> Optional[ImageFont.FreeTypeFont]: | |
| try: | |
| font_paths = [ | |
| "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf", | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", | |
| "/System/Library/Fonts/Helvetica.ttc", | |
| "C:\\Windows\\Fonts\\arial.ttf", | |
| ] | |
| for font_path in font_paths: | |
| if os.path.exists(font_path): | |
| return ImageFont.truetype(font_path, size) | |
| return ImageFont.load_default() | |
| except: | |
| return ImageFont.load_default() | |
| def _get_namespace(root: ET.Element) -> Optional[str]: | |
| if "}" in root.tag: | |
| return root.tag.split("}")[0] + "}" | |
| return None | |
| def _visualize_page_xml( | |
| draw: ImageDraw.Draw, root: ET.Element, image_size: Tuple[int, int] | |
| ): | |
| text_lines = [] | |
| for elem in root.iter(): | |
| if elem.tag.endswith("TextLine"): | |
| text_lines.append(elem) | |
| line_data = [] | |
| all_polygons = [] | |
| for text_line in text_lines: | |
| coords_elem = None | |
| for child in text_line: | |
| if child.tag.endswith("Coords"): | |
| coords_elem = child | |
| break | |
| if coords_elem is not None: | |
| points_str = coords_elem.get("points", "") | |
| points = _parse_points(points_str) | |
| if len(points) >= 3: | |
| text_content = "" | |
| confidence = None | |
| for te in text_line.iter(): | |
| if te.tag.endswith("Unicode") and te.text: | |
| text_content = te.text.strip() | |
| break | |
| for te in text_line.iter(): | |
| if te.tag.endswith("TextEquiv"): | |
| conf_str = te.get("conf") | |
| if conf_str: | |
| try: | |
| confidence = float(conf_str) | |
| except: | |
| pass | |
| break | |
| display_text = text_content | |
| if confidence is not None: | |
| display_text = f"{text_content} ({confidence:.3f})" | |
| line_data.append((points, display_text)) | |
| all_polygons.append(points) | |
| font_size = _get_dynamic_font_size(all_polygons, image_size) | |
| font = _get_font(font_size) | |
| for i, (points, text) in enumerate(line_data): | |
| color = "red" if i % 2 == 0 else "blue" | |
| draw.polygon(points, outline=color, width=2) | |
| if text: | |
| centroid_x = sum(p[0] for p in points) // len(points) | |
| centroid_y = sum(p[1] for p in points) // len(points) | |
| if font != ImageFont.load_default(): | |
| bbox = draw.textbbox( | |
| (centroid_x, centroid_y), text, font=font, anchor="mm" | |
| ) | |
| bbox = (bbox[0] - 2, bbox[1] - 2, bbox[2] + 2, bbox[3] + 2) | |
| draw.rectangle(bbox, fill=(255, 255, 255, 200), outline="black") | |
| draw.text( | |
| (centroid_x, centroid_y), text, fill="black", font=font, anchor="mm" | |
| ) | |
| else: | |
| draw.text((centroid_x, centroid_y), text, fill="black") | |
| def _visualize_alto_xml( | |
| draw: ImageDraw.Draw, root: ET.Element, image_size: Tuple[int, int] | |
| ): | |
| namespace = _get_namespace(root) | |
| text_lines = [] | |
| for elem in root.iter(): | |
| if elem.tag.endswith("TextLine"): | |
| text_lines.append(elem) | |
| line_data = [] | |
| all_polygons = [] | |
| for text_line in text_lines: | |
| points = [] | |
| for shape in text_line.iter(): | |
| if shape.tag.endswith("Shape"): | |
| for polygon in shape.iter(): | |
| if polygon.tag.endswith("Polygon"): | |
| points_str = polygon.get("POINTS", "") | |
| points = _parse_points(points_str) | |
| break | |
| break | |
| if len(points) >= 3: | |
| text_content = "" | |
| confidence = None | |
| for string_elem in text_line.iter(): | |
| if string_elem.tag.endswith("String"): | |
| text_content = string_elem.get("CONTENT", "") | |
| wc_str = string_elem.get("WC") | |
| if wc_str: | |
| try: | |
| confidence = float(wc_str) | |
| except: | |
| pass | |
| break | |
| display_text = text_content | |
| if confidence is not None: | |
| display_text = f"{text_content} ({confidence:.3f})" | |
| line_data.append((points, display_text)) | |
| all_polygons.append(points) | |
| font_size = _get_dynamic_font_size(all_polygons, image_size) | |
| font = _get_font(font_size) | |
| for i, (points, text) in enumerate(line_data): | |
| color = "red" if i % 2 == 0 else "blue" | |
| draw.polygon(points, outline=color, width=2) | |
| if text: | |
| centroid_x = sum(p[0] for p in points) // len(points) | |
| centroid_y = sum(p[1] for p in points) // len(points) | |
| if font != ImageFont.load_default(): | |
| bbox = draw.textbbox( | |
| (centroid_x, centroid_y), text, font=font, anchor="mm" | |
| ) | |
| bbox = (bbox[0] - 2, bbox[1] - 2, bbox[2] + 2, bbox[3] + 2) | |
| draw.rectangle(bbox, fill=(255, 255, 255, 200), outline="black") | |
| draw.text( | |
| (centroid_x, centroid_y), text, fill="black", font=font, anchor="mm" | |
| ) | |
| else: | |
| draw.text((centroid_x, centroid_y), text, fill="black") | |