#!/usr/bin/env python3 """5-fold WBF ensemble evaluation. Strategy: each fold model evaluated on ITS OWN val set (out-of-fold prediction is what CV gives us). For ensemble: each test image is predicted by ALL fold models, predictions fused via WBF, mAP computed against ground truth. For our 5-fold CV, we use the union of all val sets and predict each image with all 5 models then WBF. """ import os, sys, json, time from pathlib import Path import numpy as np import torch ROOT = Path('/arf/scratch/stakan/hitit-proje') FOLD_ROOT = ROOT / 'datasets/ready/detection_tablets' WEIGHTS_ROOT = ROOT / 'runs/detect/hitit_ocr/runs/h100' def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True) def main(): from ultralytics import YOLO from ensemble_boxes import weighted_boxes_fusion from PIL import Image # Tüm val setlerini birleştir (full dataset, leakage olmadan: her image kendi fold'unun val'ında) all_imgs = [] img_to_fold = {} for fold in range(5): for line in (FOLD_ROOT / f'fold_{fold}/val.txt').read_text().splitlines(): line = line.strip() if not line: continue all_imgs.append(line) img_to_fold[line] = fold log(f"Total eval images: {len(all_imgs)}") # Modelleri yükle models = {} for fold in range(5): ckpt = WEIGHTS_ROOT / f'yolo_fold{fold}/weights/best.pt' if not ckpt.exists(): log(f"WARN: {ckpt} eksik") continue models[fold] = YOLO(str(ckpt)) log(f"Loaded fold {fold}: {ckpt}") # Her image için predict (own-fold model ve ensemble karşılaştırması) own_preds = {} # img -> (boxes, scores) ensemble_preds = {} # img -> (boxes, scores) iou_thr = 0.55 skip_box_thr = 0.001 log("Predicting...") for idx, img_path in enumerate(all_imgs): if idx % 100 == 0: log(f" {idx}/{len(all_imgs)}") own_fold = img_to_fold[img_path] # Image size for normalization try: with Image.open(img_path) as im: W, H = im.size except Exception: continue boxes_list, scores_list, labels_list = [], [], [] for fold, model in models.items(): try: r = model.predict(img_path, conf=0.001, iou=0.7, max_det=2000, imgsz=1280, verbose=False, device=0 if torch.cuda.is_available() else 'cpu')[0] except Exception as e: log(f" predict skip {Path(img_path).name}: {e}"); break if r.boxes is None or len(r.boxes) == 0: continue xyxy = r.boxes.xyxy.cpu().numpy() / np.array([W, H, W, H]) # normalize 0-1 xyxy = np.clip(xyxy, 0, 1) scores = r.boxes.conf.cpu().numpy().tolist() labels = [0] * len(scores) boxes_list.append(xyxy.tolist()) scores_list.append(scores) labels_list.append(labels) if fold == own_fold: own_preds[img_path] = (xyxy, np.array(scores)) if not boxes_list: ensemble_preds[img_path] = (np.zeros((0, 4)), np.zeros(0)) continue try: wbf_boxes, wbf_scores, wbf_labels = weighted_boxes_fusion( boxes_list, scores_list, labels_list, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr) except Exception as e: log(f" WBF skip {Path(img_path).name}: {e}") ensemble_preds[img_path] = (np.zeros((0, 4)), np.zeros(0)) continue if len(wbf_boxes) == 0: ensemble_preds[img_path] = (np.zeros((0, 4)), np.zeros(0)) continue wbf_boxes_px = np.asarray(wbf_boxes) * np.array([W, H, W, H]) ensemble_preds[img_path] = (wbf_boxes_px, np.asarray(wbf_scores)) # GT loader (YOLO normalized format) LBL_DIR = ROOT / 'hitit_ocr/data/detection/labels/all' def load_gt(img_path): stem = Path(img_path).stem lbl = LBL_DIR / f'{stem}.txt' if not lbl.exists(): return np.zeros((0, 4)) boxes = [] with Image.open(img_path) as im: W, H = im.size for ln in lbl.read_text().splitlines(): p = ln.split() if len(p) < 5: continue cx, cy, w, h = float(p[1]), float(p[2]), float(p[3]), float(p[4]) x1 = (cx - w/2) * W; y1 = (cy - h/2) * H x2 = (cx + w/2) * W; y2 = (cy + h/2) * H boxes.append([x1, y1, x2, y2]) return np.array(boxes) if boxes else np.zeros((0, 4)) def iou_matrix(a, b): if len(a) == 0 or len(b) == 0: return np.zeros((len(a), len(b))) ax1, ay1, ax2, ay2 = a[:,0:1], a[:,1:2], a[:,2:3], a[:,3:4] bx1, by1, bx2, by2 = b[:,0], b[:,1], b[:,2], b[:,3] inter_x1 = np.maximum(ax1, bx1); inter_y1 = np.maximum(ay1, by1) inter_x2 = np.minimum(ax2, bx2); inter_y2 = np.minimum(ay2, by2) iw = np.clip(inter_x2-inter_x1, 0, None); ih = np.clip(inter_y2-inter_y1, 0, None) inter = iw * ih a_area = (ax2-ax1) * (ay2-ay1); b_area = (bx2-bx1) * (by2-by1) union = a_area + b_area - inter return inter / np.clip(union, 1e-9, None) def compute_map(pred_dict, iou_threshs=None): if iou_threshs is None: iou_threshs = np.arange(0.5, 1.0, 0.05) # Per-class (single class). Collect TP/FP at each conf. per_iou_aps = [] for iou_thr in iou_threshs: scores_all = []; tp_all = []; n_gt = 0 for img_path, (boxes, scores) in pred_dict.items(): gt = load_gt(img_path); n_gt += len(gt) if len(boxes) == 0: continue ious = iou_matrix(boxes, gt) # (P, G) # sort preds by score desc order = np.argsort(-scores) used_gt = np.zeros(len(gt), dtype=bool) for i in order: scores_all.append(scores[i]) if len(gt) == 0: tp_all.append(0); continue j = np.argmax(ious[i]) if ious[i, j] >= iou_thr and not used_gt[j]: tp_all.append(1); used_gt[j] = True else: tp_all.append(0) if not scores_all or n_gt == 0: per_iou_aps.append(0.0); continue scores_all = np.array(scores_all); tp_all = np.array(tp_all) order = np.argsort(-scores_all) tp_sorted = tp_all[order] cum_tp = np.cumsum(tp_sorted); cum_fp = np.cumsum(1 - tp_sorted) recall = cum_tp / n_gt precision = cum_tp / np.clip(cum_tp + cum_fp, 1, None) # 101-point AP ap = 0.0 for r in np.linspace(0, 1, 101): p = precision[recall >= r].max() if (recall >= r).any() else 0.0 ap += p ap /= 101 per_iou_aps.append(ap) return per_iou_aps log("Computing mAP for own-fold predictions (baseline)...") own_aps = compute_map(own_preds) log(f" own mAP50={own_aps[0]:.4f} mAP50-95={np.mean(own_aps):.4f}") log("Computing mAP for WBF ensemble...") ens_aps = compute_map(ensemble_preds) log(f" WBF mAP50={ens_aps[0]:.4f} mAP50-95={np.mean(ens_aps):.4f}") out = { 'n_eval_images': len(all_imgs), 'n_own_pred_images': len(own_preds), 'n_ensemble_pred_images': len(ensemble_preds), 'own_fold': {'mAP50': float(own_aps[0]), 'mAP50_95': float(np.mean(own_aps)), 'per_iou': [float(x) for x in own_aps]}, 'wbf_ensemble': {'mAP50': float(ens_aps[0]), 'mAP50_95': float(np.mean(ens_aps)), 'per_iou': [float(x) for x in ens_aps]}, } out_path = WEIGHTS_ROOT.parent.parent.parent / 'wbf_ensemble_eval.json' out_path.write_text(json.dumps(out, indent=2)) log(f"Wrote {out_path}") if __name__ == '__main__': main()