Add inference pipeline (comments stripped)

predict_ensemble.py

@@ -0,0 +1,202 @@
+import argparse
+import os
+import json
+from pathlib import Path
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import cv2
+from train_global_unet import UNet, encode_scribble, list_test_pairs, list_train_pairs, load_palette, evaluate_predictions, tta_predict, TRAIN_H, TRAIN_W, ORIG_H, ORIG_W, TEST_PRED
+
+def predict_threshold(prob, sc, predictor):
+    if predictor is None:
+        return 0.5
+    fg_scrib = (sc == 1).sum()
+    bg_scrib = (sc == 0).sum()
+    scrib_fg_ratio = fg_scrib / (fg_scrib + bg_scrib) if fg_scrib + bg_scrib > 0 else 0.5
+    pred_fg_frac = (prob > 0.5).mean()
+    prob_mean = prob.mean()
+    prob_std = prob.std()
+    p_clip = np.clip(prob, 1e-06, 1 - 1e-06)
+    entropy = -(p_clip * np.log(p_clip) + (1 - p_clip) * np.log(1 - p_clip)).mean()
+    feats = np.array([scrib_fg_ratio, pred_fg_frac, prob_mean, prob_std, entropy])
+    t = float(feats @ np.array(predictor['weights']) + predictor['bias'])
+    return float(np.clip(t, predictor['clip'][0], predictor['clip'][1]))
+
+def load_threshold_predictor():
+    p = Path('threshold_predictor.json')
+    if not p.exists():
+        return None
+    return json.load(open(p))
+
+def load_models(ckpt_specs, device):
+    models = []
+    for ckpt_dir, base, seed in ckpt_specs:
+        for fd in sorted(Path(ckpt_dir).glob('fold_*')):
+            ckpt = fd / 'best.pth'
+            if not ckpt.exists():
+                continue
+            m = UNet(in_ch=5, base=base, out_ch=1).to(device)
+            m.load_state_dict(torch.load(ckpt, map_location=device))
+            m.eval()
+            models.append(m)
+            print(f'  loaded {ckpt}  (base={base}, seed={seed})')
+    return models
+
+def parse_specs(spec_strs):
+    out = []
+    for s in spec_strs:
+        parts = s.split(':')
+        path = parts[0]
+        base = int(parts[1])
+        seed = int(parts[2]) if len(parts) > 2 else 42
+        out.append((Path(path), base, seed))
+    return out
+
+def filter_small_components(pred, min_area=200):
+    n_fg, lab_fg, stats_fg, _ = cv2.connectedComponentsWithStats(pred, connectivity=8)
+    out = pred.copy()
+    for i in range(1, n_fg):
+        if stats_fg[i, cv2.CC_STAT_AREA] < min_area:
+            out[lab_fg == i] = 0
+    inv = 1 - out
+    n_bg, lab_bg, stats_bg, _ = cv2.connectedComponentsWithStats(inv, connectivity=8)
+    for i in range(1, n_bg):
+        if stats_bg[i, cv2.CC_STAT_AREA] < min_area:
+            out[lab_bg == i] = 1
+    return out
+
+def postprocess(pred, min_area=200, close_ks=9):
+    if close_ks > 1:
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_ks, close_ks))
+        pred = cv2.morphologyEx(pred, cv2.MORPH_CLOSE, kernel)
+    pred = filter_small_components(pred, min_area)
+    return pred
+
+def predict_test1(models, device):
+    palette = load_palette()
+    predictor = load_threshold_predictor()
+    if predictor is not None:
+        print(f'Using learned per-image threshold predictor (clip={predictor['clip']})')
+    test_pairs = list_test_pairs()
+    TEST_PRED.mkdir(parents=True, exist_ok=True)
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+    for stem, img_p, sc_p in test_pairs:
+        img = np.array(Image.open(img_p).convert('RGB'))
+        sc = np.array(Image.open(sc_p).convert('L'))
+        img_r = cv2.resize(img, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_LINEAR)
+        sc_r = cv2.resize(sc, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_NEAREST)
+        img_f = (img_r.astype(np.float32) / 255.0 - mean) / std
+        x = torch.cat([torch.from_numpy(img_f.transpose(2, 0, 1)), torch.from_numpy(encode_scribble(sc_r))], 0).unsqueeze(0).to(device)
+        prob_sum = None
+        for m in models:
+            p = tta_predict(m, x, device, scales=(0.7, 1.0, 1.3))
+            prob_sum = p if prob_sum is None else prob_sum + p
+        prob = prob_sum / len(models)
+        prob_full = cv2.resize(prob, (ORIG_W, ORIG_H), interpolation=cv2.INTER_LINEAR)
+        thresh = predict_threshold(prob_full, sc, predictor)
+        pred = (prob_full > thresh).astype(np.uint8)
+        pred = postprocess(pred, min_area=200, close_ks=9)
+        pred[sc == 0] = 0
+        pred[sc == 1] = 1
+        out_img = Image.fromarray(pred.astype(np.uint8), mode='P')
+        out_img.putpalette(palette)
+        out_img.save(TEST_PRED / f'{stem}.png')
+    print(f'Wrote {len(test_pairs)} predictions to {TEST_PRED}')
+
+def eval_oof_ensemble(ckpt_specs, folds, seed, device, save=False):
+    pairs = list_train_pairs()
+
+    def fold_assignment(seed_):
+        rng = np.random.RandomState(seed_)
+        idx = np.arange(len(pairs))
+        rng.shuffle(idx)
+        fold_arr_ = np.array_split(idx, folds)
+        return {ii: k for k in range(folds) for ii in fold_arr_[k]}
+    grouped = []
+    for ckpt_dir, base, ckpt_seed in ckpt_specs:
+        fold_models = {}
+        for fd in sorted(Path(ckpt_dir).glob('fold_*')):
+            ckpt = fd / 'best.pth'
+            if not ckpt.exists():
+                continue
+            k = int(fd.name.split('_')[1])
+            m = UNet(in_ch=5, base=base, out_ch=1).to(device)
+            m.load_state_dict(torch.load(ckpt, map_location=device))
+            m.eval()
+            fold_models[k] = m
+        fold_of_seed = fold_assignment(ckpt_seed)
+        grouped.append((fold_models, fold_of_seed))
+        print(f'  {ckpt_dir} has folds: {sorted(fold_models)}  (seed={ckpt_seed})')
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+    train_pred_dir = Path('dataset/train/predictions')
+    if save:
+        train_pred_dir.mkdir(exist_ok=True)
+        palette = load_palette()
+    predictor = load_threshold_predictor()
+    if predictor is not None:
+        print(f'Using learned per-image threshold predictor')
+    all_p, all_g = ([], [])
+    for i, (stem, img_p, sc_p, gt_p) in enumerate(pairs):
+        ensemble_models = []
+        for fold_models, fold_of_seed in grouped:
+            k_dir = fold_of_seed[i]
+            if k_dir in fold_models:
+                ensemble_models.append(fold_models[k_dir])
+        if not ensemble_models:
+            continue
+        img = np.array(Image.open(img_p).convert('RGB'))
+        sc = np.array(Image.open(sc_p).convert('L'))
+        gt = (np.array(Image.open(gt_p)) > 0).astype(np.uint8)
+        img_r = cv2.resize(img, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_LINEAR)
+        sc_r = cv2.resize(sc, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_NEAREST)
+        img_f = (img_r.astype(np.float32) / 255.0 - mean) / std
+        x = torch.cat([torch.from_numpy(img_f.transpose(2, 0, 1)), torch.from_numpy(encode_scribble(sc_r))], 0).unsqueeze(0).to(device)
+        prob_sum = None
+        for m in ensemble_models:
+            p = tta_predict(m, x, device, scales=(0.7, 1.0, 1.3))
+            prob_sum = p if prob_sum is None else prob_sum + p
+        prob = prob_sum / len(ensemble_models)
+        prob_full = cv2.resize(prob, (ORIG_W, ORIG_H), interpolation=cv2.INTER_LINEAR)
+        thresh = predict_threshold(prob_full, sc, predictor)
+        pred = (prob_full > thresh).astype(np.uint8)
+        pred = postprocess(pred, min_area=200, close_ks=9)
+        pred[sc == 0] = 0
+        pred[sc == 1] = 1
+        all_p.append(pred)
+        all_g.append(gt)
+        if save:
+            out_img = Image.fromarray(pred.astype(np.uint8), mode='P')
+            out_img.putpalette(palette)
+            out_img.save(train_pred_dir / f'{stem}.png')
+    bg, fg, miou = evaluate_predictions(all_p, all_g)
+    print(f'OOF ensemble: bg={bg:.4f}  fg={fg:.4f}  mIoU={miou:.4f}  (n={len(all_p)})')
+    if save:
+        print(f'Saved {len(all_p)} OOF predictions to {train_pred_dir}')
+    return miou
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--ckpt-dirs', nargs='+', required=True, help="One or more 'path:base' entries, e.g. 'runs_global_unet:48 runs_v2:64'")
+    p.add_argument('--gpu', type=int, default=0)
+    p.add_argument('--eval', action='store_true', help='Evaluate out-of-fold ensemble on training set instead of predicting test1')
+    p.add_argument('--save', action='store_true', help='With --eval, also save predictions to dataset/train/predictions/')
+    p.add_argument('--folds', type=int, default=5)
+    p.add_argument('--seed', type=int, default=42)
+    args = p.parse_args()
+    device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu')
+    ckpt_specs = parse_specs(args.ckpt_dirs)
+    if args.eval:
+        eval_oof_ensemble(ckpt_specs, args.folds, args.seed, device, save=args.save)
+    else:
+        models = load_models(ckpt_specs, device)
+        if not models:
+            print('No models found.')
+            return
+        print(f'Total models in ensemble: {len(models)}')
+        predict_test1(models, device)
+if __name__ == '__main__':
+    main()