| """Portable frozen TESSERA optical feature extraction, readout fitting and score replay. |
| |
| Published historical research scripts are preserved separately. This entry point |
| never trains the encoder or chooses parameters using test labels. |
| """ |
| import argparse |
| import hashlib |
| import json |
| from datetime import timedelta |
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
| MAIN = Path(__file__).resolve().parents[1] |
| REPO = MAIN.parent |
| BUNDLE = MAIN / 'checkpoints/32X32_tessera_transfer_seed3407' |
| WORK = REPO / 'ablations/10_plateau_breakthrough/workspace' |
| ENCODER_SHA = '92619d4ffc2936895f145fdcf710140b58c23c860e02f01873828b10c6958d95' |
| READOUT_SHA = '5959d5f58fab2dfdff458281081f531b626965a8283458aaae61aa9857563753' |
| BAND_ORDER = [2, 0, 1, 6, 7, 3, 4, 5, 8, 9] |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for block in iter(lambda: stream.read(8 * 1024 * 1024), b''): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def verify(): |
| for filename, expected in [('frozen_satellite_encoder.pt', ENCODER_SHA), |
| ('optical_gaussian_classifier.npz', READOUT_SHA)]: |
| actual = digest(BUNDLE / filename) |
| if actual != expected: |
| raise RuntimeError(f'Hash mismatch: {filename}') |
| return {'encoder_sha256': ENCODER_SHA, 'readout_sha256': READOUT_SHA, 'passed': True} |
|
|
|
|
| def prepare(x, pos, days): |
| """Undo legacy loader normalization, preserve real observations and map bands.""" |
| import tessera_student as student |
| from dataset import get_norm, REF_DATE |
| assert len(x) == 1 and x.shape[2] == 10 |
| valid = pos[0].gt(0) |
| valid[0] = True |
| n = int(valid.sum()) |
| assert n > 0 and valid[:n].all() and not valid[n:].any() |
| mean, std = get_norm() |
| raw = x[0, :n].float() * torch.as_tensor(std + 1e-8, device=x.device)[None, :, None, None] |
| raw = raw + torch.as_tensor(mean, device=x.device)[None, :, None, None] |
| raw = raw[:, BAND_ORDER].permute(2, 3, 0, 1).reshape(-1, n, 10) |
| doy = torch.tensor([(REF_DATE + timedelta(days=int(d))).timetuple().tm_yday |
| for d in days[0, :n].tolist()], device=x.device) |
| remain = ((n + 7) // 8) * 8 - n |
| if remain: |
| groups = np.array_split(np.arange(n), remain) |
| fill = np.array([g[len(g) // 2] for g in groups]) if remain <= n else np.arange(remain) % n |
| indices = torch.as_tensor(np.concatenate([np.arange(n), fill]), device=x.device) |
| raw, doy = raw[:, indices], doy[indices] |
| bands = (raw - torch.as_tensor(student.S2_BAND_MEAN, device=x.device)) / ( |
| torch.as_tensor(student.S2_BAND_STD, device=x.device) + 1e-9) |
| return torch.cat([bands, doy[None, :, None].expand(len(bands), -1, 1)], -1) |
|
|
|
|
| @torch.inference_mode() |
| def optical_features(model, inputs): |
| chunks = [] |
| for start in range(0, len(inputs), 512): |
| with torch.autocast('cuda', dtype=torch.float16, enabled=inputs.is_cuda): |
| features = model.s2_backbone(inputs[start:start + 512]) |
| chunks.append(features.float().cpu().numpy()) |
| return np.concatenate(chunks)[:, None, :] |
|
|
|
|
| def predict(features, params): |
| """Shared-covariance Gaussian discriminant scores, then calibrated softmax.""" |
| x = (np.asarray(features[:, 0], np.float32) - params['mean']) / params['std'] |
| scores = x @ params['weight'] + params['bias'] |
| p = np.zeros((len(x), 20), np.float32) |
| p[:, 1:19] = torch.from_numpy(np.asarray(scores / float(params['temperature']), np.float32)).softmax(1).numpy() |
| return p |
|
|
|
|
| def fit(features, labels, output): |
| """Refit only the validation-frozen shrinkage=.1, temperature=4 configuration.""" |
| f = np.load(features, mmap_mode='r'); y = np.load(labels) |
| keep = (y > 0) & (y < 19); y = y[keep].astype(np.int64) - 1 |
| x = np.asarray(f[:, 0], np.float32)[keep] |
| mean, std = x.mean(0), x.std(0).clip(.01) |
| x = (x - mean) / std |
| centers = np.stack([x[y == c].mean(0) for c in range(18)]) |
| residual = x - centers[y] |
| covariance = residual.T @ residual / len(x) |
| values, vectors = np.linalg.eigh(covariance.astype(np.float64)) |
| inverse = (vectors / (.9 * values + .1 * values.mean()).clip(1e-6)) @ vectors.T |
| weight = (inverse @ centers.T).astype(np.float32) |
| bias = -.5 * (centers * weight.T).sum(1) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, depth=0, mean=mean, std=std, type='linear', weight=weight, bias=bias, temperature=4.) |
| return {'foreground_training_pixel_entries': len(x), 'output': str(output), 'encoder_updated': False} |
|
|
|
|
| def confusion(labels, prediction): |
| valid = (labels > 0) & (labels < 19) |
| return np.bincount(20 * labels[valid].astype(np.int64) + prediction[valid], minlength=400).reshape(20, 20) |
|
|
|
|
| def metrics(cm): |
| tp = cm.diagonal()[1:19].astype(float); gt = cm.sum(1)[1:19]; pr = cm.sum(0)[1:19] |
| iou = tp / np.maximum(gt + pr - tp, 1) |
| f1 = 2 * tp / np.maximum(gt + pr, 1) |
| total = cm.sum(); oa = cm.trace() / max(total, 1) |
| chance = np.dot(cm.sum(0).astype(float), cm.sum(1)) / max(total * total, 1) |
| return dict(miou=float(iou.mean()), oa=float(oa), macro_f1=float(f1.mean()), |
| kappa=float((oa - chance) / max(1 - chance, 1e-12)), valid_pixels=int(total), |
| per_class_iou=iou.tolist(), confusion_matrix=cm.tolist()) |
|
|
|
|
| def evaluate_cached(check_raw_labels=False): |
| """Replay nested probabilities exactly; optionally check every raw test target.""" |
| audit = json.loads((MAIN / 'configurations/32X32_tessera_transfer_seed3407/completion_audit.json').read_text()) |
| protocol = json.loads((WORK / 'research_goal_52/PROTOCOL.json').read_text()) |
| labels = np.load(WORK / 'convnext_research_20261010/test_labels.npy', mmap_mode='r') |
| paths = ['convnext_research_20261010/recommended_test_probabilities.npy', |
| 'score_breakthrough_20261010/gaussian_test_probabilities.npy', |
| 'score_breakthrough_20261010/structured_test_probabilities.npy', |
| 'tessera_transfer_research_20261010/32X32_pretrained_optical_lda_seed3407/test_probabilities.npy'] |
| arrays = [np.load(WORK / p, mmap_mode='r') for p in paths] |
| ids = protocol['roles']['test']['ordered_ids'] |
| assert len(ids) == len(labels) == 482 and len(set(ids)) == 482 |
| cm = np.zeros((20, 20), np.int64) |
| if check_raw_labels: |
| import config as C |
| for i, pid in enumerate(ids): |
| if check_raw_labels: |
| assert np.array_equal(labels[i], np.load(C.ANNOT_DIR / f'TARGET_{pid}.npy')[0]) |
| p0, pg, ps, pt = [a[i] for a in arrays] |
| assert all(p.shape == (128, 128, 20) and np.isfinite(p).all() for p in [p0, pg, ps, pt]) |
| incumbent = .65 * p0 + .25 * pg + .1 * ps |
| cm += confusion(labels[i], (.5 * incumbent + .5 * pt).argmax(-1)) |
| result = metrics(cm) |
| expected = audit['results']['pretrained_optical_lda']['test'] |
| assert result['confusion_matrix'] == expected['confusion_matrix'] |
| return {'passed': True, 'mode': 'saved_probability_replay', 'raw_labels_checked': check_raw_labels, |
| 'patches': 482, 'metrics': result} |
|
|
|
|
| @torch.inference_mode() |
| def extract(split, output, device): |
| """Fresh raw-data optical inference. Train/val preserve original batch/crop order.""" |
| import tessera_student as student |
| import config as C |
| from torch.utils.data import DataLoader |
| from dataset import PASTISPatchDataset, pad_collate, crop_ij, train_patch_ids, val_patch_ids, test_patch_ids |
| verify(); torch.set_num_threads(4); torch.manual_seed(3407) |
| torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True |
| model = student.load_model(str(BUNDLE / 'frozen_satellite_encoder.pt'), torch.device(device)).eval().requires_grad_(False) |
| ids = {'train': train_patch_ids, 'val': val_patch_ids, 'test': test_patch_ids}[split]() |
| dl = DataLoader(PASTISPatchDataset(ids), batch_size=1 if split == 'test' else 2, |
| shuffle=False, num_workers=2, collate_fn=pad_collate) |
| output.mkdir(parents=True, exist_ok=True) |
| count = len(ids) * (16384 if split == 'test' else 2048) |
| fw = np.lib.format.open_memmap(output / 'features.npy', mode='w+', dtype=np.float16, shape=(count, 1, 256)) |
| yw = np.lib.format.open_memmap(output / 'labels.npy', mode='w+', dtype=np.uint8, shape=(count,)) |
| pw = np.lib.format.open_memmap(output / 'probabilities.npy', mode='w+', dtype=np.float32, shape=(count, 20)) |
| params = np.load(BUNDLE / 'optical_gaussian_classifier.npz'); offset = 0; cm = np.zeros((20, 20), np.int64) |
| for bi, ((x, pos, days), y) in enumerate(dl): |
| for i, j in C.FT_CROP_IJ if split != 'test' else [(0, 0)]: |
| crop, target = crop_ij(x, y, i, j, 4 if split != 'test' else 1) |
| for sample in range(len(crop)): |
| inputs = prepare(crop[sample:sample+1].to(device), pos[sample:sample+1].to(device), days[sample:sample+1]) |
| f = optical_features(model, inputs); n = len(f) |
| |
| scored = f if split == 'test' else f.astype(np.float16).astype(np.float32) |
| p = predict(scored, params); gt = target[sample].numpy().ravel() |
| fw[offset:offset+n] = f; yw[offset:offset+n] = gt; pw[offset:offset+n] = p |
| cm += confusion(gt, p.argmax(-1)); offset += n |
| if bi % 25 == 0: print(split, bi, len(dl), flush=True) |
| assert offset == count |
| fw.flush(); yw.flush(); pw.flush() |
| return {'split': split, 'patch_entries': len(ids), 'pixel_entries': count, 'standalone_metrics': metrics(cm)} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('action', choices=['verify', 'evaluate-cached', 'extract', 'fit']) |
| parser.add_argument('--check-raw-labels', action='store_true') |
| parser.add_argument('--split', choices=['train', 'val', 'test'], default='val') |
| parser.add_argument('--device', default='cuda') |
| parser.add_argument('--output', type=Path, default=MAIN / 'reproduction') |
| parser.add_argument('--features', type=Path); parser.add_argument('--labels', type=Path) |
| args = parser.parse_args() |
| torch.set_num_threads(4) |
| if args.action == 'verify': result = verify() |
| elif args.action == 'evaluate-cached': result = evaluate_cached(args.check_raw_labels) |
| elif args.action == 'extract': result = extract(args.split, args.output, args.device) |
| else: |
| if args.features is None or args.labels is None: parser.error('fit requires --features and --labels') |
| result = fit(args.features, args.labels, args.output / 'refitted_readout.npz') |
| print(json.dumps(result, indent=2)) |
| if args.action != 'verify': |
| args.output.mkdir(parents=True, exist_ok=True) |
| (args.output / f'{args.action}_{args.split}.json').write_text(json.dumps(result, indent=2)) |
|
|
|
|
| if __name__ == '__main__': main() |
|
|
