| """Build aligned 112x112 face dataset with teacher (buffalo_l) embeddings. |
| Sources: labeled jpg faces (from clustering) + video frames (auto-labeled by centroid).""" |
| import os, sys, json, re, numpy as np, cv2, collections |
| from PIL import Image |
| from insightface.app import FaceAnalysis |
| from insightface.utils.face_align import norm_crop |
| D=os.environ.get("FACE_DATA","data") |
| labels=json.load(open("analysis/labels.json")); names=sorted(set(labels.values()))+["other"] |
| rows=json.load(open("analysis/faces_clustered.json")) |
| app=FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"]); app.prepare(ctx_id=-1, det_size=(320,320)) |
| def date_of(f): |
| m=re.match(r"(\d{8})_",f); return m.group(1) if m else "00000000" |
| imgs,embs,labs,srcs,dates=[],[],[],[],[] |
| def add(face,img,lab,src,date): |
| crop=norm_crop(img,face.kps,112); imgs.append(crop); embs.append(face.normed_embedding.astype(np.float32)); labs.append(lab); srcs.append(src); dates.append(date) |
| |
| byfile=collections.defaultdict(list) |
| for r in rows: byfile[r["file"]].append(r) |
| other_cap=400; other_n=0 |
| for i,(f,rs) in enumerate(sorted(byfile.items())): |
| im=Image.open(os.path.join(D,f)); im.draft("RGB",(1280,1280)); im=im.convert("RGB") |
| W,H=rs[0]["w"],rs[0]["h"]; s0=1280/max(W,H); s0=min(s0,1) |
| s1=im.width/W |
| img=cv2.cvtColor(np.array(im),cv2.COLOR_RGB2BGR) |
| for r in rs: |
| lab=labels.get(str(r["cluster"]),"other") |
| if lab=="other": |
| if other_n>=other_cap: continue |
| x1,y1,x2,y2=[v/s0*s1 for v in r["box"]]; m=0.6*max(x2-x1,y2-y1) |
| X1,Y1,X2,Y2=int(max(0,x1-m)),int(max(0,y1-m)),int(min(img.shape[1],x2+m)),int(min(img.shape[0],y2+m)) |
| sub=img[Y1:Y2,X1:X2] |
| if sub.size==0: continue |
| fs=app.get(sub) |
| if not fs: continue |
| fc=max(fs,key=lambda q:(q.bbox[2]-q.bbox[0])*(q.bbox[3]-q.bbox[1])) |
| add(fc,sub,lab,"jpg:"+f,date_of(f)) |
| if lab=="other": other_n+=1 |
| if i%50==0: print("jpg",i,len(byfile),len(imgs),flush=True) |
| print("jpg done",len(imgs),collections.Counter(labs),flush=True) |
| |
| E=np.array(embs); L=np.array(labs) |
| cent={n:E[L==n].mean(0) for n in names if n!="other" and (L==n).any()} |
| for n in cent: cent[n]/=np.linalg.norm(cent[n]) |
| C=np.stack([cent[n] for n in cent]); cn=list(cent) |
| |
| vid_counts=collections.Counter() |
| for f in sorted(x for x in os.listdir(D) if x.lower().endswith(".mp4")): |
| cap=cv2.VideoCapture(os.path.join(D,f)); fps=cap.get(5) or 30; n=int(cap.get(7)); step=int(round(fps)) |
| for fi in range(0,n,step): |
| cap.set(1,fi); ok,img=cap.read() |
| if not ok: continue |
| h,w=img.shape[:2]; s=960/max(h,w) |
| if s<1: img=cv2.resize(img,None,fx=s,fy=s) |
| for fc in app.get(img): |
| if max(fc.bbox[2]-fc.bbox[0],fc.bbox[3]-fc.bbox[1])<40: continue |
| sim=C@fc.normed_embedding; j=int(sim.argmax()) |
| if sim[j]>=0.45: lab=cn[j] |
| elif sim[j]<0.30: lab="other" |
| else: continue |
| add(fc,img,lab,f"vid:{f}:{fi}",date_of(f)); vid_counts[lab]+=1 |
| print("vid",f,dict(vid_counts),flush=True) |
| np.savez_compressed(os.path.join(OUT,"faces.npz"),imgs=np.array(imgs,np.uint8),embs=np.array(embs,np.float32),labels=np.array(labs),src=np.array(srcs),date=np.array(dates)) |
| print("DONE",len(imgs),collections.Counter(labs)) |
|
|
