File size: 5,833 Bytes
97b64b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """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:
# Last-resort fallback: a generic box in the lower-right quadrant,
# where signature fields conventionally sit on Western bank forms.
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,
)
|