Scripts/eval_ultimate.py

"""Ultimate WBF: All models × All scales
========================================
Maximum firepower: 5 models × 3 scales = 15 prediction sources → WBF fusion.
Also tests per-camera weighted WBF and confidence-gated variants.
"""
import sys, os, json
import numpy as np
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, 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):
        for i in range(len(boxes)):
            all_boxes.append(boxes[i])
            all_scores.append(scores[i])
    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 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]
    mAP = np.mean(recalls)
    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 available models
    exp_dir = 'runs/detect/Detection_experiments'
    all_models = []
    for ename in sorted(os.listdir(exp_dir)):
        best_path = os.path.join(exp_dir, ename, 'weights', 'best.pt')
        if os.path.exists(best_path):
            all_models.append((ename, best_path))

    print(f'Found {len(all_models)} trained models:')
    for ename, path in all_models:
        print(f'  {ename}')

    # Select ensemble models (prefer v6_1 + v12 + v14 seeds)
    ensemble_models = []
    priority_patterns = ['v6_1', 'v12_seed', 'v14_seed']
    for ename, path in all_models:
        for pat in priority_patterns:
            if pat in ename:
                ensemble_models.append((ename, path))
                break

    # Limit to 5 best (by filename, assume newer seeds are good)
    ensemble_models = ensemble_models[:5]
    print(f'\nEnsemble: {len(ensemble_models)} models')
    for ename, _ in ensemble_models:
        print(f'  {ename}')

    models = [(YOLO(p), ename) for ename, p in ensemble_models]
    scales = [1280, 1536, 1920]

    # Single model baseline
    def baseline_1536(img):
        m, _ = models[0]
        r = m.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
        if r and len(r[0].boxes): return r[0].boxes.xyxy.cpu().numpy()
        return np.array([])

    mAP_single, r75_single = evaluate('Single model', baseline_1536, val_files, val_img_dir, val_label_dir)
    print(f'\n{"="*60}')
    print(f'RESULTS:')
    print(f'  Single:   mAP50-95={mAP_single:.4f}  IoU@75={r75_single:.4f}')

    # Multi-model WBF (1536 only)
    def multi_model_wbf(img):
        boxes_list, scores_list = [], []
        for m, _ in models:
            r = m.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
            if r and len(r[0].boxes):
                boxes_list.append(r[0].boxes.xyxy.cpu().numpy())
                scores_list.append(r[0].boxes.conf.cpu().numpy())
            else:
                boxes_list.append(np.array([]))
                scores_list.append(np.array([]))
        fused, _ = wbf(boxes_list, scores_list)
        return fused

    mAP_mm, r75_mm = evaluate(f'{len(models)}model WBF', multi_model_wbf, val_files, val_img_dir, val_label_dir)
    print(f'  {len(models)}model:  mAP50-95={mAP_mm:.4f}  IoU@75={r75_mm:.4f}  (+{mAP_mm-mAP_single:+.4f})')

    # Single model multi-scale WBF
    def single_multiscale_wbf(img):
        m, _ = models[0]
        boxes_list, scores_list = [], []
        for sz in scales:
            r = m.predict(img, imgsz=sz, conf=0.25, iou=0.7, max_det=100, verbose=False)
            if r and len(r[0].boxes):
                boxes_list.append(r[0].boxes.xyxy.cpu().numpy())
                scores_list.append(r[0].boxes.conf.cpu().numpy())
            else:
                boxes_list.append(np.array([]))
                scores_list.append(np.array([]))
        fused, _ = wbf(boxes_list, scores_list)
        return fused

    mAP_sms, r75_sms = evaluate('1model 3scale', single_multiscale_wbf, val_files, val_img_dir, val_label_dir)
    print(f'  1m x 3sc:  mAP50-95={mAP_sms:.4f}  IoU@75={r75_sms:.4f}  (+{mAP_sms-mAP_single:+.4f})')

    # ULTIMATE: All models × All scales
    def ultimate_wbf(img):
        boxes_list, scores_list = [], []
        for m, _ in models:
            for sz in scales:
                r = m.predict(img, imgsz=sz, conf=0.25, iou=0.7, max_det=100, verbose=False)
                if r and len(r[0].boxes):
                    boxes_list.append(r[0].boxes.xyxy.cpu().numpy())
                    scores_list.append(r[0].boxes.conf.cpu().numpy())
                else:
                    boxes_list.append(np.array([]))
                    scores_list.append(np.array([]))
        fused, _ = wbf(boxes_list, scores_list)
        return fused

    mAP_ult, r75_ult = evaluate(f'{len(models)}m x {len(scales)}sc ULTIMATE', ultimate_wbf, val_files, val_img_dir, val_label_dir)
    print(f'  ULTIMATE: mAP50-95={mAP_ult:.4f}  IoU@75={r75_ult:.4f}  (+{mAP_ult-mAP_single:+.4f})')

    # Save results
    results = {
        'single': round(mAP_single, 4),
        'multimodel': round(mAP_mm, 4),
        'multiscale': round(mAP_sms, 4),
        'ultimate': round(mAP_ult, 4),
        'n_models': len(models),
        'n_scales': len(scales),
        'model_names': [n for n, _ in ensemble_models],
    }
    with open('logs/ultimate_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print(f'\nResults saved to logs/ultimate_results.json')


if __name__ == '__main__':
    main()