"""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) # Preserve original nesting and GPU float32 arithmetic. 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() # The original Gaussian classifier reads first-block temporal mean/std. 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) # Original validation came from float16 feature caches; test was fresh. 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)