| |
| """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 |
|
|
| |
| 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)}") |
|
|
| |
| 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}") |
|
|
| |
| own_preds = {} |
| ensemble_preds = {} |
| 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] |
| |
| 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]) |
| 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)) |
|
|
| |
| 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_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) |
| |
| 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) |
| |
| 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() |
|
|
