| """Template discovery and annotation resolution. |
| |
| Scans ``image/`` recursively for scanned bank document templates, assigns |
| each a stable ``template_id`` (``bank_doc_01``, ``bank_doc_02``, ...) based |
| on a deterministic sort of their paths, and resolves signature-box |
| annotations for each using, in order of preference: |
| |
| 1. **Mode 2 (preferred when present):** a manually authored |
| ``annotations/<template_id>.json`` file. |
| 2. **Mode 1 (fallback):** automatic CV-based detection |
| (:mod:`src.templates.box_detector`), whose result is cached to |
| ``annotations/<template_id>.json`` so subsequent runs reuse it without |
| re-running detection. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from src.templates.annotation_schema import TemplateAnnotation, TemplatePageAnnotation, SignatureBoxAnnotation |
| from src.templates.box_detector import detect_signature_boxes |
| from src.utils.image_io import discover_images, load_image_rgb |
| from src.utils.logging_utils import get_logger |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| @dataclass |
| class Template: |
| """A discovered template document: its image path and resolved annotation.""" |
|
|
| template_id: str |
| image_path: Path |
| annotation: TemplateAnnotation |
|
|
|
|
| class TemplateLoader: |
| """Discovers bank document templates and resolves their signature-box annotations. |
| |
| Example: |
| >>> loader = TemplateLoader(image_dir="image", annotations_dir="annotations") |
| >>> templates = loader.discover() |
| >>> tmpl = loader.get("bank_doc_01") |
| """ |
|
|
| def __init__(self, image_dir: "str | Path", annotations_dir: "str | Path") -> None: |
| self.image_dir = Path(image_dir) |
| self.annotations_dir = Path(annotations_dir) |
| self.annotations_dir.mkdir(parents=True, exist_ok=True) |
| self._templates: "dict[str, Template] | None" = None |
|
|
| def discover(self) -> "dict[str, Template]": |
| """Discover all templates under ``image_dir`` and resolve their annotations. |
| |
| Returns: |
| An ordered mapping of ``template_id -> Template``, sorted by id. |
| """ |
| if self._templates is not None: |
| return self._templates |
|
|
| image_paths = discover_images(self.image_dir) |
| if not image_paths: |
| raise FileNotFoundError( |
| f"No supported template images found under '{self.image_dir}'. " |
| "Expected PNG/JPG/JPEG/TIFF/BMP files." |
| ) |
|
|
| templates: "dict[str, Template]" = {} |
| for idx, image_path in enumerate(image_paths, start=1): |
| template_id = f"bank_doc_{idx:02d}" |
| annotation = self._resolve_annotation(template_id, image_path) |
| templates[template_id] = Template(template_id, image_path, annotation) |
| logger.info( |
| "Registered template '%s' -> %s (%d signature box(es), auto_detected=%s)", |
| template_id, image_path, sum(len(p.signature_boxes) for p in annotation.pages), |
| annotation.auto_detected, |
| ) |
|
|
| self._templates = templates |
| return templates |
|
|
| def get(self, template_id: str) -> Template: |
| """Return a discovered template by id, raising ``KeyError`` if unknown.""" |
| templates = self.discover() |
| if template_id not in templates: |
| raise KeyError(f"Unknown template id '{template_id}'. Known: {sorted(templates)}") |
| return templates[template_id] |
|
|
| def list_ids(self) -> list[str]: |
| """Return all discovered template ids, sorted.""" |
| return sorted(self.discover().keys()) |
|
|
| def random_template(self, rng: np.random.Generator) -> Template: |
| """Pick a uniformly random template using the given seeded generator.""" |
| ids = self.list_ids() |
| choice = ids[int(rng.integers(0, len(ids)))] |
| return self.get(choice) |
|
|
| def _resolve_annotation(self, template_id: str, image_path: Path) -> TemplateAnnotation: |
| annotation_path = self.annotations_dir / f"{template_id}.json" |
| if annotation_path.exists(): |
| data = json.loads(annotation_path.read_text(encoding="utf-8")) |
| return TemplateAnnotation.model_validate(data) |
|
|
| logger.warning( |
| "No manual annotation found for '%s'; falling back to automatic CV detection.", |
| template_id, |
| ) |
| annotation = self._auto_detect_annotation(template_id, image_path) |
| annotation_path.write_text( |
| json.dumps(annotation.model_dump(), indent=2), encoding="utf-8" |
| ) |
| logger.info("Cached auto-detected annotation to %s", annotation_path) |
| return annotation |
|
|
| def _auto_detect_annotation(self, template_id: str, image_path: Path) -> TemplateAnnotation: |
| image = load_image_rgb(image_path) |
| h, w = image.shape[:2] |
| boxes = detect_signature_boxes(image) |
|
|
| if not boxes: |
| |
| |
| boxes = [(w * 0.55, h * 0.80, w * 0.92, h * 0.90)] |
| logger.warning( |
| "'%s': no lines detected; using generic lower-right fallback box.", template_id |
| ) |
|
|
| signature_boxes = [ |
| SignatureBoxAnnotation( |
| id=f"sig{i + 1}", x1=x1, y1=y1, x2=x2, y2=y2, role="Signer", name="Signer", |
| ) |
| for i, (x1, y1, x2, y2) in enumerate(boxes) |
| ] |
| page = TemplatePageAnnotation(page=1, page_size=(w, h), dpi=300, signature_boxes=signature_boxes) |
| return TemplateAnnotation( |
| id=template_id, |
| source_image=str(image_path.as_posix()), |
| document_type="unknown", |
| pages=[page], |
| auto_detected=True, |
| ) |
|
|
