Scripts/eval_mask_wbf.py

"""Mask-weighted WBF: use background-diff masks to weight predictions.
Higher mask overlap = higher confidence = higher WBF weight.
"""
import sys,os,json,gc,numpy as np
from PIL import Image,ImageEnhance
from tqdm import tqdm
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 mask_iou(box, mask):
    """Compute IoU between a bounding box and a binary mask."""
    x1,y1,x2,y2=int(box[0]),int(box[1]),int(box[2]),int(box[3])
    x1,y1=max(0,x1),max(0,y1)
    x2,y2=min(mask.shape[1],x2),min(mask.shape[0],y2)
    if x2<=x1 or y2<=y1: return 0
    box_area=(x2-x1)*(y2-y1)
    mask_area=mask[y1:y2,x1:x2].sum()/255
    union=box_area+mask.sum()/255-mask_area
    return mask_area/(union+1e-8)

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 1.0
        for j in range(len(bx)):
            ab.append(bx[j]);as_.append(sx[j]*w)
    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,s

def main():
    val_dir='Data/Detection_dataset/images/val'
    lbl_dir='Data/Detection_dataset/labels/val'
    seg_dir='Data/Detection_dataset/labels/val_seg'
    vfs=sorted([f for f in os.listdir(val_dir) if f.endswith('.jpg')])
    exp='runs/detect/Detection_experiments'

    # Load masks for val images (only if available, val masks weren't generated)
    # For now, test on training images to validate the concept
    # Use masks as additional confidence signal

    mpaths=[('v6_1',f'{exp}/v6_1_s_refined/weights/best.pt')]
    mpaths=[(n,p) for n,p in mpaths if os.path.exists(p)]

    iou_thrs=[round(0.5+i*0.05,2) for i in range(10)]

    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} IoU@75={:.4f}'.format(name,mAP,rec[5]))
        return mAP,rec[5]

    # Baseline
    m0=YOLO(mpaths[0][1])
    bp=[]
    for vf in tqdm(vfs,desc='Baseline'):
        img=Image.open(os.path.join(val_dir,vf))
        b,_=pred_fn(m0,img,1536);bp.append(b)
    del m0;gc.collect();torch.cuda.empty_cache()
    mAP_base,r75_base=eval_boxes('Baseline',bp)

    # Quick test: single model + mask-weighted multi-scale
    m=YOLO(mpaths[0][1])
    mp=[]
    for vf in tqdm(vfs,desc='MaskWBF'):
        img=Image.open(os.path.join(val_dir,vf))
        # Load mask if exists
        mask_path=os.path.join(seg_dir,vf.replace('.jpg','.png'))
        mask=None
        if os.path.exists(mask_path):
            mask=np.array(Image.open(mask_path))
            mask=(mask>128).astype(np.uint8)*255

        bl,sl=[],[]
        box_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)
                        # Compute mask weights
                        if mask is not None:
                            w=np.ones(len(b))
                            for i in range(len(b)):
                                mi=mask_iou(b[i],mask)
                                w[i]=(1.0+mi)/2.0  # [0.5, 1.0]
                            box_weights.append(w)
                        else:
                            box_weights.append(np.ones(len(b)))

        # Apply weights to scores
        for i,(bx,sx,w) in enumerate(zip(bl,sl,box_weights)):
            sl[i]=sx*w

        mp.append(wbf_fn(bl,sl))
    mAP_mask,r75_mask=eval_boxes('MaskWBF',mp)
    del m;gc.collect();torch.cuda.empty_cache()

    sep='='*60
    print('\n{}'.format(sep))
    print('MASK-WEIGHTED WBF')
    print(sep)
    print('Baseline: {:.4f}'.format(mAP_base))
    print('MaskWBF:  {:.4f}  (+{:.4f})'.format(mAP_mask,mAP_mask-mAP_base))

    with open('logs/mask_wbf_results.json','w') as f:
        json.dump({'baseline':round(mAP_base,4),'mask_wbf':round(mAP_mask,4),'delta':round(mAP_mask-mAP_base,4)},f,indent=2)
    print('Saved.')

if __name__=='__main__':
    main()