Fix volume transformation, add observation wrapper

docs/02_notebooks/L3_slicing.ipynb

@@ -252,13 +252,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from armscan_env.slicing import slice_volume\n",
+    "from armscan_env.envs.state_action import ManipulatorAction\n",
+    "from armscan_env.volumes.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",
+    "    action=ManipulatorAction(rotation=(19.3, 0), translation=(0, 140)),\n",
     "    volume=volume,\n",
     "    slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
     ")\n",
@@ -323,9 +321,10 @@
     "    sliced_volume = slice_volume(\n",
     "        volume=volume,\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",
     "    ax2.set_title(f\"Slice {i}\")\n",

docs/02_notebooks/L4_environment.ipynb

@@ -37,8 +37,9 @@
     "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.slicing import slice_volume\n",
+    "from armscan_env.envs.state_action import ManipulatorAction\n",
     "from armscan_env.util.visualizations import show_clusters\n",
+    "from armscan_env.volumes.slicing import slice_volume\n",
     "from IPython.core.display import HTML\n",
     "\n",
     "config = get_config()"
@@ -110,8 +111,7 @@
     "    sliced_volume = slice_volume(\n",
     "        volume=volume_1,\n",
     "        slice_shape=slice_shape,\n",
-    "        z_rotation=z[i],\n",
-    "        y_trans=t[i],\n",
+    "        action=ManipulatorAction(rotation=(z[i], 0.0), translation=(0.0, t[i])),\n",
     "    )\n",
     "    sliced_img = sitk.GetArrayFromImage(sliced_volume)[:, 0, :].T\n",
     "    ax2.imshow(sliced_img.T, aspect=6, origin=\"lower\")\n",
@@ -195,14 +195,19 @@
     "    termination_criterion=LabelmapEnvTerminationCriterion(),\n",
     "    max_episode_len=10,\n",
     "    rotation_bounds=(30.0, 10.0),\n",
-    "    translation_bounds=(0.0, None),\n",
+    "    translation_bounds=(None, None),\n",
     "    render_mode=\"animation\",\n",
+    "    apply_volume_transformation=True,\n",
     ")\n",
     "\n",
     "observation, info = env.reset()\n",
     "for _ in range(50):\n",
     "    action = env.action_space.sample()\n",
-    "    observation, reward, terminated, truncated, info = env.step(action)\n",
+    "    epsilon = 0.1\n",
+    "    if np.random.rand() > epsilon:\n",
+    "        observation, reward, terminated, truncated, info = env.step(action)\n",
+    "    else:\n",
+    "        observation, reward, terminated, truncated, info = env.step_to_optimal_state()\n",
     "    env.render()\n",
     "\n",
     "    if terminated or truncated:\n",
@@ -219,13 +224,6 @@
    "source": [
     "HTML(animation.to_jshtml())"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

docs/02_notebooks/L5_linear_sweep.ipynb

@@ -193,7 +193,7 @@
     "volume_size = volume_1.GetSize()\n",
     "\n",
     "projected_env = ArmscanEnvFactory(\n",
-    "    name2volume={\"1\": volume_1},\n",
+    "    name2volume={\"2\": volume_2},\n",
     "    observation=ActionRewardObservation(action_shape=(1,)).to_array_observation(),\n",
     "    slice_shape=(volume_size[0], volume_size[2]),\n",
     "    reward_metric=LabelmapClusteringBasedReward(),\n",

scripts/armscan_array_obs.py

@@ -38,9 +38,9 @@
         num_test_envs=1,
         buffer_size=100000,
         batch_size=256,
-        step_per_collect=10,
-        update_per_step=100,
-        start_timesteps=500,
+        step_per_collect=200,
+        update_per_step=2,
+        start_timesteps=5000,
         start_timesteps_random=True,
     )
 
@@ -61,6 +61,7 @@
         termination_criterion=LabelmapEnvTerminationCriterion(min_reward_threshold=-0.1),
         reward_metric=LabelmapClusteringBasedReward(),
         project_actions_to="y",
+        apply_volume_transformation=True,
     )
 
     experiment = (

scripts/armscan_dqn_sac_hl.py

@@ -35,7 +35,7 @@
         num_epochs=1,
         step_per_epoch=1000000,
         num_train_envs=-1,
-        num_test_envs=10,
+        num_test_envs=1,
         buffer_size=1000000,
         batch_size=256,
         step_per_collect=200,

scripts/armscan_sac_hl.py

@@ -33,8 +33,8 @@
     sampling_config = SamplingConfig(
         num_epochs=1,
         step_per_epoch=1000000,
-        num_train_envs=-1,
-        num_test_envs=10,
+        num_train_envs=40,
+        num_test_envs=1,
         buffer_size=1000000,
         batch_size=256,
         step_per_collect=200,

scripts/random_volume_transformations.ipynb

@@ -0,0 +1,190 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ],
+   "id": "60c69d9345beb9d0",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import SimpleITK as sitk\n",
+    "from armscan_env import config\n",
+    "from armscan_env.clustering import TissueClusters\n",
+    "from armscan_env.envs.state_action import ManipulatorAction\n",
+    "from armscan_env.util.visualizations import show_clusters\n",
+    "from armscan_env.volumes.slicing import EulerTransform, slice_volume, transform_volume\n",
+    "\n",
+    "config = config.get_config()"
+   ],
+   "id": "bf5c60e86d1e8e19",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "volume = sitk.ReadImage(config.get_labels_path(1))\n",
+    "volume_img = sitk.GetArrayFromImage(volume)\n",
+    "plt.imshow(volume_img[40, :, :])\n",
+    "action = ManipulatorAction(rotation=(19, 0), translation=(0, 140))\n",
+    "\n",
+    "o = volume.GetOrigin()\n",
+    "x_dash = np.arange(volume_img.shape[2])\n",
+    "b = volume.TransformPhysicalPointToIndex([o[0], o[1] + action.translation[1], o[2]])[1]\n",
+    "y_dash = x_dash * np.tan(np.deg2rad(action.rotation[0])) + b\n",
+    "plt.plot(x_dash, y_dash, linestyle=\"--\", color=\"red\")\n",
+    "\n",
+    "plt.show()"
+   ],
+   "id": "ae347cf3897968a6",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "sliced_volume = slice_volume(\n",
+    "    action=action,\n",
+    "    volume=volume,\n",
+    "    slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
+    ")\n",
+    "sliced_img = sitk.GetArrayFromImage(sliced_volume)[:, 0, :]\n",
+    "print(f\"Slice value range: {np.min(sliced_img)} - {np.max(sliced_img)}\")\n",
+    "\n",
+    "slice = sliced_img\n",
+    "plt.imshow(slice, aspect=6)\n",
+    "plt.show()"
+   ],
+   "id": "cb9c333a74781d5a",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "volume_transformation = ManipulatorAction(rotation=(19, 0), translation=(15, 15))\n",
+    "transformed_volume = transform_volume(volume, volume_transformation)\n",
+    "transformed_action = EulerTransform(volume_transformation).transform_action(action)"
+   ],
+   "id": "26ffcc6d7dece611",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "transformed_img = sitk.GetArrayFromImage(transformed_volume)\n",
+    "plt.imshow(transformed_img[40, :, :])\n",
+    "\n",
+    "ot = transformed_volume.GetOrigin()\n",
+    "x_dash = np.arange(transformed_img.shape[2])\n",
+    "b = volume.TransformPhysicalPointToIndex([o[0], o[1] + transformed_action.translation[1], o[2]])[1]\n",
+    "y_dash = x_dash * np.tan(np.deg2rad(transformed_action.rotation[0])) + b\n",
+    "plt.plot(x_dash, y_dash, linestyle=\"--\", color=\"red\")\n",
+    "\n",
+    "plt.show()"
+   ],
+   "id": "db20a3ff9556e8b4",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "sliced_transformed_volume = slice_volume(\n",
+    "    action=transformed_action,\n",
+    "    volume=transformed_volume,\n",
+    "    slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
+    ")\n",
+    "sliced_img = sitk.GetArrayFromImage(sliced_transformed_volume)[:, 0, :]\n",
+    "print(f\"Slice value range: {np.min(sliced_img)} - {np.max(sliced_img)}\")\n",
+    "\n",
+    "slice = sliced_img\n",
+    "plt.imshow(slice, aspect=6)\n",
+    "plt.show()"
+   ],
+   "id": "acda09e94c3f2f2b",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "volume_2 = sitk.ReadImage(config.get_labels_path(2))\n",
+    "volume_2_img = sitk.GetArrayFromImage(volume_2)\n",
+    "spacing = volume_2.GetSpacing()\n",
+    "plt.imshow(volume_2_img[51, :, :])\n",
+    "action_2 = ManipulatorAction(rotation=(5, 0), translation=(0, 112))\n",
+    "\n",
+    "o = volume_2.GetOrigin()\n",
+    "x_dash = np.arange(volume_2_img.shape[2])\n",
+    "b = volume_2.TransformPhysicalPointToIndex([o[0], o[1] + action_2.translation[1], o[2]])[1]\n",
+    "y_dash = x_dash * np.tan(np.deg2rad(action_2.rotation[0])) + b\n",
+    "plt.plot(x_dash, y_dash, linestyle=\"--\", color=\"red\")\n",
+    "\n",
+    "plt.show()"
+   ],
+   "id": "fb6cfecff1cb7cd4",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "sliced_volume_2 = slice_volume(\n",
+    "    action=action_2,\n",
+    "    volume=volume_2,\n",
+    "    slice_shape=(volume_2.GetSize()[0], volume_2.GetSize()[2]),\n",
+    ")\n",
+    "sliced_img_2 = sitk.GetArrayFromImage(sliced_volume_2)[:, 0, :]\n",
+    "np.save(\"./array\", sliced_img_2)\n",
+    "\n",
+    "cluster = TissueClusters.from_labelmap_slice(sliced_img_2.T)\n",
+    "show_clusters(cluster, sliced_img_2.T, aspect=spacing[2] / spacing[0])\n",
+    "\n",
+    "plt.show()"
+   ],
+   "id": "6462b823c7903838",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/armscan_env/clustering.py

@@ -23,9 +23,9 @@ def find_DBSCAN_clusters(self, labelmap_slice: np.ndarray) -> list["DataCluster"
             case TissueLabel.BONES:
                 return find_DBSCAN_clusters(self, labelmap_slice, eps=4.1, min_samples=46)
             case TissueLabel.TENDONS:
-                return find_DBSCAN_clusters(self, labelmap_slice, eps=4.1, min_samples=46)
+                return find_DBSCAN_clusters(self, labelmap_slice, eps=2.5, min_samples=15)
             case TissueLabel.ULNAR:
-                return find_DBSCAN_clusters(self, labelmap_slice, eps=2.5, min_samples=18)
+                return find_DBSCAN_clusters(self, labelmap_slice, eps=2.0, min_samples=10)
             case _:
                 raise ValueError(f"Unknown tissue label: {self}")
 
@@ -142,9 +142,10 @@ def find_DBSCAN_clusters(
     label_positions = np.array(list(zip(*np.where(binary_mask), strict=True)))
     clusterer = DBSCAN(eps=eps, min_samples=min_samples)
     clusters = clusterer.fit_predict(label_positions)
-    n_clusters = (
-        len(np.unique(clusters)) - 1
-    )  # noise cluster has label -1, we don't take it into account
+    if -1 in clusters:
+        n_clusters = len(np.unique(clusters)) - 1
+    else:
+        n_clusters = len(np.unique(clusters))
     log.debug(f"Found {n_clusters} clusters")
 
     cluster_list = []

src/armscan_env/envs/base.py

@@ -91,20 +91,18 @@ def __init__(self, array_observations: list[ArrayObservation[TStateAction]]):
     def compute_observation(self, state: TStateAction) -> np.ndarray:
         return np.concatenate(
             [obs.compute_observation(state) for obs in self.array_observations],
-            axis=0,
+            axis=1,
         )
 
     @cached_property
     def observation_space(self) -> gym.spaces.Box:
+        return self.concatenate_boxes([obs.observation_space for obs in self.array_observations])
+
+    @staticmethod
+    def concatenate_boxes(boxes: list[gym.spaces.Box]) -> gym.spaces.Box:
         return gym.spaces.Box(
-            low=np.concatenate(
-                [obs.observation_space.low for obs in self.array_observations],
-                axis=0,
-            ),
-            high=np.concatenate(
-                [obs.observation_space.high for obs in self.array_observations],
-                axis=0,
-            ),
+            low=np.concatenate([box.low for box in boxes], axis=0),
+            high=np.concatenate([box.high for box in boxes], axis=0),
         )