"""Validation: the 620 points against mesh labels (px, mm), the side heads against hand labels.""" # Each function returns one report dict per call; the trainer appends it to val.json and prints the # VAL lines. Both write a picture grid through draw_grid(). import os import cv2 import numpy as np import torch from facepoints.model import LOG_SCALE_RANGE, MEAN, STD def autocast(dev, enabled): """Return a float16 autocast context for device dev (a no-op on the CPU or when disabled).""" kind = dev.split(":")[0] return torch.autocast(kind, dtype=torch.float16, enabled=bool(enabled) and kind == "cuda") def mm_scale(ds, ks): """Return (n, P) px-to-mm factors at label depth for items ks (NaN without a focal length).""" S = [] for k in ks: _, i, z = ds.meta(k) side = float(z["box"][i][2] - z["box"][i][0]) fx = float(z["fx"][i]) if "fx" in z else np.nan S.append(side / ds.res * z["z"][i] / fx) return np.array(S) def draw_grid(shown, res, path, scale_every=0): """Write up to 8 panels of labels (green) and predictions (red), with scale circles if asked.""" T = res sheet = np.zeros((2 * T, 4 * T, 3), np.uint8) for j, (x, uv, ok, o) in enumerate(shown[:8]): im = ((x.transpose(1, 2, 0) * STD + MEAN) * 255).clip(0, 255)[:, :, ::-1] im = np.ascontiguousarray(im.astype(np.uint8)) for p in np.flatnonzero(ok): cv2.circle(im, (int(uv[p, 0]), int(uv[p, 1])), 1, (0, 255, 0), -1) cv2.circle(im, (int(o[p, 0]), int(o[p, 1])), 1, (0, 0, 255), -1) if scale_every and p % scale_every == 0: rad = int(min(np.exp(np.clip(o[p, 2], *LOG_SCALE_RANGE)), T / 4)) cv2.circle(im, (int(o[p, 0]), int(o[p, 1])), rad, (0, 0, 255), 1) sheet[(j // 4) * T : (j // 4 + 1) * T, (j % 4) * T : (j % 4 + 1) * T] = im cv2.imwrite(path, sheet, [cv2.IMWRITE_JPEG_QUALITY, 88]) def evaluate(net, ds, cfg, dev, cls, region, out_dir, it): """Return per class and region the median error (px, mm) over visible and over all labelled.""" net.eval() dl = torch.utils.data.DataLoader(ds, cfg.batch, num_workers=min(cfg.workers, 4)) E, OK, VIS, MM, SC, shown = [], [], [], [], [], [] with torch.no_grad(): for x, uv, ok, vis, ks in dl: with autocast(dev, cfg.amp): o = net(x.to(dev, non_blocking=True)).float().cpu() E.append((o[..., :2] - uv).norm(dim=-1).numpy()) OK.append(ok.numpy()) VIS.append(vis.numpy()) MM.append(mm_scale(ds, ks.tolist())) SC.append(torch.exp(o[..., 2].clamp(*LOG_SCALE_RANGE)).numpy()) if len(shown) < 8: for j in range(min(8 - len(shown), len(x))): shown.append((x[j].numpy(), uv[j].numpy(), ok[j].numpy(), o[j].numpy())) E, OK, VIS, MM, SC = (np.concatenate(a) for a in (E, OK, VIS, MM, SC)) MMe = E * MM groups = {"all": np.ones(len(cls), bool)} for c in dict.fromkeys(cls): groups[str(c)] = cls == c for r in dict.fromkeys(region): if r: groups[str(r)] = region == r rep = {"it": it, "n_images": int(len(E))} for g, m in groups.items(): sel_v = OK & VIS & m[None] sel_a = OK & m[None] if sel_a.sum() == 0: continue rep[g] = { "vis_px": float(np.median(E[sel_v])) if sel_v.any() else None, "vis_mm": float(np.nanmedian(MMe[sel_v])) if sel_v.any() else None, "all_px": float(np.median(E[sel_a])), "all_mm": float(np.nanmedian(MMe[sel_a])), "sigma_vis_px": float(np.median(SC[sel_v])) if sel_v.any() else None, "n_vis": int(sel_v.sum()), "n_all": int(sel_a.sum()), } draw_grid(shown, ds.res, os.path.join(out_dir, f"pred_grid_{it}.jpg"), scale_every=8) net.train() return rep def evaluate_wild(net, side, ds, cfg, dev, out_dir, it, side_rf=None, mp=None): """Return per wild set the side head's NME and FR@10 (refined too) and the 620 head's terms.""" # NME = mean point error / outer-eye-corner distance x 100; FR@10 = share of faces above 10. # With mp (losses.Map620): medians of its four terms per face, the under-chin points' median # signed distance to the jaw (% IOD, + = below) and, where pupils are labelled, the iris error. net.eval() dl = torch.utils.data.DataLoader(ds, cfg.batch, num_workers=min(cfg.workers, 4)) E = {n: [] for n in ds.names} ER = {n: [] for n in ds.names} T = {n: {k: [] for k in ("map", "jaw", "chin", "pupil", "dense", "uc")} for n in ds.names} shown = [] with torch.no_grad(): for x, uv, _, j, _ in dl: with autocast(dev, cfg.amp): _, zt, fs = net.run(x.to(dev, non_blocking=True), 0) for k, n in enumerate(ds.names): sel = j == k if not bool(sel.any()): continue K, ey = ds.K[k], ds.lab[n]["eyes"] sd = sel.to(dev) oc = side[n](zt[sd].float()).view(-1, K, 3) o = oc.cpu() g = uv[sel, :K] io = (g[:, int(ey[0])] - g[:, int(ey[1])]).norm(dim=-1) E[n].append(((o[..., :2] - g).norm(dim=-1).mean(1) / io * 100).numpy()) if side_rf is not None and len(side_rf): r = side_rf[n]([f[sd].float() for f in fs], oc).cpu() ER[n].append(((r[..., :2] - g).norm(dim=-1).mean(1) / io * 100).numpy()) if mp is not None: mu6 = net.h[-1](zt[sd].float()).view(-1, net.n_points, 3)[..., :2] gd = g.to(dev) okd = torch.isfinite(gd).all(-1) tmap, tjaw, tdense, tuc = mp.terms(n, mu6, gd, okd, per_face=True) chin, pupil = mp.probe(n, mu6, gd, okd) terms = (("map", tmap * mp.B0), ("jaw", tjaw * mp.B0), ("chin", chin), ("pupil", pupil), ("dense", tdense), ("uc", tuc * mp.B0)) # fmt: skip for key, v in terms: T[n][key].append(v.cpu().numpy().astype(np.float64)) for q in range(min(int(sel.sum()), 4 - sum(s[0] == n for s in shown))): shown.append((n, x[sel][q].numpy(), g[q].numpy(), o[q].numpy())) rep = {} for n, e in E.items(): if not e: continue e = np.concatenate(e) rep[n] = {"nme": float(e.mean()), "fr10": float((e > 10).mean() * 100), "n": len(e)} if ER[n]: rep[n]["nme_rf"] = float(np.concatenate(ER[n]).mean()) if T[n]["map"]: t = {k: np.concatenate(v) for k, v in T[n].items()} rep[n].update( m620=float(np.median(t["map"])), j620=float(np.nanmedian(t["jaw"])), chin620=float(np.nanmedian(t["chin"])), d620=float(np.nanmedian(t["dense"])), u620=float(np.nanmedian(t["uc"])), ) if np.isfinite(t["pupil"]).any(): rep[n]["pupil620"] = float(np.nanmedian(t["pupil"])) panels = [(x, g, np.ones(len(g), bool), o) for _, x, g, o in shown[:8]] draw_grid(panels, ds.res, os.path.join(out_dir, f"wild_grid_{it}.jpg")) net.train() return rep