"""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()