| """Evaluate complete final bundles without changing the scored input protocol.""" |
| import argparse,hashlib,json,time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
| import config as C |
| from dataset import PASTISPatchDataset,val_patch_ids,test_patch_ids,pad_collate,crop_ij |
| from model import UTAE_ASPP,UTAEClassificationDual |
| from training_feature_clusters import DepthCollector,probability,metrics_cm |
| from temporal_cnn import TemporalCNN,AttentionCNN |
|
|
| MAIN=Path(__file__).resolve().parents[1] |
|
|
| def digest(path): |
| h=hashlib.sha256() |
| with path.open('rb') as f: |
| for part in iter(lambda:f.read(8*1024*1024),b''):h.update(part) |
| return h.hexdigest() |
|
|
| def manifest(name): |
| folder=MAIN/'configurations'/f'32X32_{name}_seed3407' |
| return json.loads((folder/'manifest.json').read_text()),folder |
|
|
| def load_state(path): |
| return torch.load(path,map_location='cpu',weights_only=False) |
|
|
| def load_models(names,device,verify=True): |
| specs={n:manifest(n)[0] for n in names};shared=MAIN/'checkpoints/shared' |
| for spec in specs.values(): |
| if verify: |
| for key in ['sam_checkpoint','feature_cluster_checkpoint','original_cnn_checkpoint','addition_checkpoint']: |
| file=spec[key] |
| if digest(shared/file)!=spec['sha256'][file]:raise RuntimeError('Checkpoint hash mismatch: '+file) |
| sam=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) |
| first=next(iter(specs.values()));sam.load_state_dict(load_state(shared/first['sam_checkpoint'])['model_state_dict'],strict=True) |
| sam=sam.to(device).eval().requires_grad_(False) |
| clusters=load_state(shared/first['feature_cluster_checkpoint']) |
| old=TemporalCNN(18);old.load_state_dict(load_state(shared/first['original_cnn_checkpoint'])['state'],strict=True);old=old.to(device).eval().requires_grad_(False) |
| additions={} |
| for name,spec in specs.items(): |
| model=TemporalCNN(18) if spec['addition_kind']=='cnn_control' else AttentionCNN(18) |
| ck=load_state(shared/spec['addition_checkpoint']);assert ck['seed']==3407 and ck['classes']==list(range(1,19)) |
| model.load_state_dict(ck['state'],strict=True);additions[name]=model.to(device).eval().requires_grad_(False) |
| return specs,sam,clusters,old,additions |
|
|
| @torch.inference_mode() |
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument('--configuration',choices=['recommended_validation','highest_observed_test','both'],default='recommended_validation') |
| ap.add_argument('--split',choices=['val','test'],default='test') |
| ap.add_argument('--output',type=Path,default=MAIN/'evaluation_results.json') |
| ap.add_argument('--skip_hash_check',action='store_true');args=ap.parse_args() |
| if not torch.cuda.is_available():raise RuntimeError('The scored inference recipe requires a working CUDA runtime; CPU loading is checked separately.') |
| torch.set_num_threads(4);torch.manual_seed(3407);torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True |
| names=['recommended_validation','highest_observed_test'] if args.configuration=='both' else [args.configuration] |
| specs,sam,clusters,old,additions=load_models(names,'cuda',not args.skip_hash_check) |
| collector=DepthCollector(sam.utae);mean,std=clusters['mean'].cuda(),clusters['std'].cuda() |
| state={k:v.cuda() for k,v in clusters['state'].items()} if clusters['kind']=='linear' else clusters['state'].cuda() |
| ids=val_patch_ids() if args.split=='val' else test_patch_ids();batch=2 if args.split=='val' else 1 |
| dl=DataLoader(PASTISPatchDataset(ids),batch_size=batch,shuffle=False,num_workers=2,pin_memory=True,collate_fn=pad_collate) |
| matrices={name:np.zeros((20,20),np.int64) for name in names};started=time.monotonic() |
| try: |
| for bi,((x,pos,_),y) in enumerate(dl): |
| x,pos=x.cuda(),pos.cuda() |
| for i,j in C.FT_CROP_IJ if args.split=='val' else [(0,0)]: |
| xc,yc=crop_ij(x,y,i,j,4 if args.split=='val' else 1);collector.begin(xc,pos) |
| with torch.autocast('cuda',dtype=torch.float16): |
| output=sam(xc,pos);original_logits=old(xc,pos) |
| additional_logits={name:net(xc,pos) for name,net in additions.items()} |
| features=collector.end();spatial=output['refined_logits'].float().softmax(1) |
| pp=probability(features,clusters['depth'],mean,std,clusters['kind'],state) |
| b,_,h,w=spatial.shape;cluster_probs=torch.from_numpy(pp).cuda().reshape(b,h,w,20).permute(0,3,1,2) |
| feature_blend=.5*spatial+.5*cluster_probs |
| original_probs=torch.zeros_like(spatial);original_probs[:,1:19]=original_logits.float().softmax(1) |
| reference=.75*feature_blend+.25*original_probs |
| gt=yc.numpy().ravel();valid=(gt>0)&(gt<19) |
| for name,spec in specs.items(): |
| probs=torch.zeros_like(spatial);probs[:,1:19]=additional_logits[name].float().softmax(1) |
| alpha=spec['addition_alpha'];prediction=((1-alpha)*reference+alpha*probs).argmax(1).cpu().numpy().ravel() |
| matrices[name]+=np.bincount(20*gt[valid]+prediction[valid],minlength=400).reshape(20,20) |
| if bi%50==0:print(args.split,bi,len(dl),flush=True) |
| finally:collector.close() |
| result={'split':args.split,'patch_entries':len(ids),'batch_size':batch,'seed':3407,'seconds':time.monotonic()-started,'configurations':{name:{'metrics':metrics_cm(matrices[name]),'manifest':specs[name]} for name in names}} |
| args.output.parent.mkdir(parents=True,exist_ok=True);args.output.write_text(json.dumps(result,indent=2)) |
| for name,item in result['configurations'].items():print(name,'mIoU',item['metrics']['miou'],flush=True) |
|
|
| if __name__=='__main__':main() |
|
|
