code/src/wbf_ensemble_eval.py

#!/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()