| """Preprocessor for the ConvBayes_new dose-calc model (single-channel dose output). |
| |
| Builds a depth-as-channel ``(batch, 1+depth, H, W)`` input tensor from a CT cuboid and |
| per-bixel energy, and rescales the raw network output back to physical dose [Gy]. |
| """ |
|
|
| import array_api_compat |
| import torch |
|
|
| from pyRadPlan.ai_models import BasePreprocessor |
| from pyRadPlan.core.xp_utils import to_namespace |
|
|
|
|
| class ConvDoseSinglePreprocessor(BasePreprocessor): |
| """CT+energy -> network input; network output -> physical dose.""" |
|
|
| def preprocess(self, inputs: dict) -> torch.Tensor: |
| """Assemble the depth-as-channel network input from CT cuboid + energy.""" |
| cfg = self.config |
| ct_cuboid = inputs["ct_cuboid"] |
| energy = inputs["energy"] |
|
|
| clip_lo, clip_hi = cfg["ct_clip_range"] |
| xp = array_api_compat.array_namespace(ct_cuboid) |
| ct_norm = (xp.clip(ct_cuboid, clip_lo, clip_hi) + cfg["ct_offset"]) / cfg["ct_scale"] |
| energy_norm = energy / cfg["max_energy"] |
|
|
| t = to_namespace("torch", ct_norm).float().detach() |
| e_t = to_namespace("torch", energy_norm).float().detach() |
| if e_t.ndim == 0: |
| e_t = e_t.view(1) |
|
|
| batch_size = e_t.shape[0] |
| t = t.expand(batch_size, -1, -1, -1) |
| height, width = t.shape[-2], t.shape[-1] |
| energy_map = e_t.view(-1, 1, 1, 1).expand(-1, 1, height, width) |
| return torch.cat((energy_map, t), dim=1) |
|
|
| def postprocess(self, outputs: torch.Tensor) -> dict: |
| """Rescale the normalized network output back to physical dose [Gy].""" |
| return {"physical_dose": outputs * self.config["max_dose"]} |
|
|
| def predict(self, model: torch.nn.Module, model_input: torch.Tensor) -> dict: |
| """Deterministic single forward pass. Assumes ``model_input`` is already preprocessed.""" |
| model_input = model_input.to(next(model.parameters()).device) |
| with torch.inference_mode(): |
| raw_output = model(model_input) |
| return self.postprocess(raw_output) |
|
|
