"""Stereo-constrained WBF: cross-camera consistency from East shed pairs. Learns positional mapping EastLeft<->EastRight from 197 stereo pairs, then uses it to weight predictions during WBF. """ import sys,os,json,gc,numpy as np from PIL import Image from tqdm import tqdm from collections import defaultdict PROJECT_DIR='/home/user/goat' os.chdir(PROJECT_DIR);sys.path.insert(0,PROJECT_DIR) from ultralytics import YOLO import torch def iou_fn(b1,b2): x1,y1=max(b1[0],b2[0]),max(b1[1],b2[1]) x2,y2=min(b1[2],b2[2]),min(b1[3],b2[3]) inter=max(0,x2-x1)*max(0,y2-y1) a1=(b1[2]-b1[0])*(b1[3]-b1[1]);a2=(b2[2]-b2[0])*(b2[3]-b2[1]) return inter/(a1+a2-inter+1e-8) def learn_stereo_prior(): """Learn size ratios and positional offsets between EastLeft/EastRight.""" img_dir='Data/Detection_dataset/images/train' lbl_dir='Data/Detection_dataset/labels/train' # Find matched pairs pairs={} files=sorted(os.listdir(img_dir)) for f in files: if not f.endswith('.jpg'): continue parts=f.split('_',1) if len(parts)<2: continue cam,ts=parts[0],parts[1] if cam in ['EastLeft','EastRight']: pairs.setdefault(ts,{})[cam]=f # Collect size/position stats l_sizes,r_sizes=[],[] l_centers,r_centers=[],[] for ts,data in pairs.items(): if 'EastLeft' not in data or 'EastRight' not in data: continue lf,rf=data['EastLeft'],data['EastRight'] for side,fname in [('L',lf),('R',rf)]: lp=os.path.join(lbl_dir,fname.replace('.jpg','.txt')) if not os.path.exists(lp): continue with open(lp) as fh: for line in fh: p=line.strip().split() if len(p)>=5: cx,cy,w,h=float(p[1]),float(p[2]),float(p[3]),float(p[4]) if side=='L': l_sizes.append(np.sqrt(w*h)) l_centers.append([cx,cy]) else: r_sizes.append(np.sqrt(w*h)) r_centers.append([cx,cy]) if len(l_sizes)<10: return None # Compute stats prior={ 'size_ratio': float(np.mean(r_sizes)/np.mean(l_sizes)) if l_sizes else 1.0, 'cx_offset': float(np.mean([r[0] for r in r_centers])-np.mean([l[0] for l in l_centers])), 'cy_mean_diff': float(np.mean([r[1] for r in r_centers])-np.mean([l[1] for l in l_centers])), } print(f'Stereo prior: size_ratio={prior["size_ratio"]:.3f}, cx_offset={prior["cx_offset"]:.4f}') return prior def wbf_fn(bl,sl,thr=0.55,weights=None): if not bl or all(len(b)==0 for b in bl): return np.array([]) ab,as_=[],[] for i,(bx,sx) in enumerate(zip(bl,sl)): w=weights[i] if weights else np.ones(len(bx)) for j in range(len(bx)): ab.append(bx[j]);as_.append(sx[j]*w[j]) if not ab: return np.array([]) ab=np.array(ab);as_=np.array(as_) o=np.argsort(-as_);ab=ab[o];as_=as_[o] cl,us=[],np.zeros(len(ab),dtype=bool) for i in range(len(ab)): if us[i]: continue c=[(ab[i],as_[i])];us[i]=True for j in range(i+1,len(ab)): if us[j]: continue tw=sum(s for _,s in c) ct=sum(b*s/tw for b,s in c) if iou_fn(ct.tolist(),ab[j].tolist())>thr: c.append((ab[j],as_[j]));us[j]=True cl.append(c) rb,rs=[],[] for c in cl: tw=sum(s for _,s in c) rb.append(sum(b*s/tw for b,s in c));rs.append(tw) return np.array(rb) def pred_fn(model,img,sz,flip=False,bright=1.0): ia=img if bright!=1.0: ia=ImageEnhance.Brightness(ia).enhance(bright) if flip: ia=ia.transpose(Image.FLIP_LEFT_RIGHT) r=model.predict(ia,imgsz=sz,conf=0.25,iou=0.7,max_det=100,verbose=False) if not r or len(r[0].boxes)==0: return np.array([]),np.array([]) b=r[0].boxes.xyxy.cpu().numpy();s=r[0].boxes.conf.cpu().numpy() if flip: w=img.size[0];b[:,[0,2]]=w-b[:,[2,0]] return b def main(): val_dir='Data/Detection_dataset/images/val' lbl_dir='Data/Detection_dataset/labels/val' vfs=sorted([f for f in os.listdir(val_dir) if f.endswith('.jpg')]) iou_thrs=[round(0.5+i*0.05,2) for i in range(10)] prior=learn_stereo_prior() def eval_boxes(name,boxes_per_img): tp={t:0 for t in iou_thrs};tg=0 for idx,vf in enumerate(vfs): img=Image.open(os.path.join(val_dir,vf)) gb=[] lf=vf.replace('.jpg','.txt') with open(os.path.join(lbl_dir,lf)) as f: for line in f: p=line.strip().split() if len(p)>=5: cx,cy,w,h=[float(x) for x in p[1:5]] gb.append([(cx-w/2)*img.size[0],(cy-h/2)*img.size[1],(cx+w/2)*img.size[0],(cy+h/2)*img.size[1]]) tg+=len(gb) if not gb: continue for t in iou_thrs: mt=set() for pb in boxes_per_img[idx]: if len(pb)==0: continue bi,bg=0,-1 for gi,gt in enumerate(gb): if gi in mt: continue ii=iou_fn(pb.tolist(),gt) if ii>bi: bi=ii;bg=gi if bi>=t and bg>=0: tp[t]+=1;mt.add(bg) rec=[tp[t]/tg for t in iou_thrs] mAP=np.mean(rec) print('{}: mAP50-95={:.4f}'.format(name,mAP)) return mAP # Use v6_1 for quick test m=YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt') # Baseline bp=[] for vf in tqdm(vfs,desc='Baseline'): img=Image.open(os.path.join(val_dir,vf)) b,_=pred_fn(m,img,1536);bp.append(b) mAP_base=eval_boxes('Baseline',bp) # Multi-scale WBF without stereo mp=[] for vf in tqdm(vfs,desc='MS-WBF'): img=Image.open(os.path.join(val_dir,vf)) bl,sl=[],[] for sz in [1280,1536,1920]: for fl in [False,True]: for br in [1.0,1.2]: b,s=pred_fn(m,img,sz,fl,br) if len(b)>0: bl.append(b);sl.append(s) mp.append(wbf_fn(bl,sl)) mAP_ms=eval_boxes('MS-WBF',mp) # Stereo-constrained: apply size consistency check to East shed images if prior: sp=[] for vf in tqdm(vfs,desc='StereoWBF'): img=Image.open(os.path.join(val_dir,vf));iw,ih=img.size cam=vf.split('_2025')[0] is_east=cam.startswith('East') bl,sl,weights=[],[],[] for sz in [1280,1536,1920]: for fl in [False,True]: for br in [1.0,1.2]: b,s=pred_fn(m,img,sz,fl,br) if len(b)>0: bl.append(b);sl.append(s) w=np.ones(len(b)) if is_east and prior: # Check size plausibility via stereo prior for i in range(len(b)): pred_w, pred_h = b[i][2]-b[i][0], b[i][3]-b[i][1] pred_sz = np.sqrt(pred_w*pred_h) / max(iw,ih) # Expected opposite-camera size if cam=='EastLeft': exp_sz = pred_sz * prior['size_ratio'] else: exp_sz = pred_sz / prior['size_ratio'] # If size is unusual for its camera, slight penalty z = abs(pred_sz - exp_sz) / (exp_sz + 1e-8) if z > 0.3: w[i] = 0.7 weights.append(w) sp.append(wbf_fn(bl,sl,weights=weights)) mAP_stereo=eval_boxes('StereoWBF',sp) print('\nBaseline: {:.4f}'.format(mAP_base)) print('MS-WBF: {:.4f} (+{:.4f})'.format(mAP_ms,mAP_ms-mAP_base)) print('Stereo: {:.4f} (+{:.4f})'.format(mAP_stereo,mAP_stereo-mAP_base)) del m;gc.collect();torch.cuda.empty_cache() if __name__=='__main__': main()