Scripts/eval_brn.py

"""Evaluate BRN: YOLO detection → BRN refinement → mAP
Also test: YOLO → WBF Ensemble → BRN refinement → mAP (ultimate combo)
"""
import sys, os
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)

import torch
from ultralytics import YOLO
from Scripts.modules.brn import BRN


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, iou_thr=0.55):
    if not boxes_list:
        return np.array([]), np.array([])
    all_boxes, all_scores = [], []
    for m_idx, (boxes, scores) in enumerate(zip(boxes_list, scores_list)):
        for i in range(len(boxes)):
            all_boxes.append(boxes[i])
            all_scores.append(scores[i] * weights[m_idx])
    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
            total_w = sum(s for _, s in cluster)
            center = sum(b * s / total_w 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 cluster in clusters:
        total_w = sum(s for _, s in cluster)
        avg_box = sum(b * s / total_w for b, s in cluster)
        avg_score = total_w / len(weights)
        result_boxes.append(avg_box)
        result_scores.append(avg_score)
    return np.array(result_boxes), np.array(result_scores)


def brn_refine(image, boxes, brn_model, device, input_size=96):
    """Refine boxes using BRN."""
    if len(boxes) == 0:
        return boxes

    import torchvision.transforms as T
    transform = T.Compose([
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    iw, ih = image.size
    refined = []
    for box in boxes:
        x1, y1, x2, y2 = box
        cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
        w, h = max(x2 - x1, 1), max(y2 - y1, 1)
        roi_size = max(w, h) * 1.5
        roi_size = min(roi_size, min(iw, ih))
        half = roi_size / 2
        rx1 = max(0, int(cx - half)); ry1 = max(0, int(cy - half))
        rx2 = min(iw, int(cx + half)); ry2 = min(ih, int(cy + half))
        if rx2 - rx1 < 8 or ry2 - ry1 < 8:
            refined.append(box); continue
        roi = image.crop((rx1, ry1, rx2, ry2))
        roi_t = transform(roi.resize((input_size, input_size), Image.BICUBIC))
        roi_t = roi_t.unsqueeze(0).to(device)
        delta = brn_model(roi_t)[0].cpu().detach().numpy()
        roi_w, roi_h = rx2 - rx1, ry2 - ry1
        norm = max(w, h)
        dcx = delta[0] * norm; dcy = delta[1] * norm
        dw = delta[2] * norm; dh = delta[3] * norm
        new_cx = cx + dcx; new_cy = cy + dcy
        new_w = max(4, w + dw); new_h = max(4, h + dh)
        new_x1 = max(0, new_cx - new_w/2); new_y1 = max(0, new_cy - new_h/2)
        new_x2 = min(iw, new_cx + new_w/2); new_y2 = min(ih, new_cy + new_h/2)
        refined.append([new_x1, new_y1, new_x2, new_y2])
    return np.array(refined)


def evaluate(name, boxes_list_fn):
    """Evaluate a detection pipeline on val set."""
    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')])
    iou_thrs = [round(0.5 + i * 0.05, 2) for i in range(10)]
    per_iou_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))
        iw, ih = img.size
        # GT
        lf = img_file.replace('.jpg', '.txt')
        gt_boxes = []
        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)*iw, (cy-h/2)*ih, (cx+w/2)*iw, (cy+h/2)*ih])
        total_gt += len(gt_boxes)
        if not gt_boxes: continue

        # Get predictions
        fused_boxes = boxes_list_fn(img, img_file)
        fused_boxes = np.array(fused_boxes) if len(fused_boxes) > 0 else np.array([])

        for t in iou_thrs:
            matched = set()
            for pb in fused_boxes:
                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:
                    per_iou_tp[t] += 1; matched.add(best_gi)

    recalls = []
    print(f'\n{name}:')
    for t in iou_thrs:
        r = per_iou_tp[t] / total_gt if total_gt else 0
        recalls.append(r)
    mAP = np.mean(recalls)
    print(f'  mAP50-95: {mAP:.4f}')
    print(f'  IoU@75:   {recalls[5]:.4f}')
    return mAP


def main():
    device = torch.device('cuda')

    # Load models
    yolo = YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt')
    yolo2 = YOLO('runs/detect/Detection_experiments/v12_seed_42/weights/best.pt')
    yolo3 = YOLO('runs/detect/Detection_experiments/v12_seed_123/weights/best.pt')
    yol_models = [yolo, yolo2, yolo3]
    yol_weights = [0.5125, 0.5097, 0.5113]
    yol_weights = [w / sum(yol_weights) for w in yol_weights]

    # Load BRN
    brn_model = BRN(input_size=96).to(device)
    brn_model.load_state_dict(torch.load('runs/brn/brn_best.pt', map_location=device))
    brn_model.eval()

    # ── Eval 1: YOLO single ──
    def yolo_single(img, fname):
        r = yolo.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_yolo = evaluate('YOLO single', yolo_single)

    # ── Eval 2: YOLO + BRN ──
    def yolo_brn(img, fname):
        r = yolo.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
        if r and len(r[0].boxes):
            boxes = r[0].boxes.xyxy.cpu().numpy()
            return brn_refine(img, boxes, brn_model, device)
        return np.array([])

    mAP_brn = evaluate('YOLO + BRN', yolo_brn)

    # ── Eval 3: WBF Ensemble ──
    def wbf_ensemble(img, fname):
        boxes_list, scores_list = [], []
        for model in yol_models:
            r = model.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, yol_weights)
        return fused

    mAP_wbf = evaluate('WBF Ensemble (3 models)', wbf_ensemble)

    # ── Eval 4: WBF + BRN (ultimate) ──
    def wbf_brn(img, fname):
        boxes_list, scores_list = [], []
        for model in yol_models:
            r = model.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, yol_weights)
        if len(fused) > 0:
            return brn_refine(img, fused, brn_model, device)
        return np.array([])

    mAP_ultimate = evaluate('WBF + BRN (ULTIMATE)', wbf_brn)

    # ── Summary ──
    print(f'\n{"="*60}')
    print(f'FINAL RESULTS')
    print(f'{"="*60}')
    print(f'  v6_1 single:         0.5125  (baseline)')
    print(f'  YOLO single:         {mAP_yolo:.4f}')
    print(f'  YOLO + BRN:          {mAP_brn:.4f}  (+{mAP_brn-0.5125:+.4f})')
    print(f'  WBF Ensemble:        {mAP_wbf:.4f}  (+{mAP_wbf-0.5125:+.4f})')
    print(f'  WBF + BRN ULTIMATE:  {mAP_ultimate:.4f}  (+{mAP_ultimate-0.5125:+.4f})')


if __name__ == '__main__':
    main()