| """Webcam test: detect+align with buffalo_l detector, embed with TinyFaceNet, match to centroids. |
| Low-light handling: auto-gamma + CLAHE on the frame before detection/crop; per-track majority vote over last VOTE frames. |
| usage: python live_test.py [ckpt] [--raw] (--raw = no normalization, for A/B) |
| q = quit, s = save snapshot to analysis/live_N.jpg, n = toggle normalization""" |
| import sys, time, collections, numpy as np, cv2, torch, torch.nn.functional as F |
| from insightface.app import FaceAnalysis |
| from insightface.utils.face_align import norm_crop |
| from model import TinyFaceNet |
| args=[a for a in sys.argv[1:] if not a.startswith("--")] |
| CK=args[0] if args else "dataset/best.pt"; THR=0.55; VOTE=7; NORM="--raw" not in sys.argv |
| ck=torch.load(CK,weights_only=False); names=ck["names"]; C=ck["centroids"] |
| m=TinyFaceNet(n_cls=len(names)); m.load_state_dict(ck["model"]); m.eval() |
| app=FaceAnalysis(name="buffalo_l",allowed_modules=["detection"],providers=["CPUExecutionProvider"]); app.prepare(ctx_id=-1,det_size=(320,320)) |
| clahe=cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8)) |
| def normalize(img): |
| """auto-gamma to lift dark frames + CLAHE on luminance. Returns (img, mean_brightness_before).""" |
| mean=img.mean() |
| if mean<110: |
| gamma=np.log(110/255)/np.log(max(mean,1)/255); lut=(np.linspace(0,1,256)**gamma*255).astype(np.uint8); img=cv2.LUT(img,lut) |
| lab=cv2.cvtColor(img,cv2.COLOR_BGR2LAB); lab[...,0]=clahe.apply(lab[...,0]) |
| return cv2.cvtColor(lab,cv2.COLOR_LAB2BGR),mean |
| def iou(a,b): |
| ix=max(0,min(a[2],b[2])-max(a[0],b[0])); iy=max(0,min(a[3],b[3])-max(a[1],b[1])); i=ix*iy |
| return i/((a[2]-a[0])*(a[3]-a[1])+(b[2]-b[0])*(b[3]-b[1])-i+1e-6) |
| tracks=[] |
| cap=cv2.VideoCapture(0,cv2.CAP_DSHOW); cap.set(3,640); cap.set(4,480) |
| if not cap.isOpened(): sys.exit("webcam open nahi hua") |
| n=0; t0=time.time(); fr=0 |
| while True: |
| ok,raw=cap.read() |
| if not ok: break |
| img,mean=normalize(raw) if NORM else (raw,raw.mean()) |
| for f in app.get(img): |
| crop=norm_crop(img,f.kps,112); x=(torch.from_numpy(crop).permute(2,0,1)[None].float()/255.-0.5)/0.5 |
| with torch.no_grad(): e=F.normalize(m(x),dim=1)[0] |
| sim=(C@e); j=int(sim.argmax()); s=float(sim[j]); inst=names[j] if s>=THR else "unknown" |
| |
| bb=tuple(map(int,f.bbox)); tr=max(tracks,key=lambda t:iou(t[0],bb),default=None) |
| if tr is None or iou(tr[0],bb)<0.3: tr=[bb,collections.deque(maxlen=VOTE),collections.deque(maxlen=VOTE),fr]; tracks.append(tr) |
| tr[0]=bb; tr[1].append(inst); tr[2].append(s); tr[3]=fr |
| lab,cnt=collections.Counter(tr[1]).most_common(1)[0]; conf=np.mean(tr[2]) |
| col=(0,255,0) if lab!="unknown" else (0,0,255) |
| x1,y1,x2,y2=bb; cv2.rectangle(img,(x1,y1),(x2,y2),col,2) |
| top3=" ".join(f"{names[k]}:{sim[k]:.2f}" for k in sim.argsort(descending=True)[:3]) |
| cv2.putText(img,f"{lab} {conf:.2f} ({cnt}/{len(tr[1])})",(x1,y1-8),0,0.7,col,2); cv2.putText(img,top3,(x1,y2+18),0,0.45,(255,255,0),1) |
| tracks=[t for t in tracks if fr-t[3]<15] |
| fr+=1; cv2.putText(img,f"{fr/(time.time()-t0):.1f}fps norm={'ON' if NORM else 'off'} light={mean:.0f} q=quit s=save n=norm",(8,20),0,0.5,(255,255,255),1) |
| cv2.imshow("live",img); k=cv2.waitKey(1)&0xFF |
| if k==ord("q"): break |
| if k==ord("n"): NORM=not NORM; tracks=[] |
| if k==ord("s"): n+=1; cv2.imwrite(f"analysis/live_{n}.jpg",img); print("saved",n,flush=True) |
| cap.release(); cv2.destroyAllWindows() |
|
|
