Fix volume transformation, add observation wrapper

.github/workflows/lint_and_docs.yaml

@@ -1,7 +1,19 @@
 name: PEP8, Types and Docs Check
 
-on: [push, pull_request]
-
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+        - main
+  workflow_dispatch:
+    inputs:
+      debug_enabled:
+        type: boolean
+        description: 'Run the build with tmate debugging enabled (https://github.com/marketplace/actions/debugging-with-tmate)'
+        required: false
+        default: false
 jobs:
   check:
     runs-on: ubuntu-latest
@@ -12,6 +24,10 @@ jobs:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}
     steps:
+      # Enable tmate debugging of manually-triggered workflows if the input option was provided
+      - name: Setup tmate session
+        uses: mxschmitt/action-tmate@v3
+        if: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_enabled }}
       - name: Cancel previous run
         uses: styfle/[email protected]
         with:

.github/workflows/pytest.yml

@@ -1,6 +1,12 @@
 name: Ubuntu
 
-on: [push, pull_request]
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
 
 jobs:
   cpu:

.pre-commit-config.yaml

@@ -52,5 +52,5 @@ repos:
         language: system
       - id: clean-nbs
         name: clean-nbs
-        entry: poetry run python docs/nbstripout.py
+        entry: poetry run python nbstripout.py
         language: system

docs/02_notebooks/L0_MRI_and_Labelmaps.ipynb

@@ -62,7 +62,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mri_1 = sitk.ReadImage(config.get_mri_path(1))\n",
+    "mri_1 = sitk.ReadImage(config.get_single_mri_path(1))\n",
     "mri_1_data = sitk.GetArrayFromImage(mri_1)\n",
     "print(f\"{mri_1_data.shape=}\")"
    ]
@@ -106,7 +106,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mri_1_label = sitk.ReadImage(config.get_labels_path(1))\n",
+    "mri_1_label = sitk.ReadImage(config.get_single_labelmap_path(1))\n",
     "mri_1_label_data = sitk.GetArrayFromImage(mri_1_label)\n",
     "print(f\"{mri_1_label_data.shape =}\")"
    ]

docs/02_notebooks/L1_simple_clustering.ipynb

@@ -53,7 +53,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mri_1_label = sitk.ReadImage(config.get_labels_path(1))\n",
+    "mri_1_label = sitk.ReadImage(config.get_single_labelmap_path(1))\n",
     "mri_1_label_data = sitk.GetArrayFromImage(mri_1_label)\n",
     "print(f\"{mri_1_label_data.shape=}\")"
    ]

docs/02_notebooks/L2_DBSCAN_clustering.ipynb

@@ -52,7 +52,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "path_to_mri = config.get_labels_path(1)\n",
+    "path_to_mri = config.get_single_labelmap_path(1)\n",
     "mri_1_label = sitk.ReadImage(path_to_mri)\n",
     "mri_1_label_data = sitk.GetArrayFromImage(mri_1_label)\n",
     "print(f\"{mri_1_label_data.shape=}\")"
@@ -207,18 +207,23 @@
     "zero_loss_indices = np.where(np.array(sweep_loss) == 0)[0]\n",
     "print(f\"{len(zero_loss_indices)} indices return a zero loss: \", zero_loss_indices)\n",
     "\n",
-    "fig, axes = plt.subplots(2, 4, figsize=(21, 7))\n",
-    "axes = axes.flatten()\n",
-    "for i, idx in enumerate(zero_loss_indices):\n",
-    "    axes[i] = show_clusters(\n",
-    "        tissue_clusters=clusters_list[idx],\n",
-    "        slice=mri_1_label_data[:, idx, :].T,\n",
-    "        aspect=6,\n",
-    "        ax=axes[i],\n",
-    "    )\n",
-    "    axes[i].set_title(f\"Index: {idx}, Loss: {sweep_loss[idx]:.2f}\")\n",
+    "nrows = 2\n",
+    "ncols = len(zero_loss_indices) // nrows\n",
+    "indices_to_display = nrows * ncols\n",
     "\n",
-    "plt.show()"
+    "if indices_to_display > 0:\n",
+    "    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(21, 7))\n",
+    "    axes = axes.flatten()\n",
+    "    for i, idx in enumerate(zero_loss_indices[:indices_to_display]):\n",
+    "        axes[i] = show_clusters(\n",
+    "            tissue_clusters=clusters_list[idx],\n",
+    "            slice=mri_1_label_data[:, idx, :].T,\n",
+    "            aspect=6,\n",
+    "            ax=axes[i],\n",
+    "        )\n",
+    "        axes[i].set_title(f\"Index: {idx}, Loss: {sweep_loss[idx]:.2f}\")\n",
+    "\n",
+    "    plt.show()"
    ]
   },
   {
@@ -243,7 +248,7 @@
     "\n",
     "for i in range(mri_1_label_data.shape[1]):\n",
     "    clusters = TissueClusters.from_labelmap_slice(mri_1_label_data[:, i, :].T)\n",
-    "    loss = anatomy_based_rwd(clusters, n_landmarks=[7, 2, 1])\n",
+    "    loss = anatomy_based_rwd(clusters, n_landmarks=[7, 3, 1])\n",
     "    if loss == 0:\n",
     "        zero_loss_clusters.append(clusters)\n",
     "    print(f\"Loss for slice {i}: {loss}\")\n",
@@ -289,21 +294,28 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# TODO: reduce duplication with printing above, move to a function\n",
+    "\n",
     "zero_loss_indices = np.where(np.array(sweep_loss) == 0)[0]\n",
     "print(f\"{len(zero_loss_indices)} indices return a zero loss: \", zero_loss_indices)\n",
     "\n",
-    "fig, axes = plt.subplots(2, 4, figsize=(21, 7))\n",
-    "axes = axes.flatten()\n",
-    "for i, idx in enumerate(zero_loss_indices):\n",
-    "    axes[i] = show_clusters(\n",
-    "        tissue_clusters=zero_loss_clusters[i],\n",
-    "        slice=mri_1_label_data[:, idx, :].T,\n",
-    "        aspect=6,\n",
-    "        ax=axes[i],\n",
-    "    )\n",
-    "    axes[i].set_title(f\"Index: {idx}, Loss: {sweep_loss[idx]:.2f}\")\n",
+    "nrows = 2\n",
+    "ncols = len(zero_loss_indices) // nrows\n",
+    "indices_to_display = nrows * ncols\n",
     "\n",
-    "plt.show()"
+    "if indices_to_display > 0:\n",
+    "    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(21, 7))\n",
+    "    axes = axes.flatten()\n",
+    "    for i, idx in enumerate(zero_loss_indices[:indices_to_display]):\n",
+    "        axes[i] = show_clusters(\n",
+    "            tissue_clusters=zero_loss_clusters[i],\n",
+    "            slice=mri_1_label_data[:, idx, :].T,\n",
+    "            aspect=6,\n",
+    "            ax=axes[i],\n",
+    "        )\n",
+    "        axes[i].set_title(f\"Index: {idx}, Loss: {sweep_loss[idx]:.2f}\")\n",
+    "\n",
+    "    plt.show()"
    ]
   }
  ],

docs/02_notebooks/L3_slicing.ipynb

@@ -34,7 +34,13 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import SimpleITK as sitk\n",
+    "from armscan_env.clustering import TissueClusters\n",
     "from armscan_env.config import get_config\n",
+    "from armscan_env.envs.rewards import anatomy_based_rwd\n",
+    "from armscan_env.envs.state_action import ManipulatorAction\n",
+    "from armscan_env.util.visualizations import show_clusters\n",
+    "from armscan_env.volumes.volumes import ImageVolume\n",
+    "from celluloid import Camera\n",
     "from IPython.core.display import HTML\n",
     "\n",
     "config = get_config()"
@@ -46,7 +52,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "volume = sitk.ReadImage(config.get_labels_path(1))\n",
+    "volume = sitk.ReadImage(config.get_single_labelmap_path(1))\n",
     "img_array = sitk.GetArrayFromImage(volume)\n",
     "print(f\"{img_array.shape=}\")"
    ]
@@ -164,7 +170,6 @@
     "# (cosine of the angle between the normal vector and the x axis: z-rotation)\n",
     "w = int(abs(volume_size[0] // e1[0]))\n",
     "\n",
-    "\n",
     "print(f\" {h=},\\n {w=}\")"
    ]
   },
@@ -252,17 +257,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from armscan_env.slicing import slice_volume\n",
-    "\n",
-    "sliced_volume = slice_volume(\n",
-    "    z_rotation=19.3,\n",
-    "    x_rotation=0.0,\n",
-    "    x_trans=0.0,\n",
-    "    y_trans=140.0,\n",
-    "    volume=volume,\n",
+    "volume = ImageVolume(volume)\n",
+    "sliced_volume = volume.get_volume_slice(\n",
+    "    action=ManipulatorAction(rotation=(19.3, 0), translation=(0, 140)),\n",
     "    slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
     ")\n",
-    "sliced_img = sitk.GetArrayFromImage(sliced_volume)[:, 0, :]\n",
+    "sliced_img = sitk.GetArrayFromImage(sliced_volume)\n",
     "print(f\"Slice value range: {np.min(sliced_img)} - {np.max(sliced_img)}\")\n",
     "\n",
     "slice = sliced_img\n",
@@ -276,10 +276,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from armscan_env.clustering import TissueClusters\n",
-    "from armscan_env.envs.rewards import anatomy_based_rwd\n",
-    "from armscan_env.util.visualizations import show_clusters\n",
-    "\n",
     "clusters = TissueClusters.from_labelmap_slice(sliced_img.T)\n",
     "show_clusters(clusters, sliced_img.T)\n",
     "reward = anatomy_based_rwd(clusters, (4, 2, 1))\n",
@@ -294,8 +290,6 @@
    },
    "outputs": [],
    "source": [
-    "from celluloid import Camera\n",
-    "\n",
     "# Demonstration of arbitrary slicing\n",
     "t = [160, 155, 150, 148, 146, 142, 140, 140, 115, 120, 125, 125, 130, 130, 135, 138, 140, 140, 140]\n",
     "z = [0, -5, 0, 0, 5, 15, 19.3, -10, 0, 0, 0, 5, -8, 8, 0, -10, -10, 10, 19.3]\n",
@@ -320,14 +314,14 @@
     "    ax1.plot(x_dash, y_dash, linestyle=\"--\", color=\"red\")\n",
     "\n",
     "    # Subplot 2: Function image\n",
-    "    sliced_volume = slice_volume(\n",
-    "        volume=volume,\n",
+    "    sliced_volume = volume.get_volume_slice(\n",
     "        slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
-    "        z_rotation=z[i],\n",
-    "        x_rotation=0,\n",
-    "        y_trans=t[i],\n",
+    "        action=ManipulatorAction(\n",
+    "            rotation=(z[i], 0),\n",
+    "            translation=(0, t[i]),\n",
+    "        ),\n",
     "    )\n",
-    "    sliced_img = sitk.GetArrayFromImage(sliced_volume)[:, 0, :]\n",
+    "    sliced_img = sitk.GetArrayFromImage(sliced_volume)\n",
     "    ax2.set_title(f\"Slice {i}\")\n",
     "    ax2.imshow(sliced_img, aspect=6)\n",
     "    camera.snap()\n",