| """Kitchen Sink TTA WBF — Maximum inference-time firepower |
| =========================================================== |
| Every model × scale × augmentation → WBF fusion. |
| |
| Augmentations: |
| - 3 scales: 1280, 1536, 1920 |
| - 2 flips: none, horizontal |
| - 2 brightness: normal, +20% |
| |
| Per model: 3×2×2 = 12 variants |
| With 5 models: 60 prediction sources → WBF |
| |
| Also tests subset combinations to find the best cost/accuracy tradeoff. |
| """ |
| import sys, os, json |
| import 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 |
|
|
|
|
| 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): |
| all_boxes.extend(boxes) |
| all_scores.extend(scores) |
| 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 predict_augmented(model, img, imgsz, flip=False, brighten=1.0): |
| """Predict on (possibly augmented) image.""" |
| img_aug = img |
| if brighten != 1.0: |
| enhancer = ImageEnhance.Brightness(img) |
| img_aug = enhancer.enhance(brighten) |
| if flip: |
| img_aug = img_aug.transpose(Image.FLIP_LEFT_RIGHT) |
|
|
| r = model.predict(img_aug, imgsz=imgsz, 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([]) |
|
|
| boxes = r[0].boxes.xyxy.cpu().numpy() |
| scores = r[0].boxes.conf.cpu().numpy() |
|
|
| if flip: |
| w = img.size[0] |
| boxes[:, [0, 2]] = w - boxes[:, [2, 0]] |
|
|
| return boxes, scores |
|
|
|
|
| def evaluate(name, get_boxes_fn, val_files, val_img_dir, val_label_dir): |
| 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] |
| return np.mean(recalls), recalls[5] |
|
|
|
|
| def main(): |
| 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')]) |
|
|
| |
| exp_dir = 'runs/detect/Detection_experiments' |
| models = [] |
| priority = ['v6_1_s_refined', 'v12_seed_42', 'v12_seed_123', 'v14_seed_789', 'v14_seed_999'] |
| for name in priority: |
| path = os.path.join(exp_dir, name, 'weights', 'best.pt') |
| if os.path.exists(path): |
| models.append((YOLO(path), name)) |
| print(f'Loaded: {name}') |
|
|
| if not models: |
| print('No models found!') |
| return |
|
|
| print(f'\n{len(models)} models ready') |
|
|
| scales = [1280, 1536, 1920] |
| flips = [False, True] |
| brights = [1.0, 1.2] |
|
|
| |
| def baseline(img): |
| m, _ = models[0] |
| boxes, scores = predict_augmented(m, img, 1536) |
| return boxes |
|
|
| mAP_base, r75_base = evaluate('Single baseline', baseline, val_files, val_img_dir, val_label_dir) |
|
|
| |
| def single_ks(img): |
| m, _ = models[0] |
| all_boxes, all_scores = [], [] |
| for sz in scales: |
| for fl in flips: |
| for br in brights: |
| b, s = predict_augmented(m, img, sz, fl, br) |
| if len(b) > 0: |
| all_boxes.append(b); all_scores.append(s) |
| return wbf(all_boxes, all_scores)[0] |
|
|
| mAP_sks, r75_sks = evaluate(f'1m Kitchen Sink ({len(scales)*len(flips)*len(brights)}x)', single_ks, val_files, val_img_dir, val_label_dir) |
|
|
| |
| def all_ks(img): |
| all_boxes, all_scores = [], [] |
| for m, _ in models: |
| for sz in scales: |
| for fl in flips: |
| for br in brights: |
| b, s = predict_augmented(m, img, sz, fl, br) |
| if len(b) > 0: |
| all_boxes.append(b); all_scores.append(s) |
| return wbf(all_boxes, all_scores)[0] |
|
|
| n_sources = len(models) * len(scales) * len(flips) * len(brights) |
| mAP_aks, r75_aks = evaluate(f'ALL Kitchen Sink ({n_sources}x)', all_ks, val_files, val_img_dir, val_label_dir) |
|
|
| print(f'\n{"="*60}') |
| print(f'KITCHEN SINK RESULTS') |
| print(f'{"="*60}') |
| print(f'Baseline (1m, 1536): {mAP_base:.4f} IoU@75={r75_base:.4f}') |
| print(f'1m Kitchen Sink (12x): {mAP_sks:.4f} IoU@75={r75_sks:.4f} (+{mAP_sks-mAP_base:+.4f})') |
| print(f'ALL Kitchen Sink ({n_sources}x): {mAP_aks:.4f} IoU@75={r75_aks:.4f} (+{mAP_aks-mAP_base:+.4f})') |
|
|
| |
| with open('logs/kitchen_sink_results.json', 'w') as f: |
| json.dump({ |
| 'baseline': round(mAP_base, 4), |
| 'single_kitchen_sink': round(mAP_sks, 4), |
| 'all_kitchen_sink': round(mAP_aks, 4), |
| 'n_models': len(models), |
| 'n_sources': n_sources, |
| }, f, indent=2) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
