File size: 5,147 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
"""Renders a multi-panel debug figure for one generated sample.

Shows the full rendered document with the target box, detected bounding
box, and detected polygon overlaid; a crop of the signature field; a mask
reconstructed from the detected ink polygon; the resolved coverage metrics;
and a difficulty-component breakdown bar chart.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import cv2
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from src.utils.logging_utils import get_logger

logger = get_logger(__name__)


def _draw_overlays(image: np.ndarray, metadata: dict[str, Any]) -> np.ndarray:
    out = image.copy()
    geometry = metadata["geometry"]

    def as_int_box(box):
        return tuple(int(round(v)) for v in box)

    tb = as_int_box(geometry["target_box"])
    ab = as_int_box(geometry["actual_box"])
    cv2.rectangle(out, (tb[0], tb[1]), (tb[2], tb[3]), (0, 200, 0), 2)
    cv2.rectangle(out, (ab[0], ab[1]), (ab[2], ab[3]), (30, 80, 230), 2)

    polygon = geometry.get("signature_polygon_px") or []
    if len(polygon) >= 3:
        pts = np.array(polygon, dtype=np.int32).reshape(-1, 1, 2)
        cv2.polylines(out, [pts], isClosed=True, color=(220, 30, 30), thickness=2)

    return out


def _crop_region(image: np.ndarray, box: list, pad: int = 25) -> np.ndarray:
    h, w = image.shape[:2]
    x1, y1, x2, y2 = (int(round(v)) for v in box)
    x1, y1 = max(x1 - pad, 0), max(y1 - pad, 0)
    x2, y2 = min(x2 + pad, w), min(y2 + pad, h)
    if x2 <= x1 or y2 <= y1:
        return image
    return image[y1:y2, x1:x2]


def _reconstruct_mask(shape: tuple[int, int], box: list, polygon: list) -> np.ndarray:
    mask = np.zeros(shape, dtype=np.uint8)
    if len(polygon) >= 3:
        pts = np.array(polygon, dtype=np.int32).reshape(-1, 1, 2)
        cv2.fillPoly(mask, [pts], 255)
    x1, y1, x2, y2 = (int(round(v)) for v in box)
    pad = 25
    x1c, y1c = max(x1 - pad, 0), max(y1 - pad, 0)
    x2c, y2c = min(x2 + pad, shape[1]), min(y2 + pad, shape[0])
    if x2c <= x1c or y2c <= y1c:
        return mask
    return mask[y1c:y2c, x1c:x2c]


def save_visualization(path: "str | Path", final_image: np.ndarray, metadata: dict[str, Any]) -> None:
    """Render and save a multi-panel debug figure for one generated sample.

    Args:
        path: Output PNG path.
        final_image: The fully rendered (and augmented) document.
        metadata: The sample's full metadata dict, as built by
            :class:`~src.metadata.metadata_exporter.MetadataExporter`.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    geometry = metadata["geometry"]
    difficulty = metadata["difficulty"]
    metrics = metadata["metrics"]

    overlaid = _draw_overlays(final_image, metadata)
    crop = _crop_region(final_image, geometry["actual_box"])
    mask = _reconstruct_mask(final_image.shape[:2], geometry["actual_box"], geometry["signature_polygon_px"])

    fig = plt.figure(figsize=(15, 9.5), dpi=110)
    fig.suptitle(
        f"{metadata['document_id']}  |  template={metadata['template']}  |  box={metadata['geometry']['target_box']}",
        x=0.02, ha="left", fontsize=11, family="monospace",
    )
    grid = fig.add_gridspec(2, 3, width_ratios=[2.0, 1.0, 1.0], height_ratios=[1.0, 1.0])

    ax_full = fig.add_subplot(grid[:, 0])
    ax_full.imshow(overlaid)
    ax_full.axis("off")

    ax_crop = fig.add_subplot(grid[0, 1])
    ax_crop.imshow(crop)
    ax_crop.set_title("Signature field (cropped)")
    ax_crop.axis("off")

    ax_mask = fig.add_subplot(grid[0, 2])
    ax_mask.imshow(mask, cmap="gray")
    ax_mask.set_title("Detected ink mask")
    ax_mask.axis("off")

    ax_metrics = fig.add_subplot(grid[1, 1])
    ax_metrics.axis("off")
    metric_lines = [
        f"box_coverage_pct:  {metrics['box_coverage_pct']:.2f}",
        f"containment_pct:   {metrics['containment_pct']:.2f}",
        f"ink_fill_ratio:    {metrics['ink_fill_ratio']:.3f}",
        f"polygon_overlap:   {metrics['polygon_overlap_proxy']:.3f}",
        f"dist_from_center:  {metrics['distance_from_box_center_px']:.1f}px",
        f"margin_util:       {metrics['margin_utilization']:.3f}",
        f"mean_dark_int:     {metrics['mean_dark_intensity']:.1f}",
        f"gray_stddev:       {metrics['gray_stddev']:.1f}",
        f"signer:            {metadata['signer']['name']} ({metadata['signer']['role']})",
    ]
    ax_metrics.text(0.0, 1.0, "\n".join(metric_lines), va="top", ha="left", family="monospace", fontsize=9)
    ax_metrics.set_title("Coverage metrics")

    ax_diff = fig.add_subplot(grid[1, 2])
    components = ["misplacement", "whitespace", "faintness", "smallness", "low_texture"]
    values = [difficulty[c] for c in components]
    colors = ["#d62728" if v >= 0.66 else "#ff7f0e" if v >= 0.33 else "#2ca02c" for v in values]
    ax_diff.barh(components, values, color=colors)
    ax_diff.set_xlim(0, 1)
    ax_diff.set_title(f"Difficulty: {difficulty['score']:.3f} ({difficulty['tier']})")

    fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.94))
    fig.savefig(path)
    plt.close(fig)