File size: 4,028 Bytes
f211247 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | #!/usr/bin/env python3
"""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)
# Build confusion matrix
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
# Row sums (per-class n)
n_per = confusion.sum(-1).tolist()
# Find confusion pairs
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))
# Group into merge groups via union-find
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) # keep larger idx as root
groups = defaultdict(list)
for i in range(C): groups[find(i)].append(i)
# Build merge_map: label → merged_label
merge_map = {}
for root, members in groups.items():
if len(members) == 1:
lab = idx_to_label[root]
merge_map[lab] = lab
else:
# Sort by n desc, compound name = "A|B|C" (alphabetic stable)
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
# Apply to manifest
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()
|