verify_rotation.py

"""
verify_rotation.py — scratch check for P1-4 default tile rotation.

Computes the wall-aligned rotation from a bundle (same algorithm as
estimate_default_rotation in app.py), then renders a brick grid with
rotation 0 (left) vs the computed rotation (right) using the exact frontend
rotation math, to confirm the sign and magnitude visually.

Usage:
    python verify_rotation.py /tmp/hf_new_bundle.json
"""

import base64
import io
import json
import sys

import numpy as np
from PIL import Image


def estimate_default_rotation(mask, H, vp2=None):
    h, w = mask.shape[:2]

    def _to_plane(img_pts):
        ones = np.ones((len(img_pts), 1), dtype=np.float64)
        p = np.hstack([img_pts.astype(np.float64), ones]) @ H.T
        zs = p[:, 2]
        valid = np.abs(zs) > 1e-9
        if valid.sum() < 2:
            return None
        return p[valid, 0] / zs[valid], p[valid, 1] / zs[valid]

    pts = []
    for x in range(int(w * 0.2), int(w * 0.8), 4):
        col = np.where(mask[:, x] > 0)[0]
        if len(col) > 10:
            pts.append((float(x), float(col[0])))

    if len(pts) >= 12:
        plane = _to_plane(np.array(pts))
        if plane is not None:
            px, py = plane
            full_span = px.max() - px.min()
            keep = np.ones(len(px), dtype=bool)
            angle = None
            for i in range(3):
                if keep.sum() < 12:
                    break
                sub_x, sub_y = px[keep], py[keep]
                span = sub_x.max() - sub_x.min()
                if span < max(1e-6, full_span * 0.3):
                    break
                slope, intercept = np.polyfit(sub_x, sub_y, 1)
                resid = np.abs(py - (slope * px + intercept))
                rms = float(np.sqrt(np.mean(resid[keep] ** 2)))
                print(f"  pass {i}: n={keep.sum()} span={span:.0f} rms={rms:.1f} "
                      f"({rms / span:.3f} of span) angle={np.degrees(np.arctan(slope)):.1f}")
                if rms < span * 0.04 and keep.sum() >= len(px) * 0.6:
                    angle = float(np.degrees(np.arctan(slope)))
                    break
                keep &= resid <= np.percentile(resid[keep], 70)
            if angle is not None and abs(angle) < 25.0:
                return angle
    return 0.0


def decode_pixels(bundle):
    w, h = bundle["width"], bundle["height"]
    raw = base64.b64decode(bundle["pixels"])
    if len(raw) == w * h * 4:
        return np.frombuffer(raw, np.uint8).reshape(h, w, 4)[:, :, :3].copy()
    return np.array(Image.open(io.BytesIO(raw)).convert("RGB"))


def render(bundle, rotation_deg):
    w, h = bundle["width"], bundle["height"]
    img = decode_pixels(bundle)
    seg = bundle["segments"][0]
    idx = np.frombuffer(base64.b64decode(seg["mask"]), dtype=np.uint32)
    H = np.array(seg["homography"], dtype=np.float64).reshape(3, 3)
    p = seg["plane"]

    ys, xs = idx // w, idx % w
    pw, ph = p["width"], p["height"]
    cx, cy = p["x"] + pw / 2, p["y"] + ph / 2
    repeat = max(32.0, pw * 0.18)

    pts = np.column_stack([xs, ys, np.ones(len(xs))]) @ H.T
    fx = pts[:, 0] / pts[:, 2]
    fy = pts[:, 1] / pts[:, 2]

    # frontend rotation math: rad = rot*pi/180; cos(-rad), sin(-rad)
    rad = rotation_deg * np.pi / 180.0
    c, s = np.cos(-rad), np.sin(-rad)
    dx, dy = fx - cx, fy - cy
    rx = dx * c - dy * s
    ry = dx * s + dy * c

    u = rx / repeat
    v = ry / repeat
    row = np.floor(v).astype(int)
    uu = u + (row % 2) * 0.5
    cell = ((np.floor(uu).astype(int) + row) % 2).astype(bool)
    fu, fv = uu - np.floor(uu), v - np.floor(v)
    grout = (fu < 0.06) | (fv < 0.06)

    color = np.where(cell[:, None], [184, 115, 51], [222, 184, 135]).astype(np.uint8)
    color[grout] = (60, 60, 60)
    out = img.copy()
    out[ys, xs] = (0.75 * color + 0.25 * out[ys, xs]).astype(np.uint8)
    return out


def main():
    bundle_path = sys.argv[1]
    with open(bundle_path) as f:
        bundle = json.load(f)

    w, h = bundle["width"], bundle["height"]
    seg = bundle["segments"][0]
    idx = np.frombuffer(base64.b64decode(seg["mask"]), dtype=np.uint32)
    mask = np.zeros(h * w, np.uint8)
    mask[idx] = 1
    mask = mask.reshape(h, w)
    H = np.array(seg["homography"], dtype=np.float64).reshape(3, 3)

    rot = estimate_default_rotation(mask, H)
    print(f"defaultRotation = {rot:.2f} deg")

    a = render(bundle, 0.0)
    b = render(bundle, rot)
    Image.fromarray(np.hstack([a, b])).save("verify_out/rotation_compare.png")
    print("saved verify_out/rotation_compare.png  (left=rotation 0, right=defaultRotation)")


if __name__ == "__main__":
    main()