Facepoints 1.0.0: code, model, prep data and card

88607d8 verified 7 days ago

7.4 kB

"""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

	"""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