← back to Wallco Ai

scripts/edges-scan.py

220 lines

#!/usr/bin/env python3
"""edges-agent — 6-lens seamless-tile defect scanner.

Samples 2-pixel-thick strips at four wrap edges (top/bottom/left/right) and the
two internal midlines (h_mid, v_mid). Computes mean LAB ΔE between the two
sides of each lens — low score = seamless at that lens, high score = visible
seam.

Outputs JSON to stdout (default) or a human-readable scorecard (--text).
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

# Tunable thresholds (1024x1024 SDXL output baseline).
PASS_MAX = 5.0
WARN_MAX = 12.0

# Strip thickness on each side of the lens line.
STRIP_PX = 2

LENSES = ['top', 'bottom', 'left', 'right', 'h_mid', 'v_mid']


def verdict_for(score: float) -> str:
    if score <= PASS_MAX:
        return 'PASS'
    if score <= WARN_MAX:
        return 'WARN'
    return 'FAIL'


def rgb_to_lab_strip(arr):
    """Convert a numpy RGB strip (H, W, 3) uint8 → LAB float64. Pure-numpy
    sRGB→LAB so we don't pull in skimage."""
    import numpy as np
    rgb = arr.astype(np.float64) / 255.0
    # sRGB → linear RGB
    a = 0.055
    linear = np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + a) / (1 + a)) ** 2.4)
    # linear RGB → XYZ (D65)
    M = np.array([
        [0.4124564, 0.3575761, 0.1804375],
        [0.2126729, 0.7151522, 0.0721750],
        [0.0193339, 0.1191920, 0.9503041],
    ])
    xyz = linear @ M.T
    # XYZ → LAB (D65 white)
    white = np.array([0.95047, 1.0, 1.08883])
    xyz_n = xyz / white
    eps = 216 / 24389
    kappa = 24389 / 27
    f = np.where(xyz_n > eps, np.cbrt(xyz_n), (kappa * xyz_n + 16) / 116)
    L = 116 * f[..., 1] - 16
    a_ch = 500 * (f[..., 0] - f[..., 1])
    b_ch = 200 * (f[..., 1] - f[..., 2])
    return np.stack([L, a_ch, b_ch], axis=-1)


def lens_score(img_np, lens: str) -> float:
    """Mean ΔE76 across the lens strip."""
    import numpy as np
    H, W = img_np.shape[:2]

    if lens == 'top':
        a = img_np[0:STRIP_PX, :, :]
        b = img_np[H - STRIP_PX:H, :, :]
    elif lens == 'bottom':
        a = img_np[H - STRIP_PX:H, :, :]
        b = img_np[0:STRIP_PX, :, :]
    elif lens == 'left':
        a = img_np[:, 0:STRIP_PX, :]
        b = img_np[:, W - STRIP_PX:W, :]
    elif lens == 'right':
        a = img_np[:, W - STRIP_PX:W, :]
        b = img_np[:, 0:STRIP_PX, :]
    elif lens == 'h_mid':
        m = H // 2
        a = img_np[m - STRIP_PX:m, :, :]
        b = img_np[m:m + STRIP_PX, :, :]
    elif lens == 'v_mid':
        m = W // 2
        a = img_np[:, m - STRIP_PX:m, :]
        b = img_np[:, m:m + STRIP_PX, :]
    else:
        raise ValueError(f'unknown lens {lens}')

    # Resize sides to identical shape for safety (top/bottom both have STRIP_PX
    # rows × W cols, left/right both have H rows × STRIP_PX cols — already
    # matched; this is just defensive).
    if a.shape != b.shape:
        raise ValueError(f'shape mismatch in lens {lens}: {a.shape} vs {b.shape}')

    lab_a = rgb_to_lab_strip(a)
    lab_b = rgb_to_lab_strip(b)
    delta = lab_a - lab_b
    de76 = (delta ** 2).sum(axis=-1) ** 0.5
    return float(de76.mean())


def scan_path(png_path: Path) -> dict:
    from PIL import Image
    import numpy as np

    if not png_path.exists():
        return {'ok': False, 'error': f'file not found: {png_path}'}
    img = Image.open(png_path).convert('RGB')
    arr = np.asarray(img)
    if arr.ndim != 3 or arr.shape[2] != 3:
        return {'ok': False, 'error': f'unexpected image shape: {arr.shape}'}
    H, W = arr.shape[:2]
    if min(H, W) < 2 * STRIP_PX + 4:
        return {'ok': False, 'error': f'image too small ({W}x{H})'}

    lenses = {}
    for lens in LENSES:
        score = lens_score(arr, lens)
        lenses[lens] = {'score': round(score, 2), 'verdict': verdict_for(score)}

    overall_max = max(l['score'] for l in lenses.values())
    overall_verdict = verdict_for(overall_max)

    edges_max = max(lenses[k]['score'] for k in ('top', 'bottom', 'left', 'right'))
    mids_max = max(lenses[k]['score'] for k in ('h_mid', 'v_mid'))

    if overall_verdict == 'PASS':
        summary = 'Clean seamless tile — every lens under the PASS threshold.'
    elif mids_max > WARN_MAX and edges_max <= PASS_MAX:
        summary = (f'Internal-grid defect — midline lens(es) > {WARN_MAX:.0f} '
                   f'while edges are clean. Classic SDXL latent-grid leak.')
    elif edges_max > WARN_MAX and mids_max <= PASS_MAX:
        summary = (f'Edge-wrap defect — edge lens(es) > {WARN_MAX:.0f} while '
                   f'midlines clean. Tile-prep regression (e.g. center-crop on '
                   f'a seamless source).')
    elif overall_verdict == 'FAIL':
        summary = f'Multiple lens failures — both edges and midlines flag.'
    else:
        summary = 'Borderline — some lenses in WARN range.'

    return {
        'ok': True,
        'path': str(png_path),
        'width': W, 'height': H,
        'verdict': overall_verdict,
        'lenses': lenses,
        'edges_max': round(edges_max, 2),
        'mids_max': round(mids_max, 2),
        'summary': summary,
    }


def scan_id(design_id: int) -> dict:
    sql = f"SELECT local_path FROM spoon_all_designs WHERE id={design_id};"
    try:
        r = subprocess.run(
            ['psql', 'dw_unified', '-At', '-c', sql],
            check=True, capture_output=True, text=True, timeout=10,
        )
    except subprocess.CalledProcessError as e:
        # Try sudo-postgres path (Linux prod)
        try:
            r = subprocess.run(
                ['sudo', '-n', '-u', 'postgres', 'psql', 'dw_unified', '-At', '-c', sql],
                check=True, capture_output=True, text=True, timeout=10,
            )
        except Exception:
            return {'ok': False, 'error': f'psql failed: {e.stderr.strip()}'}
    path_str = (r.stdout or '').strip()
    if not path_str:
        return {'ok': False, 'error': f'design {design_id} not found in PG'}
    out = scan_path(Path(path_str))
    out['id'] = design_id
    return out


def render_text(result: dict) -> str:
    if not result.get('ok'):
        return f"ERROR: {result.get('error', 'unknown')}"
    lines = []
    lines.append('')
    lines.append(f"  edges-agent · {result.get('id', '?')} · {result['width']}x{result['height']}")
    lines.append(f"  ────────────────────────────────────")
    for lens, v in result['lenses'].items():
        bar = '█' * min(60, int(v['score'] * 2))
        lines.append(f"  {lens:<7} ΔE={v['score']:>6.2f}  {v['verdict']:<4}  {bar}")
    lines.append(f"  ────────────────────────────────────")
    lines.append(f"  VERDICT: {result['verdict']}")
    lines.append(f"  {result['summary']}")
    lines.append('')
    return '\n'.join(lines)


def main():
    p = argparse.ArgumentParser(description='6-lens seamless-tile defect scanner.')
    p.add_argument('--id', type=int, help='Design id to look up in PG')
    p.add_argument('--path', type=str, help='Direct path to a PNG')
    p.add_argument('--json', action='store_true', help='JSON only (default)')
    p.add_argument('--text', action='store_true', help='Human-readable scorecard')
    args = p.parse_args()

    if args.id:
        result = scan_id(args.id)
    elif args.path:
        result = scan_path(Path(args.path))
    else:
        p.error('Pass --id or --path')

    if args.text:
        print(render_text(result))
    else:
        print(json.dumps(result, indent=2 if args.text else None))


if __name__ == '__main__':
    main()