animation/characters/work/lena/20_atlas_probe.py

# Stage 20 (read-only probe): characterise the CURRENT UV atlas, so "the texture looks cut up
# and pasted together" becomes a measurement instead of an impression.
#
#   blender --background --python 20_atlas_probe.py -- <in.blend> <out_dir>
#
# Reports, for the mesh's active UV layer:
#   - which image the material actually samples (bpy.data.images holds stale duplicates)
#   - UV-vertex count vs mesh-vertex count  => how much the atlas is cut apart
#   - island count + size distribution      => "pasted together" from how many pieces
#   - per-island texel density (px per mm)  => whether pieces are at inconsistent scale
#   - atlas coverage + wasted area
# Writes: basecolor.png (the real one), islands.png (island map), density.png (px/mm heat),
#         uvgrid.png (UV wireframe), and uv_cache.npz for later stages.
import bpy, sys, os, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND = argv[0]
OUTDIR = os.path.abspath(argv[1])
os.makedirs(OUTDIR, exist_ok=True)
t0 = time.time()


def log(m):
    print(f"[atlas {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)
log(f"mesh '{ob.name}': {n_v}v {n_l}loops {n_f}faces  uv_layers={[l.name for l in me.uv_layers]}")

# ---- which image does the material ACTUALLY sample? follow the node link ----
sampled = {}
for slot in ob.material_slots:
    mat = slot.material
    if not mat or not mat.use_nodes:
        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
            src = node.inputs[sock].links[0].from_node
            seen = set()
            while src and src.type != 'TEX_IMAGE' and id(src) not in seen:
                seen.add(id(src))
                nxt = None
                for i in src.inputs:
                    if i.links:
                        nxt = i.links[0].from_node
                        break
                src = nxt
            if src and src.type == 'TEX_IMAGE' and src.image:
                sampled[key] = src.image
for k, im in sampled.items():
    log(f"SAMPLED {k}: '{im.name}' {im.size[0]}x{im.size[1]} packed={bool(im.packed_file)}")
print("ALL IMAGES IN FILE (duplicates are stale):")
for im in bpy.data.images:
    if im.size[0]:
        print(f"   '{im.name}' {im.size[0]}x{im.size[1]}")

base = sampled.get("base")
if base is None:
    print("!! material samples no basecolor image"); sys.exit(1)
W, H = base.size
buf = np.empty(W * H * 4, dtype=np.float32)
base.pixels.foreach_get(buf)
tex = buf.reshape(H, W, 4)

# ---- UV data ----
loops_v = np.empty(n_l, dtype=np.int32)
me.loops.foreach_get("vertex_index", loops_v)
uv = np.empty(n_l * 2, dtype=np.float64)
me.uv_layers.active.data.foreach_get("uv", uv)
uv = uv.reshape(-1, 2)
co = np.empty(n_v * 3)
me.vertices.foreach_get("co", co)
co = co.reshape(-1, 3)

print(f"UV range u {uv[:,0].min():.4f}..{uv[:,0].max():.4f}  "
      f"v {uv[:,1].min():.4f}..{uv[:,1].max():.4f}")

# ---- uv-vertices: a mesh vertex split across N atlas pieces becomes N uv-vertices ----
Q = 1 << 20
key = (loops_v.astype(np.int64) * Q * Q
       + np.round(np.clip(uv[:, 0], 0, 1) * (Q - 1)).astype(np.int64) * Q
       + np.round(np.clip(uv[:, 1], 0, 1) * (Q - 1)).astype(np.int64))
_, uvv = np.unique(key, return_inverse=True)
n_uvv = uvv.max() + 1
splits = np.bincount(uvv, minlength=n_uvv)
per_vert = np.bincount(loops_v, weights=np.zeros(n_l))  # placeholder
# how many atlas copies does each mesh vertex have?
vk = np.unique(np.stack([loops_v, uvv], axis=1), axis=0)
copies = np.bincount(vk[:, 0], minlength=n_v)
print(f"\n=== CUT-APART ===")
print(f"uv-vertices {n_uvv} for {n_v} mesh vertices  -> {100.0*(n_uvv-n_v)/n_v:+.1f}% duplication")
print(f"vertices on a UV seam: {int((copies > 1).sum())} ({100.0*(copies>1).sum()/n_v:.1f}%)  "
      f"max copies {int(copies.max())}")

# ---- islands = connected components of the uv-mesh ----
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)
tri = l_tot == 3
log(f"faces: {int(tri.sum())} tris, {int((~tri).sum())} n-gons")
li = l_start[tri]
T = np.stack([uvv[li], uvv[li + 1], uvv[li + 2]], axis=1)

parent = np.arange(n_uvv, dtype=np.int64)


def find(x):
    r = x
    while parent[r] != r:
        r = parent[r]
    while parent[x] != r:
        parent[x], x = r, parent[x]
    return r


for a, b, c in T:
    ra, rb, rc = find(a), find(b), find(c)
    if ra != rb:
        parent[rb] = ra
    if ra != rc:
        parent[rc] = ra
log("union-find done")
roots = np.array([find(i) for i in range(n_uvv)])
_, isl = np.unique(roots, return_inverse=True)
n_isl = isl.max() + 1

# per-island geometry: UV area and 3D area
Puv = np.clip(uv, 0, 1)
tri_uv = np.stack([Puv[li], Puv[li + 1], Puv[li + 2]], axis=1)          # (F,3,2)
auv = 0.5 * np.abs((tri_uv[:, 1, 0] - tri_uv[:, 0, 0]) * (tri_uv[:, 2, 1] - tri_uv[:, 0, 1])
                   - (tri_uv[:, 2, 0] - tri_uv[:, 0, 0]) * (tri_uv[:, 1, 1] - tri_uv[:, 0, 1]))
P3 = co[np.stack([loops_v[li], loops_v[li + 1], loops_v[li + 2]], axis=1)]  # (F,3,3)
cr = np.cross(P3[:, 1] - P3[:, 0], P3[:, 2] - P3[:, 0])
a3 = 0.5 * np.linalg.norm(cr, axis=1)
fisl = isl[T[:, 0]]
isl_auv = np.bincount(fisl, weights=auv, minlength=n_isl)
isl_a3 = np.bincount(fisl, weights=a3, minlength=n_isl)
isl_nf = np.bincount(fisl, minlength=n_isl)

order = np.argsort(-isl_auv)
UNIT = 1.815      # 1 mesh unit = 1.815 m (body is 0.979 units for 1.777 m)
print(f"\n=== PASTED TOGETHER ===")
print(f"islands: {n_isl}")
print(f"atlas UV area used: {isl_auv.sum()*100:.1f}%  (rest is padding/waste)")
cum = np.cumsum(isl_auv[order]) / max(isl_auv.sum(), 1e-12)
for frac in (0.5, 0.9, 0.99):
    print(f"  {int(np.searchsorted(cum, frac))+1} islands cover {frac*100:.0f}% of the used area")
tiny = int((isl_nf < 20).sum())
print(f"  islands with <20 faces: {tiny} ({100.0*tiny/n_isl:.1f}%)")
print("\ntop 20 islands (px/mm = texel density at 4096):")
print("  #     faces    uv_area%   3D area cm2   px/mm   uv centre")
for i in order[:20]:
    if isl_a3[i] <= 0:
        continue
    dens = np.sqrt(isl_auv[i] / isl_a3[i]) * W / (UNIT * 1000.0)
    m = fisl == i
    cu = tri_uv[m].reshape(-1, 2).mean(axis=0)
    print(f"  {i:6d} {isl_nf[i]:7d}  {isl_auv[i]*100:8.3f}   "
          f"{isl_a3[i]*UNIT*UNIT*1e4:10.1f}   {dens:6.2f}   ({cu[0]:.3f},{cu[1]:.3f})")

big = order[:max(1, int(np.searchsorted(cum, 0.99)) + 1)]
dens_all = np.where(isl_a3 > 0, np.sqrt(np.maximum(isl_auv, 0) / np.maximum(isl_a3, 1e-12))
                    * W / (UNIT * 1000.0), np.nan)
d = dens_all[big]
d = d[np.isfinite(d)]
print(f"\ntexel density over the 99%-area islands: min {d.min():.2f} median {np.median(d):.2f} "
      f"max {d.max():.2f} px/mm  -> {d.max()/max(d.min(),1e-9):.1f}x spread")

# ---- pictures ----
def save(arr, name):
    h, w = arr.shape[:2]
    img = bpy.data.images.new(name, w, h, alpha=False, float_buffer=False)
    a = np.ones((h, w, 4), dtype=np.float32)
    a[:, :, :3] = arr.astype(np.float32)
    img.pixels.foreach_set(a.reshape(-1))
    p = os.path.join(OUTDIR, name + ".png")
    img.file_format = 'PNG'
    img.filepath_raw = p
    img.save(filepath=p)
    log(f"wrote {p} exists={os.path.exists(p)}")


save(tex[:, :, :3], "basecolor")

R = 1024
sc = R / float(W)
rng = np.random.RandomState(3)
pal = rng.rand(n_isl, 3) * 0.75 + 0.2
IS = np.zeros((R, R, 3), dtype=np.float64)
DN = np.zeros((R, R), dtype=np.float64)
GR = np.zeros((R, R), dtype=np.float64)
tri_px = tri_uv * np.array([(R - 1), (R - 1)])
for fi in range(len(tri_px)):
    P = tri_px[fi]
    x0, x1 = int(P[:, 0].min()), int(np.ceil(P[:, 0].max()))
    y0, y1 = int(P[:, 1].min()), int(np.ceil(P[:, 1].max()))
    if x1 < x0 or y1 < y0 or x1 - x0 > 64 or y1 - y0 > 64:
        continue
    dd = ((P[1, 1] - P[2, 1]) * (P[0, 0] - P[2, 0]) + (P[2, 0] - P[1, 0]) * (P[0, 1] - P[2, 1]))
    if abs(dd) < 1e-12:
        continue
    gx, gy = np.meshgrid(np.arange(x0, min(x1, R - 1) + 1), np.arange(y0, min(y1, R - 1) + 1))
    aa = ((P[1, 1] - P[2, 1]) * (gx - P[2, 0]) + (P[2, 0] - P[1, 0]) * (gy - P[2, 1])) / dd
    bb = ((P[2, 1] - P[0, 1]) * (gx - P[2, 0]) + (P[0, 0] - P[2, 0]) * (gy - P[2, 1])) / dd
    cc = 1.0 - aa - bb
    ins = (aa >= 0) & (bb >= 0) & (cc >= 0)
    if not ins.any():
        continue
    IS[gy[ins], gx[ins]] = pal[fisl[fi]]
    DN[gy[ins], gx[ins]] = dens_all[fisl[fi]] if np.isfinite(dens_all[fisl[fi]]) else 0
    # edge pixels -> wireframe
    ed = ins & ((aa < 0.06) | (bb < 0.06) | (cc < 0.06))
    GR[gy[ed], gx[ed]] = 1.0
log("rasterised island map")
save(IS, "islands")
dv = DN / max(np.percentile(DN[DN > 0], 98), 1e-9)
save(np.stack([np.clip(dv, 0, 1), np.clip(1 - np.abs(dv - 0.5) * 2, 0, 1),
               np.clip(1 - dv, 0, 1)], axis=2), "density")
tsm = tex[::W // R, ::W // R, :3]
save(np.clip(tsm * (1 - GR[:, :, None] * 0.8) + GR[:, :, None] * np.array([0.0, 1.0, 0.2]), 0, 1),
     "uvgrid")

np.savez_compressed(os.path.join(OUTDIR, "uv_cache.npz"),
                    isl=isl, uvv=uvv, fisl=fisl, isl_auv=isl_auv, isl_a3=isl_a3,
                    isl_nf=isl_nf, copies=copies)
print("ATLAS_PROBE_DONE")