| """Enroll identities: build per-person centroids from folders of aligned 112x112 face crops. |
| Layout: <root>/PERSON_A/*.jpg <root>/PERSON_B/*.jpg ... (any folder names; they become the labels) |
| usage: python enroll.py <root> [tinyfacenet_v1.pt] [centroids.npz] |
| Then at inference: emb = normalize(model(x)); sim = centroids @ emb; label = names[argmax] if max >= THR (0.45) else "unknown". |
| Crops must be 5-point aligned (e.g. insightface norm_crop) and RGB.""" |
| import sys, os, glob, numpy as np, torch, torch.nn.functional as F |
| from PIL import Image |
| from model import TinyFaceNet |
| root=sys.argv[1]; CK=sys.argv[2] if len(sys.argv)>2 else "tinyfacenet_v1.pt"; OUT=sys.argv[3] if len(sys.argv)>3 else "centroids.npz" |
| m=TinyFaceNet(); m.load_state_dict(torch.load(CK,weights_only=True)["model"],strict=False); m.eval() |
| def prep(paths): |
| x=np.stack([np.asarray(Image.open(p).convert("RGB").resize((112,112))) for p in paths]).astype(np.float32) |
| return (torch.from_numpy(x).permute(0,3,1,2)/255.-0.5)/0.5 |
| names,cents=[],[] |
| for d in sorted(os.listdir(root)): |
| ps=sorted(glob.glob(os.path.join(root,d,"*.jpg"))+glob.glob(os.path.join(root,d,"*.png"))) |
| if not ps: continue |
| with torch.no_grad(): E=torch.cat([F.normalize(m(prep(ps[i:i+64])),dim=1) for i in range(0,len(ps),64)]) |
| names.append(d); cents.append(F.normalize(E.mean(0),dim=0).numpy()); print(d,len(ps),"faces") |
| np.savez(OUT,names=np.array(names),centroids=np.stack(cents)); print("wrote",OUT,len(names),"identities") |
|
|
