hitit-cuneiform-ocr / code /src /enhancements /active_learning_export.py
savastakan's picture
Initial upload: code + 5 record checkpoints + fuse
f211247 verified
Raw
History Blame Contribute Delete
4.22 kB
#!/usr/bin/env python3
"""Export margin-low samples for expert review.
Rows to review (margin = p_top1 - p_top2):
margin < 0.2 → most informative
entropy > 2.5 → high uncertainty
flipped by pair-MLP → suspected misclassification
Output: CSV + image symlinks in review folder.
"""
import os, json, csv, argparse, shutil
from pathlib import Path
import torch
ROOT = Path("/arf/scratch/stakan/hitit-proje")
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--probs', required=True)
ap.add_argument('--manifest', required=True)
ap.add_argument('--label-to-idx', required=True,
help='any train ckpt providing label_to_idx')
ap.add_argument('--val-fold', type=int, default=0)
ap.add_argument('--margin-threshold', type=float, default=0.2)
ap.add_argument('--entropy-threshold', type=float, default=2.5)
ap.add_argument('--output-csv', required=True)
ap.add_argument('--output-dir', default=None,
help='If set, symlink low-margin images here for UI review')
args = ap.parse_args()
d = torch.load(args.probs, map_location='cpu', weights_only=False)
probs, targets = d['probs'], d['targets']
ck = torch.load(args.label_to_idx, map_location='cpu', weights_only=False)
label_to_idx = ck['label_to_idx']; idx_to_label = {v: k for k, v in label_to_idx.items()}
# Respect min_samples filter (same as prototype_net / training dataset)
from collections import Counter
cls_count = Counter()
for line in open(args.manifest):
r = json.loads(line)
if r.get('task') != 'classification' or not r.get('unified_label'): continue
cls_count[r['unified_label']] += 1
MIN_SAMPLES = 10 # same default used by training/prototype_net
records = []
with open(args.manifest) as f:
for line in f:
r = json.loads(line)
if r.get('task') != 'classification': continue
if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
if not r.get('path') or r.get('storage') != 'fs': continue
if r.get('integrity_ok') is False: continue
if cls_count[r['unified_label']] < MIN_SAMPLES: continue
if r.get('tablet_view_fold', 0) != args.val_fold: continue
records.append(r)
if len(records) != probs.size(0):
print(f"WARNING: records={len(records)} vs probs={probs.size(0)}; truncating to min")
records = records[:probs.size(0)]
probs = probs[:len(records)]
targets = targets[:len(records)]
top2 = probs.topk(2, dim=-1)
margin = top2.values[:, 0] - top2.values[:, 1]
entropy = -(probs * probs.clamp_min(1e-9).log()).sum(-1)
rows = []
for i in range(len(records)):
if margin[i] < args.margin_threshold or entropy[i] > args.entropy_threshold:
rows.append({
'idx': i,
'path': records[i]['path'],
'true_label': records[i]['unified_label'],
'pred_label': idx_to_label[int(top2.indices[i, 0])],
'alt_label': idx_to_label[int(top2.indices[i, 1])],
'p_pred': float(top2.values[i, 0]),
'p_alt': float(top2.values[i, 1]),
'margin': float(margin[i]),
'entropy': float(entropy[i]),
'misclassified': int(top2.indices[i, 0]) != int(targets[i]),
})
rows.sort(key=lambda r: r['margin'])
Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True)
with open(args.output_csv, 'w') as f:
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()) if rows else ['idx'])
w.writeheader()
w.writerows(rows)
print(f"Exported {len(rows)} rows for review → {args.output_csv}")
if args.output_dir:
d = Path(args.output_dir); d.mkdir(parents=True, exist_ok=True)
for r in rows[:300]: # cap for sanity
src = Path(r['path']); dst = d / f"{r['idx']:05d}_{r['true_label']}_vs_{r['pred_label']}.png"
if src.exists() and not dst.exists():
try: shutil.copy(src, dst)
except Exception: pass
if __name__ == '__main__':
main()