| """Fresh raw-PASTIS inference for every component of the 58.6346% ensemble. |
| |
| Paths are explicit to support uploaded/local checkpoints. The published recipe |
| is fixed; this module never selects weights with test labels. Use the released |
| component hashes to reproduce the recorded system, not a newly trained replacement. |
| """ |
| import argparse |
| import gc |
| import json |
| import time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import config as C |
| from dataset import PASTISPatchDataset, pad_collate, val_patch_ids, test_patch_ids, crop_ij |
| from model import UTAE_ASPP, UTAEClassificationDual |
| from temporal_cnn import TemporalCNN |
| from training_feature_clusters import DepthCollector, probability |
| import foundation_transfer as F |
| import tessera_student as student |
|
|
|
|
| def load(path): |
| return torch.load(path, map_location='cpu', weights_only=False) |
|
|
|
|
| def spatial_network(path): |
| net = UTAEClassificationDual(UTAE_ASPP(n_channels=C.N_CHANNELS, d_model=C.D_MODEL, |
| n_heads=C.N_HEADS), d_model=C.D_MODEL, |
| num_classes=C.N_CLASSES) |
| net.load_state_dict(load(path)['model_state_dict'], strict=True) |
| return net.eval().requires_grad_(False) |
|
|
|
|
| def class_probabilities(logits): |
| b, _, h, w = logits.shape |
| result = torch.zeros((b, 20, h, w), device=logits.device) |
| result[:, 1:19] = logits.float().softmax(1) |
| return result |
|
|
|
|
| class CompleteEnsemble: |
| """Sequential GPU residency keeps the two spatial networks from overlapping.""" |
| def __init__(self, paths): |
| if not torch.cuda.is_available(): |
| raise RuntimeError('Fresh complete-ensemble inference requires CUDA.') |
| self.paths = {k: Path(v) for k, v in paths.items()} |
| for name, path in self.paths.items(): |
| if not path.is_file(): raise FileNotFoundError(f'{name}: {path}') |
| self.sam = spatial_network(self.paths['sam']) |
| self.structured = spatial_network(self.paths['structured']) |
| self.collector = DepthCollector(self.sam.utae) |
| self.cluster = load(self.paths['clusters']) |
| self.old = TemporalCNN(18).cuda().eval().requires_grad_(False) |
| self.old.load_state_dict(load(self.paths['old_cnn'])['state'], strict=True) |
| self.new = TemporalCNN(18).cuda().eval().requires_grad_(False) |
| self.new.load_state_dict(load(self.paths['new_cnn'])['state'], strict=True) |
| self.foundation = student.load_model(str(self.paths['tessera']), torch.device('cuda')).eval().requires_grad_(False) |
| self.original_gaussian = dict(np.load(self.paths['original_gaussian'])) |
| self.optical_gaussian = dict(np.load(self.paths['optical_gaussian'])) |
| self.mean, self.std = self.cluster['mean'].cuda(), self.cluster['std'].cuda() |
| self.cluster_state = ({k: v.cuda() for k, v in self.cluster['state'].items()} |
| if self.cluster['kind'] == 'linear' else self.cluster['state'].cuda()) |
|
|
| @torch.inference_mode() |
| def predict(self, x, pos, days, split, return_features=False): |
| """Return B×H×W×20 probabilities; labels are never an argument.""" |
| self.sam.cuda(); self.collector.begin(x, pos) |
| with torch.autocast('cuda', dtype=torch.float16): |
| output = self.sam(x, pos) |
| old_logits = self.old(x, pos) |
| new_logits = self.new(x, pos) |
| features = self.collector.end() |
| semantic = output['refined_logits'].float().softmax(1) |
| cluster = probability(features, self.cluster['depth'], self.mean, self.std, |
| self.cluster['kind'], self.cluster_state) |
| b, _, h, w = semantic.shape |
| cluster = torch.from_numpy(cluster).cuda().reshape(b, h, w, 20).permute(0, 3, 1, 2) |
| |
| spatial_cluster = .5 * semantic + .5 * cluster |
| old_reference = .75 * spatial_cluster + .25 * class_probabilities(old_logits) |
| reference = .75 * old_reference + .25 * class_probabilities(new_logits) |
| reference = reference.permute(0, 2, 3, 1).cpu().numpy() |
| del output, semantic, cluster, spatial_cluster, old_reference, old_logits, new_logits |
| self.sam.cpu(); torch.cuda.empty_cache() |
| |
| p = self.original_gaussian |
| z = (features[:, 0].astype(np.float32) - p['mean']) / p['std'] |
| scores = z @ p['weight'] + p['bias'] |
| gaussian = np.zeros((len(z), 20), np.float32) |
| gaussian[:, 1:19] = torch.from_numpy(scores / float(p['temperature'])).softmax(1).numpy() |
| gaussian = gaussian.reshape(b, h, w, 20) |
| self.structured.cuda() |
| with torch.autocast('cuda', dtype=torch.float16): |
| structured = self.structured(x, pos)['refined_logits'].float().softmax(1) |
| structured = structured.permute(0, 2, 3, 1).cpu().numpy() |
| self.structured.cpu(); torch.cuda.empty_cache() |
| optical = []; optical_features = [] |
| for i in range(b): |
| inputs = F.prepare(x[i:i+1], pos[i:i+1], days[i:i+1]) |
| f = F.optical_features(self.foundation, inputs) |
| |
| if split == 'val': f = f.astype(np.float16).astype(np.float32) |
| optical_features.append(f.reshape(h, w, -1)) |
| optical.append(F.predict(f, self.optical_gaussian).reshape(h, w, 20)) |
| optical = np.stack(optical) |
| incumbent = .65 * reference + .25 * gaussian + .1 * structured |
| final = .5 * incumbent + .5 * optical |
| result = {'reference': reference, 'gaussian': gaussian, 'structured': structured, |
| 'optical': optical, 'incumbent': incumbent, 'final': final} |
| if return_features: result['optical_features'] = np.stack(optical_features) |
| return result |
|
|
| def close(self): |
| self.collector.close() |
| for name in ['sam', 'structured', 'old', 'new', 'foundation']: |
| getattr(self, name).cpu() |
| gc.collect(); torch.cuda.empty_cache() |
|
|
|
|
| def evaluate(paths, split='val', output=None, max_batches=0, save_probabilities=False): |
| """A limited smoke run is explicitly marked incomplete and has no headline score.""" |
| from torch.utils.data import DataLoader |
| assert split in ['val', 'test'] |
| torch.set_num_threads(4); torch.manual_seed(3407) |
| torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True |
| ids = val_patch_ids() if split == 'val' else test_patch_ids() |
| dl = DataLoader(PASTISPatchDataset(ids), batch_size=2 if split == 'val' else 1, |
| shuffle=False, num_workers=2, pin_memory=True, collate_fn=pad_collate) |
| model = CompleteEnsemble(paths); matrices = {}; offset = 0; entries = 0; chunks = [] |
| started = time.monotonic() |
| try: |
| for bi, ((x, pos, days), y) in enumerate(dl): |
| if max_batches and bi >= max_batches: break |
| x, pos = x.cuda(), pos.cuda(); entries += len(x) |
| for i, j in C.FT_CROP_IJ if split == 'val' else [(0, 0)]: |
| crop, target = crop_ij(x, y, i, j, 4 if split == 'val' else 1) |
| parts = model.predict(crop, pos, days, split) |
| for name, p in parts.items(): |
| matrices.setdefault(name, np.zeros((20, 20), np.int64)) |
| matrices[name] += F.confusion(target.numpy().ravel(), p.argmax(-1).ravel()) |
| if save_probabilities: chunks.append(parts['final'].reshape(-1, 20)) |
| offset += target.numel() |
| if bi % 10 == 0: print(split, bi, len(dl), 'seconds', time.monotonic() - started, flush=True) |
| finally: |
| model.close() |
| complete = entries == len(ids) |
| result = {'split': split, 'complete': complete, 'patch_entries': entries, |
| 'expected_patch_entries': len(ids), 'seed': 3407, 'pixel_entries': offset, |
| 'seconds': time.monotonic() - started, |
| 'components': {k: F.metrics(v) for k, v in matrices.items()}, |
| 'paths': {k: str(v) for k, v in paths.items()}, |
| 'selection': 'published fixed recipe; no test tuning'} |
| if output: |
| output = Path(output); output.mkdir(parents=True, exist_ok=True) |
| (output / f'fresh_{split}_results.json').write_text(json.dumps(result, indent=2)) |
| if save_probabilities: np.save(output / f'fresh_{split}_probabilities.npy', np.concatenate(chunks)) |
| return result |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--paths', type=Path, required=True, help='JSON mapping eight named component paths') |
| parser.add_argument('--split', choices=['val', 'test'], default='val') |
| parser.add_argument('--output', type=Path, required=True) |
| parser.add_argument('--max_batches', type=int, default=0) |
| parser.add_argument('--save_probabilities', action='store_true') |
| args = parser.parse_args() |
| results = evaluate(json.loads(args.paths.read_text()), args.split, args.output, |
| args.max_batches, args.save_probabilities) |
| print('complete', results['complete'], 'mIoU', results['components']['final']['miou'], flush=True) |
|
|
