Scripts/eval_stereo_wbf.py

"""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()