| |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from c80_models import SurfaceCrownNet |
|
|
| DEFAULT_CHECKPOINT=Path(__file__).resolve().parents[1]/'models/c80/c80_v3_step12320.pt' |
|
|
| def valid_fdi(value): |
| if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or not (1 <= value//10 <= 4 and 1 <= value%10 <= 8): |
| raise ValueError('请指定恒牙FDI牙位：11–18、21–28、31–38或41–48') |
| return int(value) |
|
|
| def load_c80(checkpoint=DEFAULT_CHECKPOINT, device='cuda'): |
| saved=torch.load(checkpoint,map_location='cpu',weights_only=True) |
| if saved.get('step') != 12320 or saved.get('architecture') != 'c80': |
| raise ValueError('C80 checkpoint must be v3 step12320') |
| model=SurfaceCrownNet(**saved['config']).float().to(device).eval() |
| model.load_state_dict(saved['model'],strict=True) |
| return model |
|
|
| @torch.inference_mode() |
| def predict_support(model, occupancy_xyz, fdi): |
| fdi=valid_fdi(fdi); a=np.asarray(occupancy_xyz,np.float32) |
| if a.shape != (64,64,64) or not np.isfinite(a).all() or not np.isin(a,[0,1]).all() or a.sum()==0: |
| raise ValueError('C80 jaw must be nonempty binary 64^3 occupancy in XYZ order') |
| device=next(model.parameters()).device |
| jaw=torch.from_numpy(np.ascontiguousarray(a.transpose(2,1,0)))[None,None].to(device) |
| condition=model.encode_condition(jaw,fdi=torch.tensor([fdi],device=device)) |
| _,surface=model.predict_outputs(condition) |
| logits_dhw=surface[0,0].float().cpu().numpy() |
| if not np.isfinite(logits_dhw).all():raise ValueError('Nonfinite C80 prediction') |
| coords_xyz=np.argwhere(logits_dhw>=0)[:,[2,1,0]].astype(np.int64) |
| return np.ascontiguousarray(logits_dhw.transpose(2,1,0)), coords_xyz |
|
|
