"""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')]) # Discover models 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] # Baseline 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) # Config 1: Single model Kitchen Sink 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) # Config 2: All models Kitchen Sink 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})') # Save 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()