animation/characters/work/lena/54_crotch_refill.py

# Stage 54 (v02): kill the rusty crotch patch on the v02 atlas — the focused re-run of what
# 36_refill.py did for v01.
#
#   blender --background --python 54_crotch_refill.py -- <in.blend> <out.blend>
#
# Why not just re-run 36: it rebuilds its mask from masks.npz + 00_welded.blend, and the welded
# reference didn't survive the prune. It's also not needed here — the v02 body texture comes from
# the 34_v04 master whose garment regions were already patch-filled clean; the ONLY colour
# regression is the donor-crotch rust, and that region is a geometric box (the same box 36 used:
# |x| < 0.075, 0.340 < z < 0.480, in the 0.98-unit body frame).
#
# Method is 36's, unchanged in the ways that mattered:
#   - the fill tone is a HARMONIC solve on the mesh with Dirichlet boundaries (CG, not Jacobi) —
#     equal to her real skin at the mask edge by construction, and seam-proof across the
#     torso/leg chart border the crotch straddles;
#   - grain transplanted from clean skin tiles so it is not a decal;
#   - normal flattened and rm set to surrounding-skin median over the same texels, which is what
#     keeps the region featureless.
# Texture-only: works directly on the RIGGED master, no re-rig needed.
import bpy, sys, os, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
GROW = 3
GRAIN_T = 16
FEATHER = 4
t0 = time.time()


def log(m):
    print(f"[cr54 {time.time()-t0:6.1f}s] {m}", flush=True)


bpy.ops.wm.open_mainfile(filepath=BLEND)
ob = max([o for o in bpy.data.objects if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
me = ob.data
n_v, n_l, n_f = len(me.vertices), len(me.loops), len(me.polygons)
co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
z = co[:, 2] - co[:, 2].min()
log(f"{n_v}v {n_f}f")

# ---- the mask: 36's crotch box, slightly extended down the inner thigh, + grow rings ----
mask = (np.abs(co[:, 0]) < 0.075) & (z > 0.320) & (z < 0.480)
ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
for _ in range(GROW):
    hit = mask[ev[:, 0]] | mask[ev[:, 1]]
    mask[ev[hit, 0]] = True
    mask[ev[hit, 1]] = True
log(f"crotch mask: {int(mask.sum())} verts")

# ---- images through the material graph ----
src = {}
for slot in ob.material_slots:
    mat = slot.material
    if not mat or not mat.node_tree:
        continue
    for node in mat.node_tree.nodes:
        if node.type != 'BSDF_PRINCIPLED':
            continue
        for sock, key in (("Base Color", "base"), ("Normal", "normal"), ("Roughness", "rm")):
            if sock not in node.inputs or not node.inputs[sock].links:
                continue
            nd = node.inputs[sock].links[0].from_node
            seen = set()
            while nd and nd.type != 'TEX_IMAGE' and id(nd) not in seen:
                seen.add(id(nd))
                nxt = None
                for i in nd.inputs:
                    if i.links:
                        nxt = i.links[0].from_node
                        break
                nd = nxt
            if nd and nd.type == 'TEX_IMAGE' and nd.image:
                src[key] = nd.image
base = src["base"]
W, H = base.size
buf = np.empty(W * H * 4, dtype=np.float32)
base.pixels.foreach_get(buf)
tex = buf.reshape(H, W, 4)
rgb = tex[:, :, :3].astype(np.float64)
log(f"atlas '{base.name}' {W}x{H}")

loops_v = np.empty(n_l, dtype=np.int32); me.loops.foreach_get("vertex_index", loops_v)
uv = np.empty(n_l * 2); me.uv_layers.active.data.foreach_get("uv", uv); uv = uv.reshape(-1, 2)
l_start = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_start", l_start)
l_tot = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_total", l_tot)
li = l_start[l_tot == 3]

px = np.clip(np.round(uv[:, 0] * (W - 1)).astype(np.int32), 0, W - 1)
py = np.clip(np.round(uv[:, 1] * (H - 1)).astype(np.int32), 0, H - 1)
lc = rgb[py, px]
vc = np.zeros((n_v, 3))
for c in range(3):
    vc[:, c] = np.bincount(loops_v, weights=lc[:, c], minlength=n_v)
vn = np.maximum(np.bincount(loops_v, minlength=n_v), 1)
vc /= vn[:, None]

# ---- harmonic fill (CG on the graph Laplacian, Dirichlet boundary = her real skin) ----
o_ = np.concatenate([ev[:, 0], ev[:, 1]])
n_ = np.concatenate([ev[:, 1], ev[:, 0]])
deg = np.maximum(np.bincount(o_, minlength=n_v).astype(np.float64), 1.0)
mf = mask.astype(np.float64)


def A_mul(x):
    xm = x * mf
    return (deg * xm - np.bincount(o_, weights=xm[n_], minlength=n_v)) * mf


fill = vc.copy()
for c in range(3):
    known = vc[:, c] * (1.0 - mf)
    b = np.bincount(o_, weights=known[n_], minlength=n_v) * mf
    x = np.zeros(n_v)
    r = b - A_mul(x)
    p = r.copy()
    rs = float(r @ r); r0 = rs
    for it in range(4000):
        if rs <= max(r0 * 1e-12, 1e-20):
            break
        Ap = A_mul(p)
        d_ = float(p @ Ap)
        if abs(d_) < 1e-30:
            break
        al = rs / d_
        x += al * p; r -= al * Ap
        rs2 = float(r @ r)
        p = r + (rs2 / rs) * p
        rs = rs2
    fill[mask, c] = x[mask]
    log(f"  ch{c}: CG {it+1} iters, residual {np.sqrt(rs/max(r0,1e-30)):.2e}")
bnd = mask & (np.bincount(o_, weights=(1.0 - mf)[n_], minlength=n_v) > 0)
step = np.abs(fill[bnd] - vc[bnd]).max(axis=1)
log(f"boundary agreement: mean {step.mean():.4f} p99 {np.percentile(step,99):.4f}")

# ---- rasterise masked faces ----
IDX = np.stack([li, li + 1, li + 2], axis=1)
V = loops_v[IDX]
face_any = mask[V].any(axis=1)
P = np.stack([uv[IDX][:, :, 0] * (W - 1), uv[IDX][:, :, 1] * (H - 1)], axis=2)
out = rgb.copy()
paint = np.zeros((H, W), dtype=bool)
for f in np.nonzero(face_any)[0]:
    p3 = P[f]
    x0, x1 = int(p3[:, 0].min()), int(np.ceil(p3[:, 0].max()))
    y0, y1 = int(p3[:, 1].min()), int(np.ceil(p3[:, 1].max()))
    if x1 < x0 or y1 < y0 or x1 - x0 > 512 or y1 - y0 > 512:
        continue
    det = ((p3[1, 1] - p3[2, 1]) * (p3[0, 0] - p3[2, 0])
           + (p3[2, 0] - p3[1, 0]) * (p3[0, 1] - p3[2, 1]))
    if abs(det) < 1e-12:
        continue
    gx, gy = np.meshgrid(np.arange(max(x0, 0), min(x1, W - 1) + 1),
                         np.arange(max(y0, 0), min(y1, H - 1) + 1))
    if gx.size == 0:
        continue
    a = ((p3[1, 1] - p3[2, 1]) * (gx - p3[2, 0]) + (p3[2, 0] - p3[1, 0]) * (gy - p3[2, 1])) / det
    b_ = ((p3[2, 1] - p3[0, 1]) * (gx - p3[2, 0]) + (p3[0, 0] - p3[2, 0]) * (gy - p3[2, 1])) / det
    c_ = 1.0 - a - b_
    ins = (a >= -0.02) & (b_ >= -0.02) & (c_ >= -0.02)
    if not ins.any():
        continue
    aa, bb, cc = a[ins], b_[ins], c_[ins]
    w = aa * mf[V[f, 0]] + bb * mf[V[f, 1]] + cc * mf[V[f, 2]]
    col = (aa[:, None] * fill[V[f, 0]] + bb[:, None] * fill[V[f, 1]] + cc[:, None] * fill[V[f, 2]])
    yy, xx = gy[ins], gx[ins]
    hard = w > 0.5
    if hard.any():
        out[yy[hard], xx[hard]] = col[hard]
        paint[yy[hard], xx[hard]] = True
log(f"repainted {int(paint.sum())} texels ({100.0*paint.mean():.3f}%)")


def box(a, r):
    def b1(v, ax):
        pad = [(0, 0)] * v.ndim
        pad[ax] = (r, r)
        cs = np.cumsum(np.pad(v, pad, mode="edge"), axis=ax)
        return (np.take(cs, np.arange(2 * r, cs.shape[ax]), axis=ax)
                - np.take(cs, np.arange(0, cs.shape[ax] - 2 * r), axis=ax)) / (2 * r)
    return b1(b1(a, 0), 1)


# grain transplant
grain = np.stack([rgb[:, :, c] - box(rgb[:, :, c], 5) for c in range(3)], axis=2)
tone = np.median(out[paint], axis=0)
cand = []
for ty in range(0, H - GRAIN_T, GRAIN_T):
    for tx in range(0, W - GRAIN_T, GRAIN_T):
        if paint[ty:ty + GRAIN_T, tx:tx + GRAIN_T].any():
            continue
        t = rgb[ty:ty + GRAIN_T, tx:tx + GRAIN_T].reshape(-1, 3)
        if t.min() < 0.02:
            continue
        if np.abs(t.mean(axis=0) - tone).max() < 0.10:
            cand.append((ty, tx))
rng = np.random.RandomState(11)
if cand:
    for ty in range(0, H - GRAIN_T + 1, GRAIN_T):
        for tx in range(0, W - GRAIN_T + 1, GRAIN_T):
            tm = paint[ty:ty + GRAIN_T, tx:tx + GRAIN_T]
            if not tm.any():
                continue
            sy, sx = cand[rng.randint(len(cand))]
            out[ty:ty + GRAIN_T, tx:tx + GRAIN_T][tm] += \
                grain[sy:sy + GRAIN_T, sx:sx + GRAIN_T][tm] * 0.85
log(f"grain from {len(cand)} tiles")

# feather rim
a_ = np.ones((H, W))
edge = paint.copy()
for k in range(FEATHER):
    grown = edge.copy()
    grown[1:-1, 1:-1] |= (edge[:-2, 1:-1] | edge[2:, 1:-1] | edge[1:-1, :-2] | edge[1:-1, 2:])
    ring = grown & ~edge
    a_[ring] = (k + 1) / (FEATHER + 1.0)
    edge = grown
blend = np.where(paint, 1.0, 1.0 - a_)[:, :, None]
final = np.clip(out * blend + rgb * (1 - blend), 0, 1)
b4 = tex.copy()
b4[:, :, :3] = final.astype(np.float32)
base.pixels.foreach_set(b4.reshape(-1))
base.pack()

# normal flat + rm to surrounding median over the same texels
def dil(m, k):
    g = m.copy()
    for _ in range(k):
        n2 = g.copy()
        n2[1:, :] |= g[:-1, :]; n2[:-1, :] |= g[1:, :]
        n2[:, 1:] |= g[:, :-1]; n2[:, :-1] |= g[:, 1:]
        g = n2
    return g


soft = dil(paint, 2)
for key in ("normal", "rm"):
    if key not in src:
        continue
    im = src[key]
    if tuple(im.size) != (W, H):
        continue
    a2 = np.empty(W * H * 4, dtype=np.float32)
    im.pixels.foreach_get(a2)
    arr = a2.reshape(H, W, 4)
    if key == "normal":
        arr[soft, 0] = 0.5; arr[soft, 1] = 0.5; arr[soft, 2] = 1.0
    else:
        med = np.median(arr[~dil(paint, 8)][:, :3], axis=0)
        arr[soft, 0] = med[0]; arr[soft, 1] = med[1]; arr[soft, 2] = med[2]
    im.pixels.foreach_set(arr.reshape(-1))
    im.pack()
    log(f"  {key}: {int(soft.sum())} texels neutralised")

bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("CR54_DONE")