Add slice_axis to SlidingWindowInferer

Signed-off-by: Suraj <b.pai@maastrichtuniversity.nl>

Author: surajpaib

modules/2d_inference_3d_volume.ipynb

@@ -74,13 +74,19 @@
     "\n",
     "\n",
     "class YourSlidingWindowInferer(SlidingWindowInferer):\n",
-    "    def __init__(self, *args, **kwargs):\n",
+    "    def __init__(self, slice_axis: int = 0, \n",
+    "                 *args, **kwargs):\n",
+    "        # Set axis to slice the volume across, for example, `0` could slide over axial slices, `1` over coronal slices \n",
+    "        # and `2` over sagittal slices.\n",
+    "        self.slice_axis = slice_axis\n",
+    "\n",
     "        super().__init__(*args, **kwargs)\n",
     "\n",
     "    def __call__(\n",
     "        self,\n",
     "        inputs: torch.Tensor,\n",
     "        network: Callable[..., torch.Tensor],\n",
+    "        slice_axis: int = 0,\n",
     "        *args: Any,\n",
     "        **kwargs: Any,\n",
     "    ) -> torch.Tensor:\n",
@@ -90,10 +96,11 @@
     "\n",
     "            # If they mismatch and roi_size is 2D add another dimension to roi size\n",
     "            if len(self.roi_size) == 2:\n",
-    "                self.roi_size = [1, *self.roi_size]\n",
+    "                self.roi_size = list(self.roi_size)\n",
+    "                self.roi_size.insert(self.slice_axis, 1)\n",
     "            else:\n",
-    "                raise RuntimeError(\"Unsupported roi size, cannot broadcast to volume. \")\n",
-    "\n",
+    "                raise RuntimeError(\"Currently, only 2D `roi_size` is supported, cannot broadcast to volume. \")\n",
+    "                \n",
     "        return super().__call__(inputs, lambda x: self.network_wrapper(network, x))\n",
     "\n",
     "    def network_wrapper(self, network, x, *args, **kwargs):\n",
@@ -103,12 +110,13 @@
     "        \"\"\"\n",
     "        # If depth dim is 1 in [D, H, W] roi size, then the input is 2D and needs\n",
     "        # be handled accordingly\n",
-    "        if self.roi_size[0] == 1:\n",
-    "            #  Pass [N, C, H, W] to the model as it is 2D.\n",
-    "            x = x.squeeze(dim=2)\n",
+    "\n",
+    "        if self.roi_size[self.slice_axis] == 1:\n",
+    "            #  Pass 4D input [N, C, H, W]/[N, C, D, W]/[N, C, D, H] to the model as it is 2D.\n",
+    "            x = x.squeeze(dim=self.slice_axis + 2)\n",
     "            out = network(x, *args, **kwargs)\n",
     "            #  Unsqueeze the network output so it is [N, C, D, H, W] as expected by the default SlidingWindowInferer class\n",
-    "            return out.unsqueeze(dim=2)\n",
+    "            return out.unsqueeze(dim=self.slice_axis + 2)\n",
     "\n",
     "        else:\n",
     "            return network(x, *args, **kwargs)"
@@ -128,7 +136,16 @@
    "execution_count": 3,
    "id": "85b15305",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Axial Inferer Output Shape:  torch.Size([1, 1, 64, 256, 256])\n",
+      "Coronal Inferer Output Shape:  torch.Size([1, 1, 64, 256, 256])\n"
+     ]
+    }
+   ],
    "source": [
     "# Create a 2D UNet with randomly initialized weights for testing purposes\n",
     "from monai.networks.nets import UNet\n",
@@ -144,37 +161,36 @@
     ")\n",
     "\n",
     "# Initialize a dummy 3D tensor volume with shape (N,C,D,H,W)\n",
-    "input_volume = torch.ones(1, 1, 30, 256, 256)\n",
+    "input_volume = torch.ones(1, 1, 64, 256, 256)\n",
     "\n",
-    "# Create an instance of YourSlidingWindowInferer with roi_size as the 256x256 (HxW)\n",
-    "inferer = YourSlidingWindowInferer(roi_size=(256, 256),\n",
+    "# Create an instance of YourSlidingWindowInferer with roi_size as the 256x256 (HxW) and sliding over D axis\n",
+    "axial_inferer = YourSlidingWindowInferer(roi_size=(256, 256),\n",
     "                                   sw_batch_size=1,\n",
     "                                   cval=-1)\n",
     "\n",
-    "output = inferer(input_volume, net)"
+    "output = axial_inferer(input_volume, net)\n",
+    "\n",
+    "# Output is a 3D volume with 2D slices aggregated\n",
+    "print(\"Axial Inferer Output Shape: \", output.shape)\n",
+    "\n",
+    "# Create an instance of YourSlidingWindowInferer with roi_size as the 64x256 (DxW) and sliding over H axis\n",
+    "coronal_inferer = YourSlidingWindowInferer(roi_size=(64, 256),\n",
+    "                                   sw_batch_size=1, slice_axis=1, # Slice axis is added here\n",
+    "                                   cval=-1)\n",
+    "\n",
+    "output = coronal_inferer(input_volume, net)\n",
+    "\n",
+    "# Output is a 3D volume with 2D slices aggregated\n",
+    "print(\"Coronal Inferer Output Shape: \", output.shape)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "id": "5ad96534",
+   "execution_count": null,
+   "id": "454c353f",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([1, 1, 30, 256, 256])"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# Output is a 3D volume with 2D slices aggregated\n",
-    "output.shape"
-   ]
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {