| |
| """Auto-merge highly confused class pairs. |
| |
| Reads probs from v4 eval; finds pairs where >=30% of true-class-A samples |
| are predicted class-B (or vice versa). Merges A+B into A|B compound class. |
| |
| Output: manifest with unified_label replaced by merged form, plus label_merge_map.json. |
| """ |
| import json, argparse |
| from pathlib import Path |
| from collections import defaultdict, Counter |
| import torch |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--probs', required=True, help='pt with probs/targets/label_to_idx') |
| ap.add_argument('--manifest', required=True) |
| ap.add_argument('--output-manifest', required=True) |
| ap.add_argument('--output-map', required=True) |
| ap.add_argument('--confusion-frac', type=float, default=0.25, |
| help='Merge if A→B confusion >= frac of A samples') |
| ap.add_argument('--min-sym-frac', type=float, default=0.15, |
| help='Also require B→A confusion >= frac (symmetric)') |
| args = ap.parse_args() |
|
|
| d = torch.load(args.probs, map_location='cpu', weights_only=False) |
| probs = d['probs']; targets = d['targets'] |
| label_to_idx = d['label_to_idx']; idx_to_label = {v: k for k, v in label_to_idx.items()} |
|
|
| pred = probs.argmax(-1) |
| |
| C = len(label_to_idx) |
| confusion = torch.zeros(C, C, dtype=torch.long) |
| for i in range(len(targets)): |
| confusion[int(targets[i]), int(pred[i])] += 1 |
|
|
| |
| n_per = confusion.sum(-1).tolist() |
|
|
| |
| merge_pairs = [] |
| for a in range(C): |
| if n_per[a] < 5: continue |
| for b in range(a+1, C): |
| if n_per[b] < 5: continue |
| ab = confusion[a, b].item() / max(1, n_per[a]) |
| ba = confusion[b, a].item() / max(1, n_per[b]) |
| if ab >= args.confusion_frac and ba >= args.min_sym_frac: |
| merge_pairs.append((a, b, ab, ba)) |
| elif ba >= args.confusion_frac and ab >= args.min_sym_frac: |
| merge_pairs.append((b, a, ba, ab)) |
|
|
| |
| parent = list(range(C)) |
| def find(x): |
| while parent[x] != x: x = parent[x] |
| return x |
| for a, b, _, _ in merge_pairs: |
| ra, rb = find(a), find(b) |
| if ra != rb: parent[min(ra, rb)] = max(ra, rb) |
| groups = defaultdict(list) |
| for i in range(C): groups[find(i)].append(i) |
|
|
| |
| merge_map = {} |
| for root, members in groups.items(): |
| if len(members) == 1: |
| lab = idx_to_label[root] |
| merge_map[lab] = lab |
| else: |
| |
| sorted_labs = sorted([idx_to_label[m] for m in members]) |
| merged = '|'.join(sorted_labs) |
| for m in members: |
| merge_map[idx_to_label[m]] = merged |
|
|
| |
| n_merged = sum(1 for v in merge_map.values() if '|' in v) |
| n_groups = len(set(merge_map.values())) |
| print(f"Confusion pairs found: {len(merge_pairs)}") |
| print(f"Original classes: {C}, After merge: {n_groups}, Labels in merged groups: {n_merged}") |
|
|
| records = [json.loads(l) for l in open(args.manifest)] |
| with open(args.output_manifest, 'w') as f: |
| for r in records: |
| lab = r.get('unified_label') |
| if lab and lab in merge_map: |
| r['unified_label'] = merge_map[lab] |
| r['original_label'] = lab |
| f.write(json.dumps(r) + '\n') |
|
|
| Path(args.output_map).parent.mkdir(parents=True, exist_ok=True) |
| json.dump({'merge_map': merge_map, 'n_original': C, 'n_merged': n_groups, |
| 'pairs': [{'a': idx_to_label[a], 'b': idx_to_label[b], |
| 'ab': ab, 'ba': ba} for a, b, ab, ba in merge_pairs]}, |
| open(args.output_map, 'w'), indent=2) |
| print(f"Saved: {args.output_manifest}, map: {args.output_map}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
