File size: 4,984 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 | """Multi-Scale WBF: Single model at 3 scales → WBF Ensemble
============================================================
Each model predicts at 1280, 1536, 1920. WBF fuses all 3 predictions.
Zero training cost. Tested on v6_1.
"""
import sys, os
import numpy as np
from PIL import Image
from tqdm import tqdm
PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)
from ultralytics import YOLO
def compute_iou(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 wbf(boxes_list, scores_list, iou_thr=0.55):
if not boxes_list or all(len(b) == 0 for b in boxes_list):
return np.array([]), np.array([])
all_boxes, all_scores = [], []
for boxes, scores in zip(boxes_list, scores_list):
for i in range(len(boxes)):
all_boxes.append(boxes[i])
all_scores.append(scores[i])
if not all_boxes: return np.array([]), np.array([])
all_boxes = np.array(all_boxes); all_scores = np.array(all_scores)
order = np.argsort(-all_scores)
all_boxes, all_scores = all_boxes[order], all_scores[order]
clusters, used = [], np.zeros(len(all_boxes), dtype=bool)
for i in range(len(all_boxes)):
if used[i]: continue
cluster = [(all_boxes[i], all_scores[i])]
used[i] = True
for j in range(i+1, len(all_boxes)):
if used[j]: continue
tw = sum(s for _, s in cluster)
center = sum(b*s/tw for b, s in cluster)
if compute_iou(center.tolist(), all_boxes[j].tolist()) > iou_thr:
cluster.append((all_boxes[j], all_scores[j]))
used[j] = True
clusters.append(cluster)
result_boxes, result_scores = [], []
for cl in clusters:
tw = sum(s for _, s in cl)
avg_b = sum(b*s/tw for b, s in cl)
result_boxes.append(avg_b)
result_scores.append(tw)
return np.array(result_boxes), np.array(result_scores)
def eval_config(name, get_boxes_fn):
val_img_dir = 'Data/Detection_dataset/images/val'
val_label_dir = 'Data/Detection_dataset/labels/val'
val_files = sorted([f for f in os.listdir(val_img_dir) if f.endswith('.jpg')])
iou_thrs = [round(0.5+i*0.05, 2) for i in range(10)]
tp = {t:0 for t in iou_thrs}
total_gt = 0
for img_file in tqdm(val_files, desc=name):
img = Image.open(os.path.join(val_img_dir, img_file))
gt_boxes = []
lf = img_file.replace('.jpg','.txt')
with open(os.path.join(val_label_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]]
gt_boxes.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]])
total_gt += len(gt_boxes)
if not gt_boxes: continue
preds = get_boxes_fn(img)
for t in iou_thrs:
matched = set()
for pb in preds:
best_iou, best_gi = 0, -1
for gi, gb in enumerate(gt_boxes):
if gi in matched: continue
iou = compute_iou(pb.tolist(), gb)
if iou > best_iou: best_iou = iou; best_gi = gi
if best_iou >= t and best_gi >= 0:
tp[t] += 1; matched.add(best_gi)
recalls = [tp[t]/total_gt for t in iou_thrs]
mAP = np.mean(recalls)
print(f' {name}: mAP50-95={mAP:.4f}, IoU@75={recalls[5]:.4f}')
return mAP
def main():
model = YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt')
scales = [1280, 1536, 1920]
# Baseline: single scale 1536
def single_1536(img):
r = model.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
if r and len(r[0].boxes): return r[0].boxes.xyxy.cpu().numpy()
return np.array([])
mAP_single = eval_config('Single 1536', single_1536)
# Multi-scale WBF
def multiscale_wbf(img):
boxes_list, scores_list = [], []
for sz in scales:
r = model.predict(img, imgsz=sz, conf=0.25, iou=0.7, max_det=100, verbose=False)
if r and len(r[0].boxes):
boxes_list.append(r[0].boxes.xyxy.cpu().numpy())
scores_list.append(r[0].boxes.conf.cpu().numpy())
else:
boxes_list.append(np.array([]))
scores_list.append(np.array([]))
fused, _ = wbf(boxes_list, scores_list)
return fused
mAP_ms = eval_config(f'MultiScale WBF ({scales})', multiscale_wbf)
print(f'\n{"="*50}')
print(f'Single 1536: {mAP_single:.4f}')
print(f'MultiScale WBF: {mAP_ms:.4f} (+{mAP_ms-mAP_single:+.4f})')
if __name__ == '__main__':
main()
|