"""Visualize OCR results — overlay recognized text directly on detected regions. Features: - Text overlaid on word bounding boxes, scaled to fit - Semi-transparent background behind text for readability - Color-coded line bounding boxes - Confidence heat-map coloring (green=high, red=low) - Summary panel with statistics """ import sys import time from pathlib import Path import numpy as np from PIL import Image, ImageDraw, ImageFont sys.path.insert(0, str(Path(__file__).parent.parent)) from ocr.engine_onnx import OcrEngineOnnx from ocr.models import BoundingRect # ─── Color helpers ──────────────────────────────────────────────────────── def _conf_color(conf: float) -> tuple[int, int, int]: """Map confidence 0..1 → red..yellow..green.""" if conf >= 0.85: return (40, 180, 40) elif conf >= 0.6: t = (conf - 0.6) / 0.25 return (int(220 * (1 - t)), int(180 * t + 100), 40) else: return (220, 60, 40) _LINE_COLORS = [ (70, 130, 255), (255, 100, 70), (50, 200, 100), (255, 180, 40), (180, 80, 255), (40, 200, 200), (255, 80, 180), (160, 200, 60), ] # ─── Font helpers ──────────────────────────────────────────────────────── def _load_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: """Try to load TrueType font, fallback to default.""" for name in ("arial.ttf", "Arial.ttf", "segoeui.ttf", "msyh.ttc", "NotoSansCJK-Regular.ttc", "DejaVuSans.ttf"): try: return ImageFont.truetype(name, size) except Exception: pass return ImageFont.load_default() def _fit_font_size( text: str, box_w: float, box_h: float, min_size: int = 8, max_size: int = 120, ) -> int: """Binary search for font size that fits text into box_w × box_h.""" lo, hi = min_size, max_size best = min_size while lo <= hi: mid = (lo + hi) // 2 font = _load_font(mid) bbox = font.getbbox(text) tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1] if tw <= box_w * 0.95 and th <= box_h * 0.88: best = mid lo = mid + 1 else: hi = mid - 1 return best # ─── Drawing helpers ───────────────────────────────────────────────────── def _draw_quad( draw: ImageDraw.ImageDraw, b: BoundingRect, color: tuple, width: int = 2, ) -> None: """Draw a quadrilateral outline.""" pts = [(b.x1, b.y1), (b.x2, b.y2), (b.x3, b.y3), (b.x4, b.y4)] draw.polygon(pts, outline=color, width=width) def _overlay_text_on_word( overlay: Image.Image, word_text: str, b: BoundingRect, conf: float, ) -> None: """Draw text overlaid inside the word bounding box with semi-transparent bg.""" xs = [b.x1, b.x2, b.x3, b.x4] ys = [b.y1, b.y2, b.y3, b.y4] x_min, x_max = min(xs), max(xs) y_min, y_max = min(ys), max(ys) box_w = x_max - x_min box_h = y_max - y_min if box_w < 3 or box_h < 3: return # Fit font to box font_size = _fit_font_size(word_text, box_w, box_h) font = _load_font(font_size) # Measure text bbox = font.getbbox(word_text) tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1] # Center text in box tx = x_min + (box_w - tw) / 2 ty = y_min + (box_h - th) / 2 - bbox[1] # Semi-transparent white background behind text bg = Image.new("RGBA", overlay.size, (0, 0, 0, 0)) bg_draw = ImageDraw.Draw(bg) pad = 2 bg_draw.rectangle( [tx - pad, y_min + (box_h - th) / 2 - pad, tx + tw + pad, y_min + (box_h + th) / 2 + pad], fill=(255, 255, 255, 170), ) overlay.alpha_composite(bg) # Text color based on confidence text_color = _conf_color(conf) draw = ImageDraw.Draw(overlay) draw.text((tx, ty), word_text, fill=(*text_color, 255), font=font) # ─── Summary panel ─────────────────────────────────────────────────────── def _draw_summary( overlay: Image.Image, n_lines: int, n_words: int, avg_conf: float, angle: float, elapsed: float, img_size: tuple[int, int], ) -> None: """Draw summary statistics panel at the top of the image.""" font = _load_font(16) stats = ( f"Lines: {n_lines} | Words: {n_words} | " f"Conf: {avg_conf:.1%} | Angle: {angle:.1f}\u00b0 | " f"Time: {elapsed:.0f}ms | {img_size[0]}\u00d7{img_size[1]}" ) bbox = font.getbbox(stats) th = bbox[3] - bbox[1] panel_h = th + 12 # Semi-transparent dark panel bg = Image.new("RGBA", overlay.size, (0, 0, 0, 0)) bg_draw = ImageDraw.Draw(bg) bg_draw.rectangle([0, 0, overlay.width, panel_h], fill=(0, 0, 0, 180)) overlay.alpha_composite(bg) draw = ImageDraw.Draw(overlay) draw.text((8, 4), stats, fill=(255, 255, 255, 255), font=font) # ═══════════════════════════════════════════════════════════════════════════ # Main visualization # ═══════════════════════════════════════════════════════════════════════════ def visualize( image_path: str, output_path: str = "result_ocr.png", show_word_boxes: bool = True, show_line_boxes: bool = True, show_text_overlay: bool = True, show_confidence: bool = True, ) -> None: """Run OCR and visualize results with text overlay. Args: image_path: Input image path. output_path: Output path for annotated image. show_word_boxes: Draw word-level bounding boxes. show_line_boxes: Draw line-level bounding boxes. show_text_overlay: Overlay recognized text on words. show_confidence: Show confidence % below words. """ img = Image.open(image_path).convert("RGBA") engine = OcrEngineOnnx() t0 = time.perf_counter() result = engine.recognize_pil(img.convert("RGB")) elapsed_ms = (time.perf_counter() - t0) * 1000 if result.error: print(f"Error: {result.error}") return overlay = img.copy() draw = ImageDraw.Draw(overlay) n_words = sum(len(l.words) for l in result.lines) for i, line in enumerate(result.lines): lc = _LINE_COLORS[i % len(_LINE_COLORS)] # Line-level bounding box (thicker) if show_line_boxes and line.bounding_rect: _draw_quad(draw, line.bounding_rect, color=lc, width=3) for word in line.words: if not word.bounding_rect: continue b = word.bounding_rect # Word-level bounding box (confidence-colored) if show_word_boxes: wc = _conf_color(word.confidence) _draw_quad(draw, b, color=wc, width=2) # Overlay text inside the word box if show_text_overlay: _overlay_text_on_word(overlay, word.text, b, word.confidence) draw = ImageDraw.Draw(overlay) # refresh after alpha_composite # Confidence label below word box if show_confidence: xs = [b.x1, b.x2, b.x3, b.x4] ys = [b.y1, b.y2, b.y3, b.y4] cx = sum(xs) / 4 y_bot = max(ys) + 2 conf_font = _load_font(11) label = f"{word.confidence:.0%}" lbbox = conf_font.getbbox(label) lw = lbbox[2] - lbbox[0] draw.text( (cx - lw / 2, y_bot), label, fill=(*_conf_color(word.confidence), 220), font=conf_font, ) # Summary panel _draw_summary( overlay, n_lines=len(result.lines), n_words=n_words, avg_conf=result.average_confidence, angle=result.text_angle or 0.0, elapsed=elapsed_ms, img_size=(img.width, img.height), ) # Save as RGB final = Image.new("RGB", overlay.size, (255, 255, 255)) final.paste(overlay, mask=overlay.split()[3]) final.save(output_path, quality=95) print(f"\nSaved: {output_path}") print(f"Text: \"{result.text}\"") print(f"Lines: {len(result.lines)}, Words: {n_words}, " f"Conf: {result.average_confidence:.1%}, Time: {elapsed_ms:.0f}ms") for i, line in enumerate(result.lines): words_info = " ".join( f'"{w.text}"({w.confidence:.0%})' for w in line.words ) print(f" L{i}: {words_info}") if __name__ == "__main__": image_path = sys.argv[1] if len(sys.argv) > 1 else "test3.png" output_path = sys.argv[2] if len(sys.argv) > 2 else "result_ocr.png" visualize(image_path, output_path)