#!/usr/bin/env python # train.py import sys import numpy as np import pickle from scipy.fft import idct from sklearn.ensemble import HistGradientBoostingClassifier from sklearn.utils.class_weight import compute_class_weight IGNORE_LABELS = {1, 4, 7} NON_CANCER = [0, 2] CANCER = [3, 5, 6] BCKG = [8] DCT_SIZE = 32 N_CHANNELS = 3 N_DCT_COEFFS = DCT_SIZE * DCT_SIZE # ── Features ────────────────────────────────────────────────────────────────── def apply_idct_2d(block): return idct(idct(block, axis=0, norm='ortho'), axis=1, norm='ortho') def dct_zone_energies(block): total = np.sum(block**2) + 1e-8 z1 = np.sum(block[:4, :4 ]**2) z2 = np.sum(block[:8, :8 ]**2) - z1 z3 = np.sum(block[:16, :16]**2) - z1 - z2 z4 = np.sum(block**2) - z1 - z2 - z3 diag = sum(block[i, i]**2 for i in range(min(block.shape))) return [z1/total, z2/total, z3/total, z4/total, diag/total] def spatial_texture(spatial): padded = np.pad(spatial, 1, mode='reflect') diffs = [] for dy in [-1, 0, 1]: for dx in [-1, 0, 1]: if dy == 0 and dx == 0: continue neighbor = padded[1+dy:1+dy+DCT_SIZE, 1+dx:1+dx+DCT_SIZE] diffs.append(neighbor - spatial) lbp_var = np.var(np.stack(diffs, axis=0), axis=0) return [float(np.mean(lbp_var)), float(np.percentile(lbp_var, 75)), float(np.percentile(lbp_var, 95))] def build_features_row(row): spatial_channels = [] eng = [] for c in range(N_CHANNELS): block = row[c * N_DCT_COEFFS:(c+1) * N_DCT_COEFFS].reshape(DCT_SIZE, DCT_SIZE) spatial = apply_idct_2d(block) spatial_channels.append(spatial) eng += [float(block[0, 0])] eng += dct_zone_energies(block) flat = spatial.flatten() eng += [float(np.mean(flat)), float(np.std(flat)), float(np.percentile(flat, 10)), float(np.percentile(flat, 90)), float(np.percentile(flat, 25)), float(np.percentile(flat, 75))] # Higher order moments — key for dcis mean = np.mean(flat) std = np.std(flat) + 1e-8 eng += [float(np.mean(((flat - mean)/std)**3)), # skewness float(np.mean(((flat - mean)/std)**4))] # kurtosis gy, gx = np.gradient(spatial) grad_mag = np.sqrt(gx**2 + gy**2) eng += [float(np.mean(grad_mag)), float(np.std(grad_mag)), float(np.max(grad_mag)), float(np.percentile(grad_mag, 90))] angles = np.arctan2(gy, gx) hist, _ = np.histogram(angles, bins=8, range=(-np.pi, np.pi), weights=grad_mag) hist = hist / (hist.sum() + 1e-8) eng += hist.tolist() eng += spatial_texture(spatial) # Radial DCT energy bands cy, cx = DCT_SIZE // 2, DCT_SIZE // 2 yy, xx = np.ogrid[:DCT_SIZE, :DCT_SIZE] radius = np.sqrt((yy - cy)**2 + (xx - cx)**2) total_e = float(np.sum(block**2)) + 1e-8 for r0, r1 in [(0, 4), (4, 8), (8, 12), (12, 16)]: mask = (radius >= r0) & (radius < r1) eng.append(float(np.sum(block[mask]**2)) / total_e) # Center vs border contrast h, w = spatial.shape center = spatial[h//4:3*h//4, w//4:3*w//4] border_mask = np.ones_like(spatial, dtype=bool) border_mask[h//4:3*h//4, w//4:3*w//4] = False eng += [float(np.mean(center) - np.mean(spatial[border_mask])), float(np.std(center)), float(np.std(spatial[border_mask]))] c0, c1, c2 = [ch.flatten() for ch in spatial_channels] eng += [ float(np.corrcoef(c0, c1)[0, 1]), float(np.corrcoef(c1, c2)[0, 1]), float(np.corrcoef(c0, c2)[0, 1]), float(np.mean(np.abs(c0 - c1))), float(np.mean(np.abs(c1 - c2))), float(np.mean(np.abs(c0 - c2))), float(np.mean(np.abs(c0)) / (np.mean(np.abs(c1)) + 1e-8)), float(np.mean(np.abs(c1)) / (np.mean(np.abs(c2)) + 1e-8)), ] spatial_flat = np.concatenate([c0, c1, c2]) eng += [float(np.mean(np.abs(spatial_flat))), float(np.std(np.abs(spatial_flat)))] return np.concatenate([row, spatial_flat, np.array(eng, dtype=np.float32)]) def build_features(X): return np.array([build_features_row(r) for r in X], dtype=np.float32) # ── Data ────────────────────────────────────────────────────────────────────── def load_csv(path): labels, feats = [], [] with open(path) as f: lines = f.readlines() for line in lines[1:]: p = line.strip().split(',') labels.append(int(float(p[0]))) feats.append(np.array(p[1:], dtype=np.float32)) return np.array(labels), np.array(feats) def filter_labels(y, X): mask = ~np.isin(y, list(IGNORE_LABELS)) return y[mask], X[mask] # ── Training ────────────────────────────────────────────────────────────────── def train_weighted(Xf, y, extra_weights=None, n_iter=300, lr=0.05, depth=10): unique_classes = np.unique(y) class_map = {c: i for i, c in enumerate(unique_classes)} inv_map = {i: c for c, i in class_map.items()} y_mapped = np.array([class_map[v] for v in y]) classes = np.unique(y_mapped) weights = compute_class_weight('balanced', classes=classes, y=y_mapped) cw = {} for i, c in enumerate(classes): orig = inv_map[c] w = weights[i] if extra_weights and orig in extra_weights: w *= extra_weights[orig] cw[c] = w model = HistGradientBoostingClassifier( max_depth=depth, max_iter=n_iter, learning_rate=lr, class_weight=cw, random_state=42, min_samples_leaf=12, l2_regularization=0.1, verbose=0, ) model.fit(Xf, y_mapped) return model, class_map, inv_map # ── Main ────────────────────────────────────────────────────────────────────── def main(argv): train_csv, dev_csv, out_file = argv[1], argv[2], argv[3] y_tr, X_tr = load_csv(train_csv) y_tr, X_tr = filter_labels(y_tr, X_tr) y_dv, X_dv = load_csv(dev_csv) y_dv, X_dv = filter_labels(y_dv, X_dv) y = np.concatenate([y_tr, y_dv]) X = np.concatenate([X_tr, X_dv]) print(f"Dataset: {len(y)} samples (train + dev)") print(f"Label counts: { {int(v): int(np.sum(y==v)) for v in np.unique(y)} }") print("Building features...") Xf = build_features(X) print(f"Feature shape: {Xf.shape}") # ── Model A: hierarchy, optimized for norm/nneo/infl/indc ───────────── print("Training Model A (hierarchical)...") print(" Stage 1...") y_s1 = np.where(np.isin(y, BCKG), 2, np.where(np.isin(y, NON_CANCER), 0, 1)) model_s1, _, inv_s1 = train_weighted(Xf, y_s1, extra_weights={2: 0.4}) print(" Stage 2a (norm vs nneo)...") mask_nc = np.isin(y, NON_CANCER) model_nc, _, inv_nc = train_weighted( Xf[mask_nc], y[mask_nc], extra_weights={2: 2.5} ) print(" Stage 2b (cancer subtypes)...") mask_c = np.isin(y, CANCER) model_c, _, inv_c = train_weighted( Xf[mask_c], y[mask_c], extra_weights={5: 6.0, 6: 3.0, 3: 1.5}, n_iter=400, depth=12 ) print(" Stage 2c (bckg)...") mask_bg = np.isin(y, BCKG) model_bg, _, inv_bg = train_weighted(Xf[mask_bg], y[mask_bg]) # ── Model B: flat, aggressively biased toward dcis ──────────────────── # Trains on all 6 classes at once. dcis at 15x forces it to carve # out a dcis region even at the cost of nneo false positives. # We only use Model B's dcis probability at predict time. print("Training Model B (flat dcis-aggressive)...") model_flat, _, inv_flat = train_weighted( Xf, y, extra_weights={5: 15.0, 6: 4.0, 3: 2.0, 2: 1.5, 0: 1.0, 8: 0.5}, n_iter=400, depth=12 ) print("All training done.") with open(out_file, 'wb') as f: pickle.dump({ 'stage1': model_s1, 's1_inv': inv_s1, 'nc': model_nc, 'nc_inv': inv_nc, 'cancer': model_c, 'c_inv': inv_c, 'bckg': model_bg, 'bg_inv': inv_bg, 'flat': model_flat,'flat_inv': inv_flat, }, f) print("Saved model.") if __name__ == "__main__": main(sys.argv)