| """WBF Ensemble Evaluation |
| ========================== |
| Weighted Box Fusion across multiple models for higher mAP50-95. |
| Combines predictions from different seed models using IoU-based clustering. |
| |
| Models: v6_1 + v12_seed_42 + v12_seed_123 |
| Weights: proportional to each model's mAP50-95 |
| """ |
| import sys, os |
| import numpy as np |
| from pathlib import Path |
| 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 = max(b1[0], b2[0]); y1 = max(b1[1], b2[1]) |
| x2 = min(b1[2], b2[2]); y2 = 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, weights, iou_thr=0.55): |
| """Weighted Box Fusion.""" |
| if not boxes_list: |
| return np.array([]), np.array([]) |
|
|
| all_boxes = [] |
| all_scores = [] |
|
|
| for m_idx, (boxes, scores) in enumerate(zip(boxes_list, scores_list)): |
| w = weights[m_idx] |
| for i in range(len(boxes)): |
| all_boxes.append(boxes[i]) |
| all_scores.append(scores[i] * w) |
|
|
| 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_boxes[order] |
| all_scores = 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 |
| total_w = sum(s for _, s in cluster) |
| center = np.zeros(4) |
| for b, s in cluster: |
| center += b * s / total_w |
| 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 cluster in clusters: |
| total_w = sum(s for _, s in cluster) |
| avg_box = np.zeros(4) |
| for b, s in cluster: |
| avg_box += b * s / total_w |
| avg_score = total_w / len(weights) |
| result_boxes.append(avg_box) |
| result_scores.append(avg_score) |
|
|
| return np.array(result_boxes), np.array(result_scores) |
|
|
|
|
| def main(): |
| |
| model_paths = [ |
| 'runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt', |
| 'runs/detect/Detection_experiments/v12_seed_42/weights/best.pt', |
| 'runs/detect/Detection_experiments/v12_seed_123/weights/best.pt', |
| ] |
| weights = [0.5125, 0.5097, 0.5113] |
| weights = [w / sum(weights) for w in weights] |
|
|
| models = [] |
| for i, mp in enumerate(model_paths): |
| if os.path.exists(mp): |
| models.append((YOLO(mp), weights[i])) |
| print(f'Loaded: {mp} (weight={weights[i]:.3f})') |
| else: |
| print(f'MISSING: {mp}') |
|
|
| if len(models) < 2: |
| print('Need at least 2 models') |
| return |
|
|
| 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_thresholds = [round(0.5 + i * 0.05, 2) for i in range(10)] |
| per_iou_tp = {t: 0 for t in iou_thresholds} |
| total_gt = 0 |
|
|
| for img_file in tqdm(val_files, desc='Evaluating'): |
| img_path = os.path.join(val_img_dir, img_file) |
| img = Image.open(img_path) |
| iw, ih = img.size |
|
|
| |
| lf = img_file.replace('.jpg', '.txt') |
| lpath = os.path.join(val_label_dir, lf) |
| gt_boxes = [] |
| if os.path.exists(lpath): |
| with open(lpath) 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) * iw, (cy - h/2) * ih, |
| (cx + w/2) * iw, (cy + h/2) * ih |
| ]) |
| total_gt += len(gt_boxes) |
| if not gt_boxes: |
| continue |
|
|
| |
| boxes_list = [] |
| scores_list = [] |
| for model, _ in models: |
| results = model.predict( |
| source=img_path, imgsz=1536, conf=0.25, iou=0.7, |
| max_det=100, save=False, verbose=False |
| ) |
| if results and len(results[0].boxes): |
| boxes_list.append(results[0].boxes.xyxy.cpu().numpy()) |
| scores_list.append(results[0].boxes.conf.cpu().numpy()) |
| else: |
| boxes_list.append(np.array([])) |
| scores_list.append(np.array([])) |
|
|
| |
| fused_boxes, fused_scores = wbf(boxes_list, scores_list, |
| [w for _, w in models]) |
|
|
| |
| for t in iou_thresholds: |
| matched = set() |
| for pb in fused_boxes: |
| 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: |
| per_iou_tp[t] += 1 |
| matched.add(best_gi) |
|
|
| |
| print(f'\n{"="*60}') |
| print(f'WBF Ensemble Results (3 models)') |
| print(f'{"="*60}') |
| recalls = [] |
| for t in iou_thresholds: |
| r = per_iou_tp[t] / total_gt if total_gt else 0 |
| recalls.append(r) |
| print(f' IoU@{t:.2f}: Recall={r:.4f}') |
| map5095 = np.mean(recalls) |
| print(f'\n Approx mAP50-95: {map5095:.4f}') |
| print(f' vs v6_1 single (0.5125): {map5095-0.5125:+.4f}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
