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