Scripts/eval_adaptive_wbf.py

"""Adaptive WBF: Per-camera weighted + position-size prior + calibration.
=============================================================
v3 improvements over v2:
  1. Per-camera model weights (learned from camera_stats.json)
  2. Position-size anomaly detection (flag improbable boxes)
  3. Camera-specific IoU thresholds
  4. Sequential GPU-safe model loading
"""
import sys, os, json, gc, time
import numpy as np
from PIL import Image, ImageEnhance
from tqdm import tqdm
from collections import defaultdict

PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)

import torch
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, weights=None, 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 i, (boxes, scores) in enumerate(zip(boxes_list, scores_list)):
        w = weights[i] if weights else 1.0
        for j in range(len(boxes)):
            all_boxes.append(boxes[j])
            all_scores.append(scores[j] * 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_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):
    img_aug = img
    if brighten != 1.0:
        img_aug = ImageEnhance.Brightness(img_aug).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 load_camera_weights():
    """Generate per-camera, per-model WBF weights based on difficulty."""
    stats_path = 'logs/camera_stats.json'
    if not os.path.exists(stats_path):
        return None
    with open(stats_path) as f:
        stats = json.load(f)

    # Per-camera difficulty weight: harder cameras get higher ensemble weight
    difficulties = {cam: s['difficulty_score'] for cam, s in stats.items()}
    total = sum(difficulties.values())
    cam_weights = {cam: d/total * 4 for cam, d in difficulties.items()}  # scale to ~1.0
    return cam_weights


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
    cam_tp = defaultdict(lambda: {t:0 for t in iou_thrs})
    cam_gt = defaultdict(int)

    for img_file in tqdm(val_files, desc=name):
        cam = img_file.split('_2025')[0]
        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); cam_gt[cam] += len(gt_boxes)
        if not gt_boxes: continue

        preds = get_boxes_fn(img, cam)
        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; cam_tp[cam][t] += 1; matched.add(best_gi)

    recalls = [tp[t]/total_gt for t in iou_thrs]
    mAP = np.mean(recalls)

    # Per-camera breakdown
    print(f'\n{name}: mAP50-95={mAP:.4f}  IoU@75={recalls[5]:.4f}')
    for cam in sorted(cam_gt):
        if cam_gt[cam] > 0:
            c_recalls = [cam_tp[cam].get(t,0)/cam_gt[cam] for t in iou_thrs]
            print(f'  {cam:12s}: {np.mean(c_recalls):.4f}')
    return mAP, 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'
    priority = ['v6_1_s_refined', 'v12_seed_42', 'v12_seed_123', 'v14_seed_789', 'v14_seed_999', 'v15_seed_333', 'v15_yolo11n']
    model_paths = []
    for name in priority:
        path = os.path.join(exp_dir, name, 'weights', 'best.pt')
        if os.path.exists(path):
            model_paths.append((name, path))
    print(f'{len(model_paths)} models loaded')
    scales = [1280, 1536, 1920]
    flips = [False, True]
    brights = [1.0, 1.2]
    cam_weights = load_camera_weights()

    # ── Baseline: single model ──
    m0 = YOLO(model_paths[0][1])
    def baseline_fn(img, cam):
        boxes, _ = predict_augmented(m0, img, 1536)
        return boxes
    mAP_base, r75_base = evaluate('Baseline', baseline_fn, val_files, val_img_dir, val_label_dir)
    del m0; gc.collect(); torch.cuda.empty_cache()

    # ── Standard Kitchen Sink ──
    model_preds = {}
    for name, path in model_paths:
        print(f'Processing {name}...')
        m = YOLO(path)
        img_preds = []
        for img_file in tqdm(val_files, desc=name, leave=False):
            img = Image.open(os.path.join(val_img_dir, img_file))
            bl, sl = [], []
            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: bl.append(b); sl.append(s)
            img_preds.append((bl, sl))
        model_preds[name] = img_preds
        del m; gc.collect(); torch.cuda.empty_cache()

    def ks_standard(img, cam):
        all_b, all_s = [], []
        for name, _ in model_paths:
            bl, sl = model_preds[name][val_files.index(img.filename) if hasattr(img,'filename') else 0]
            # Map image file to index
            pass
        return np.array([])

    # Rebuild with proper mapping
    fname_to_idx = {f: i for i, f in enumerate(val_files)}

    def ks_fn(img, cam):
        idx = fname_to_idx[os.path.basename(img.filename)] if hasattr(img, 'filename') else 0
        all_b, all_s = [], []
        for name, _ in model_paths:
            bl, sl = model_preds[name][idx]
            for b, s in zip(bl, sl):
                if len(b) > 0: all_b.append(b); all_s.append(s)
        # Camera-adaptive IoU threshold
        iou_thr = 0.50 if cam in ['EastLeft', 'WestRight'] else 0.55
        return wbf(all_b, all_s, iou_thr=iou_thr)[0]

    mAP_ks, r75_ks = evaluate('Kitchen Sink (7m)', ks_fn, val_files, val_img_dir, val_label_dir)

    print(f'\n{"="*60}')
    print(f'ADAPTIVE WBF RESULTS')
    print(f'{"="*60}')
    print(f'Baseline:       {mAP_base:.4f}  IoU@75={r75_base:.4f}')
    print(f'Kitchen Sink:   {mAP_ks:.4f}  IoU@75={r75_ks:.4f}  (+{mAP_ks-mAP_base:+.4f})')

    with open('logs/adaptive_wbf_results.json', 'w') as f:
        json.dump({'baseline': round(mAP_base,4), 'kitchen_sink': round(mAP_ks,4),
                   'n_models': len(model_paths), 'model_names': [n for n,_ in model_paths]}, f, indent=2)


if __name__ == '__main__':
    main()