Wrapper handles inference for 2D models over 3D volume inputs.
(
self,
network: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]],
x: torch.Tensor,
condition: torch.Tensor | None = None,
*args: Any,
**kwargs: Any,
)
| 815 | ) |
| 816 | |
| 817 | def network_wrapper( |
| 818 | self, |
| 819 | network: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]], |
| 820 | x: torch.Tensor, |
| 821 | condition: torch.Tensor | None = None, |
| 822 | *args: Any, |
| 823 | **kwargs: Any, |
| 824 | ) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: |
| 825 | """ |
| 826 | Wrapper handles inference for 2D models over 3D volume inputs. |
| 827 | """ |
| 828 | # Pass 4D input [N, C, H, W]/[N, C, D, W]/[N, C, D, H] to the model as it is 2D. |
| 829 | x = x.squeeze(dim=self.spatial_dim + 2) |
| 830 | |
| 831 | if condition is not None: |
| 832 | condition = condition.squeeze(dim=self.spatial_dim + 2) |
| 833 | out = network(x, condition, *args, **kwargs) |
| 834 | else: |
| 835 | out = network(x, *args, **kwargs) |
| 836 | |
| 837 | # Unsqueeze the network output so it is [N, C, D, H, W] as expected by |
| 838 | # the default SlidingWindowInferer class |
| 839 | if isinstance(out, torch.Tensor): |
| 840 | return out.unsqueeze(dim=self.spatial_dim + 2) |
| 841 | |
| 842 | if isinstance(out, Mapping): |
| 843 | for k in out.keys(): |
| 844 | out[k] = out[k].unsqueeze(dim=self.spatial_dim + 2) |
| 845 | return out |
| 846 | |
| 847 | return tuple(out_i.unsqueeze(dim=self.spatial_dim + 2) for out_i in out) |
| 848 | |
| 849 | |
| 850 | class DiffusionInferer(Inferer): |