File size: 3,587 Bytes
902df3d e0c1479 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 | """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: # dark frame -> gamma lift towards 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=[] # [bbox, deque(labels), deque(scores), last_seen_frame]
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"
# --- track + vote ---
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()
|