File size: 6,450 Bytes
6a5bb7e | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """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()
|