"""OCR evaluation for outputs produced by controlnet_benchmark.py.""" from __future__ import annotations import argparse import json import math from pathlib import Path from PIL import Image, ImageChops, ImageDraw, ImageFilter, ImageFont, ImageOps from benchmarks.text_accuracy import ( centered_text_crop, character_error_rate, normalize_ocr, run_ocr, ) def _normalized_outline( image: Image.Image, size: tuple[int, int] = (384, 128), *, light_ink: bool = False, ) -> Image.Image: grayscale = image.convert("L") if light_ink: ink = grayscale.point(lambda value: 255 if value >= 128 else 0) else: ink = grayscale.point(lambda value: 255 if value < 128 else 0) bbox = ink.getbbox() if bbox is None: return Image.new("L", size, 0) cropped = ImageOps.expand(ink.crop(bbox), border=4, fill=0) normalized = cropped.resize(size, Image.Resampling.LANCZOS) normalized = normalized.point(lambda value: 255 if value >= 128 else 0) outer = normalized.filter(ImageFilter.MaxFilter(3)) inner = normalized.filter(ImageFilter.MinFilter(3)) return ImageChops.subtract(outer, inner) def glyph_similarity(image: Image.Image, control: Image.Image) -> float: generated = _normalized_outline(image.crop(centered_text_crop(image))) expected = _normalized_outline(control, light_ink=True) generated_dilated = generated.filter(ImageFilter.MaxFilter(7)) expected_dilated = expected.filter(ImageFilter.MaxFilter(7)) generated_pixels = sum(1 for value in generated.getdata() if value) expected_pixels = sum(1 for value in expected.getdata() if value) if not generated_pixels or not expected_pixels: return 0.0 generated_hit = sum( 1 for edge, nearby in zip(generated.getdata(), expected_dilated.getdata()) if edge and nearby ) expected_hit = sum( 1 for edge, nearby in zip(expected.getdata(), generated_dilated.getdata()) if edge and nearby ) precision = generated_hit / generated_pixels recall = expected_hit / expected_pixels return round(2 * precision * recall / max(1e-9, precision + recall), 4) def light_ink_margin(image: Image.Image) -> int: ink = image.convert("L").point(lambda value: 255 if value >= 128 else 0) bbox = ink.getbbox() if bbox is None: return -1 width, height = image.size return min(bbox[0], bbox[1], width - bbox[2], height - bbox[3]) def write_contact_sheets( results: list[dict[str, object]], output_dir: Path, *, columns: int = 5, rows: int = 4, ) -> list[str]: page_size = columns * rows thumb_size = 240 label_height = 44 font_path = Path("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf") font = ImageFont.truetype(str(font_path), 22) if font_path.exists() else ImageFont.load_default() names: list[str] = [] for page_index in range(math.ceil(len(results) / page_size)): page_rows = results[page_index * page_size : (page_index + 1) * page_size] sheet = Image.new( "RGB", (columns * thumb_size, rows * (thumb_size + label_height)), "white", ) draw = ImageDraw.Draw(sheet) for cell_index, result in enumerate(page_rows): image_path = output_dir / str(result["image"]) with Image.open(image_path) as source: thumbnail = source.convert("RGB") thumbnail.thumbnail((thumb_size, thumb_size), Image.Resampling.LANCZOS) x = (cell_index % columns) * thumb_size y = (cell_index // columns) * (thumb_size + label_height) sheet.paste( thumbnail, (x + (thumb_size - thumbnail.width) // 2, y), ) draw.text((x + 6, y + thumb_size + 4), str(result["text"]), fill="black", font=font) name = f"contact-sheet-{page_index + 1:02d}.jpg" sheet.save(output_dir / name, quality=90) names.append(name) return names def summarize(results: list[dict[str, object]]) -> dict[str, dict[str, float | int]]: strategies = sorted({str(result["strategy"]) for result in results}) summary: dict[str, dict[str, float | int]] = {} for strategy in strategies: rows = [result for result in results if result["strategy"] == strategy] summary[strategy] = { "samples": len(rows), "exact": sum(bool(row["exact_match"]) for row in rows), "mean_cer": round(sum(float(row["cer"]) for row in rows) / len(rows), 4), "mean_glyph_similarity": round( sum(float(row.get("glyph_similarity", 0.0)) for row in rows) / len(rows), 4, ), "edge_touching_controls": sum( int(row.get("control_margin_px", 1)) <= 0 for row in rows ), } return summary def run(report_path: Path) -> dict[str, object]: generation = json.loads(report_path.read_text(encoding="utf-8")) output_dir = report_path.parent results: list[dict[str, object]] = [] for row in generation["results"]: image_path = output_dir / row["image"] with Image.open(image_path) as image: crop = centered_text_crop(image) actual = run_ocr(image_path, "rus+eng", crop=crop) gold = normalize_ocr(row["text"]) similarity = 0.0 control_margin_px: int | None = None if row.get("control"): with Image.open(image_path) as image, Image.open(output_dir / row["control"]) as control: similarity = glyph_similarity(image, control) control_margin_px = light_ink_margin(control) results.append( { **row, "gold": gold, "ocr": actual, "exact_match": actual == gold, "cer": round(character_error_rate(gold, actual), 4), "glyph_similarity": similarity, "control_margin_px": control_margin_px, } ) contact_sheets = write_contact_sheets(results, output_dir) report: dict[str, object] = { "generation_report": report_path.name, "summary": summarize(results), "contact_sheets": contact_sheets, "results": results, } (output_dir / "accuracy-report.json").write_text( json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8", ) return report def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("generation_report", type=Path) return parser.parse_args() if __name__ == "__main__": args = parse_args() completed = run(args.generation_report.resolve()) print(json.dumps(completed, ensure_ascii=False))