File size: 3,561 Bytes
902df3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """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") # flat folder of jpg/mp4, filenames YYYYMMDD_*; OUT="dataset"
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)
# ---- jpg faces ----
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) # scale used in scan
s1=im.width/W # scale after draft
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)
# ---- centroids for auto-label ----
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)
# ---- video frames @1fps ----
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 # ambiguous, skip
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))
|