Create visualizer.py

visualizer.py

@@ -0,0 +1,257 @@
+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")