File size: 7,615 Bytes
9f7ad84 | 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 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """Benchmark different ensemble configs on GPU. Measures combined mAP score."""
import json
import time
import itertools
import numpy as np
from pathlib import Path
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
# Patch torch.load
import torch
_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe
from ultralytics import YOLO
from ensemble_boxes import weighted_boxes_fusion
ANN = "input/train/annotations.json"
IMG_DIR = Path("input/train/images")
# Load ground truth
coco_gt = COCO(ANN)
img_ids = sorted(coco_gt.getImgIds())
img_infos = coco_gt.loadImgs(img_ids)
def get_predictions_single(model, imgsz=1280, conf=0.001, iou=0.6, max_det=500):
"""Run single model, return predictions list."""
preds = []
for img_info in img_infos:
img_path = IMG_DIR / img_info["file_name"]
results = model.predict(str(img_path), verbose=False, imgsz=imgsz,
conf=conf, iou=iou, max_det=max_det, augment=False)
for r in results:
if r.boxes is None: continue
for box, sc, cl in zip(r.boxes.xyxy.cpu().numpy(),
r.boxes.conf.cpu().numpy(),
r.boxes.cls.cpu().numpy()):
x1, y1, x2, y2 = box
preds.append({
"image_id": img_info["id"], "category_id": int(cl),
"bbox": [float(x1), float(y1), float(x2-x1), float(y2-y1)],
"score": float(sc)
})
return preds
def get_predictions_ensemble(models, imgsz=1280, conf=0.001, iou=0.7,
max_det=500, wbf_iou=0.55, skip_thr=0.001):
"""Run ensemble with WBF, return predictions list."""
preds = []
for img_info in img_infos:
img_path = IMG_DIR / img_info["file_name"]
all_boxes, all_scores, all_labels = [], [], []
img_w, img_h = None, None
for model in models:
results = model.predict(str(img_path), verbose=False, imgsz=imgsz,
conf=conf, iou=iou, max_det=max_det, augment=False)
boxes, scores, labels = [], [], []
for r in results:
if r.boxes is None: continue
if img_w is None: img_h, img_w = r.orig_shape
for box, sc, cl in zip(r.boxes.xyxy.cpu().numpy(),
r.boxes.conf.cpu().numpy(),
r.boxes.cls.cpu().numpy()):
x1, y1, x2, y2 = box
boxes.append([x1/img_w, y1/img_h, x2/img_w, y2/img_h])
scores.append(float(sc))
labels.append(int(cl))
all_boxes.append(boxes if boxes else [[0,0,0,0]])
all_scores.append(scores if scores else [0])
all_labels.append(labels if labels else [0])
if img_w is None: continue
fb, fs, fl = weighted_boxes_fusion(all_boxes, all_scores, all_labels,
iou_thr=wbf_iou, skip_box_thr=skip_thr)
for box, sc, lb in zip(fb, fs, fl):
x1, y1 = box[0]*img_w, box[1]*img_h
x2, y2 = box[2]*img_w, box[3]*img_h
w, h = x2-x1, y2-y1
if w > 0 and h > 0:
preds.append({
"image_id": img_info["id"], "category_id": int(lb),
"bbox": [float(x1), float(y1), float(w), float(h)],
"score": float(sc)
})
return preds
def compute_score(preds):
"""Compute 0.7*det_mAP + 0.3*cls_mAP."""
if not preds:
return 0, 0, 0
# Cap at 49000
if len(preds) > 49000:
preds.sort(key=lambda x: x["score"], reverse=True)
preds = preds[:49000]
gt_data = json.load(open(ANN))
# Detection mAP (category-agnostic)
det_preds = [dict(p, category_id=1) for p in preds]
det_gt = dict(gt_data)
det_gt["annotations"] = [dict(a, category_id=1) for a in gt_data["annotations"]]
det_gt["categories"] = [{"id": 1, "name": "product"}]
with open("/tmp/det_gt.json", "w") as f: json.dump(det_gt, f)
with open("/tmp/det_pr.json", "w") as f: json.dump(det_preds, f)
coco_d = COCO("/tmp/det_gt.json")
ev = COCOeval(coco_d, coco_d.loadRes("/tmp/det_pr.json"), "bbox")
ev.params.iouThrs = [0.5]
ev.evaluate(); ev.accumulate(); ev.summarize()
det_map = ev.stats[0]
# Classification mAP
with open("/tmp/cls_pr.json", "w") as f: json.dump(preds, f)
coco_c = COCO(ANN)
ev2 = COCOeval(coco_c, coco_c.loadRes("/tmp/cls_pr.json"), "bbox")
ev2.params.iouThrs = [0.5]
ev2.evaluate(); ev2.accumulate(); ev2.summarize()
cls_map = ev2.stats[0]
combined = 0.7 * det_map + 0.3 * cls_map
return combined, det_map, cls_map
# Load models
print("Loading models...")
m1 = YOLO("bench/model1.onnx", task="detect") # run4: s123, 1280, SGD (best)
m2 = YOLO("bench/model2.onnx", task="detect") # run1: s42, 1024, AdamW
m3 = YOLO("bench/model3.onnx", task="detect") # run2: s42, 1280, AdamW
# Also load the .pt models for more diversity
pt_models = {}
for pt in Path(".").glob("train_*/run/weights/best.pt"):
name = pt.parts[0]
pt_models[name] = YOLO(str(pt), task="detect")
print(f" Loaded {name}")
print(f"\nLoaded 3 ONNX + {len(pt_models)} PT models")
print("=" * 60)
results = []
# Test 1: Single models
print("\n--- SINGLE MODELS ---")
for name, model in [("model1(s123_1280_SGD)", m1), ("model2(s42_1024_AdamW)", m2), ("model3(s42_1280_AdamW)", m3)]:
t = time.time()
preds = get_predictions_single(model)
score, det, cls = compute_score(preds)
elapsed = time.time() - t
print(f"{name}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
results.append((name, score, det, cls))
# Test 2: 3-model ensemble with different WBF iou thresholds
print("\n--- ENSEMBLE WBF IOU SWEEP ---")
for wbf_iou in [0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7]:
t = time.time()
preds = get_predictions_ensemble([m1, m2, m3], wbf_iou=wbf_iou)
score, det, cls = compute_score(preds)
elapsed = time.time() - t
print(f"wbf_iou={wbf_iou}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
results.append((f"ensemble_wbf{wbf_iou}", score, det, cls))
# Test 3: Different max_det
print("\n--- MAX_DET SWEEP ---")
for max_det in [300, 500, 800, 1000]:
t = time.time()
preds = get_predictions_ensemble([m1, m2, m3], max_det=max_det, wbf_iou=0.55)
score, det, cls = compute_score(preds)
elapsed = time.time() - t
print(f"max_det={max_det}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
results.append((f"maxdet_{max_det}", score, det, cls))
# Test 4: Different conf thresholds
print("\n--- CONF SWEEP ---")
for conf in [0.0001, 0.001, 0.005, 0.01]:
t = time.time()
preds = get_predictions_ensemble([m1, m2, m3], conf=conf, wbf_iou=0.55)
score, det, cls = compute_score(preds)
elapsed = time.time() - t
print(f"conf={conf}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
results.append((f"conf_{conf}", score, det, cls))
# Summary
print("\n" + "=" * 60)
print("RANKED RESULTS:")
results.sort(key=lambda x: x[1], reverse=True)
for i, (name, score, det, cls) in enumerate(results):
marker = " <-- BEST" if i == 0 else ""
print(f" {i+1}. {name}: {score:.4f} (det={det:.4f} cls={cls:.4f}){marker}")
|