mu-net-sample-notebook-inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebb7c42d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:27.104617Z",
     "iopub.status.busy": "2023-12-25T03:29:27.104342Z",
     "iopub.status.idle": "2023-12-25T03:29:52.724213Z",
     "shell.execute_reply": "2023-12-25T03:29:52.723179Z"
    },
    "papermill": {
     "duration": 25.631618,
     "end_time": "2023-12-25T03:29:52.726516",
     "exception": false,
     "start_time": "2023-12-25T03:29:27.094898",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import plotly.graph_objects as go  # for 3d plots\n",
    "\n",
    "from models.agg_3d import Agg3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a86ecab",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:52.745707Z",
     "iopub.status.busy": "2023-12-25T03:29:52.745187Z",
     "iopub.status.idle": "2023-12-25T03:29:57.167031Z",
     "shell.execute_reply": "2023-12-25T03:29:57.166012Z"
    },
    "papermill": {
     "duration": 4.433772,
     "end_time": "2023-12-25T03:29:57.169198",
     "exception": false,
     "start_time": "2023-12-25T03:29:52.735426",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "strategy = tf.distribute.MirroredStrategy()\n",
    "    \n",
    "print(\"REPLICAS: \", strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb617aa9",
   "metadata": {
    "papermill": {
     "duration": 0.008419,
     "end_time": "2023-12-25T03:29:57.186482",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.178063",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Loading data and defining some functions for data visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f30aed9a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.205182Z",
     "iopub.status.busy": "2023-12-25T03:29:57.204590Z",
     "iopub.status.idle": "2023-12-25T03:29:57.211064Z",
     "shell.execute_reply": "2023-12-25T03:29:57.210193Z"
    },
    "papermill": {
     "duration": 0.017808,
     "end_time": "2023-12-25T03:29:57.212976",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.195168",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List of materials by radiation length\n",
    "radiation_lengths = [\n",
    "    1000000000000000000000000, # air\n",
    "    49.834983498349835,  # benzene\n",
    "    49.82309830679809,   # methanol\n",
    "    36.08,               # water\n",
    "    14.385057471264368,  # magnesium\n",
    "    11.552173913043479,  # concrete\n",
    "    10.607758620689655,  # gypsum\n",
    "    10.412903225806451,  # calcium\n",
    "    9.75,                # sulfur\n",
    "    9.368827823100043,   # silicon\n",
    "    8.895887365690998,   # aluminium\n",
    "    4.436732514682328,   # caesium\n",
    "    1.967741935483871,   # manganese\n",
    "    1.7576835153670307,  # iron\n",
    "    1.7200811359026373,  # iodine\n",
    "    1.4243990114580993,  # nickel\n",
    "    0.9589041095890413,  # molybdenum\n",
    "    0.8542857142857143,  # silver\n",
    "    0.6609442060085837,  # polonium\n",
    "    0.5612334801762114,  # lead\n",
    "    0.33436853002070394, # gold\n",
    "    0.316622691292876    # uranium\n",
    "]\n",
    "\n",
    "inverse_radiation_length = [1/x for x in radiation_lengths]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eb8b51a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.238087Z",
     "iopub.status.busy": "2023-12-25T03:29:57.237272Z",
     "iopub.status.idle": "2023-12-25T03:29:57.261511Z",
     "shell.execute_reply": "2023-12-25T03:29:57.259954Z"
    },
    "papermill": {
     "duration": 0.042277,
     "end_time": "2023-12-25T03:29:57.264234",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.221957",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a description of the features.\n",
    "feature_description = {\n",
    "    'x': tf.io.FixedLenFeature([], tf.string),\n",
    "    'y': tf.io.FixedLenFeature([], tf.string)\n",
    "}\n",
    "\n",
    "def _parse_example(example_proto):\n",
    "    res = tf.io.parse_single_example(example_proto, feature_description)\n",
    "    x = tf.io.parse_tensor(res['x'], out_type=tf.double)\n",
    "    y = tf.io.parse_tensor(res['y'], out_type=tf.int32)\n",
    "    y.set_shape((64, 64, 64))\n",
    "    \n",
    "    x = tf.cast(x, dtype=tf.float32)\n",
    "    return x, y\n",
    "\n",
    "def set_dosage(x, y, dosage):\n",
    "    x = tf.random.shuffle(x)\n",
    "    x = x[:dosage]\n",
    "    x.set_shape((dosage, 12))\n",
    "    return x, y\n",
    "\n",
    "def construct_ds(dosage, detector_resolution=-1, p_error=0.2, batch_size=8, dataset=\"test\"):\n",
    "    train_path = \"/path/to/dataset/voxels_prediction.tfrecord\"\n",
    "    test_path = \"/path/to/dataset/voxels_prediction_test.tfrecord\"\n",
    "    val_path = \"/path/to/dataset/voxels_prediction_val.tfrecord\"\n",
    "    \n",
    "    if dataset == \"test\": path = test_path\n",
    "    elif dataset == \"val\": path = val_path\n",
    "    else: path = train_path\n",
    "    \n",
    "    return (\n",
    "        tf.data.TFRecordDataset(path, compression_type=\"GZIP\")\n",
    "        .map(_parse_example)\n",
    "        .filter(lambda x, y: len(x) >= dosage)\n",
    "        .map(lambda x, y: set_dosage(x, y, dosage))\n",
    "        .map(\n",
    "            lambda x, y: (\n",
    "                tf.concat([\n",
    "                    (x[:, :3] / 1000 + 0.5) if detector_resolution < 0 else tf.cast(tf.math.rint(x[:, :3] / 1000 * detector_resolution), tf.float32) / detector_resolution + 0.5, \n",
    "                    x[:, 3:6] / tf.norm(x[:, 3:6], axis=-1, keepdims=True),\n",
    "                    (x[:, 6:9] / 1000 + 0.5) if detector_resolution < 0 else tf.cast(tf.math.rint(x[:, 6:9] / 1000 * detector_resolution), tf.float32) / detector_resolution + 0.5,\n",
    "                    x[:, 9:12],\n",
    "                    (tf.norm(x[:, 3:6], axis=-1, keepdims=True) * tf.random.normal((tf.shape(x[:, :1])), 1, p_error) - 5585.2666) / 13839.263 if p_error > -1e-8 else 0 * tf.norm(x[:, 3:6], axis=-1, keepdims=True),\n",
    "                    tf.norm(x[:, 9:12] - x[:, 3:6] / tf.norm(x[:, 3:6], axis=-1, keepdims=True), axis=-1, keepdims=True),\n",
    "                ], axis=1), tf.gather_nd(\n",
    "                    inverse_radiation_length, \n",
    "                    tf.cast(y[..., tf.newaxis], tf.int32)\n",
    "                )[..., tf.newaxis]\n",
    "            )\n",
    "        ).batch(batch_size).shuffle(200)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f03edfd",
   "metadata": {
    "papermill": {
     "duration": 0.008501,
     "end_time": "2023-12-25T03:29:57.282245",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.273744",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Plot the scattering points predicted by point of closest approach (POCA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ed6a4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.300756Z",
     "iopub.status.busy": "2023-12-25T03:29:57.300483Z",
     "iopub.status.idle": "2023-12-25T03:29:57.306863Z",
     "shell.execute_reply": "2023-12-25T03:29:57.305814Z"
    },
    "papermill": {
     "duration": 0.01818,
     "end_time": "2023-12-25T03:29:57.309128",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.290948",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_scattering_points(pts):\n",
    "    fig = go.Figure(data=[go.Scatter3d(x=pts[:, 0], y=pts[:, 1], z=pts[:, 2], mode='markers')])\n",
    "    fig.update_layout(\n",
    "        scene = dict(\n",
    "            xaxis = dict(nticks=4, range=[0,1]),\n",
    "            yaxis = dict(nticks=4, range=[0,1]),\n",
    "            zaxis = dict(nticks=4, range=[0,1]),\n",
    "            aspectratio={\"x\": 1, \"y\": 1, \"z\": 1}\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b60bf05",
   "metadata": {
    "papermill": {
     "duration": 0.009311,
     "end_time": "2023-12-25T03:29:57.327696",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.318385",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Plot the voxels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20af4df4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.347327Z",
     "iopub.status.busy": "2023-12-25T03:29:57.346655Z",
     "iopub.status.idle": "2023-12-25T03:29:57.355326Z",
     "shell.execute_reply": "2023-12-25T03:29:57.353862Z"
    },
    "papermill": {
     "duration": 0.021565,
     "end_time": "2023-12-25T03:29:57.358363",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.336798",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_voxels(voxels, resolution=64, maximum=3.4, file=None):\n",
    "    x_vals, y_vals, z_vals = tf.meshgrid(tf.range(0,resolution,1), tf.range(0,resolution,1), tf.range(0,resolution,1))\n",
    "    x_vals = tf.reshape(x_vals, (resolution*resolution*resolution,))\n",
    "    y_vals = tf.reshape(y_vals, (resolution*resolution*resolution,))\n",
    "    z_vals = tf.reshape(z_vals, (resolution*resolution*resolution,))\n",
    "    values = tf.reshape(voxels, (resolution*resolution*resolution,))\n",
    "\n",
    "    fig = go.Figure(data=go.Volume(\n",
    "        x=x_vals/resolution,\n",
    "        y=y_vals/resolution,\n",
    "        z=z_vals/resolution,\n",
    "        value=values,\n",
    "        isomin=0,\n",
    "        isomax=maximum,\n",
    "        opacityscale=[[0,0],[1,1],[1,1]],\n",
    "        surface_count=20, # needs to be a large number for good volume rendering,\n",
    "        colorscale=\"blues\",\n",
    "    ))\n",
    "    \n",
    "    if file is None:\n",
    "        fig.show()\n",
    "    else:\n",
    "        fig.write_image(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd51f5a2",
   "metadata": {
    "papermill": {
     "duration": 0.009298,
     "end_time": "2023-12-25T03:29:57.379613",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.370315",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now, we will build the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a659c4b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.403430Z",
     "iopub.status.busy": "2023-12-25T03:29:57.403127Z",
     "iopub.status.idle": "2023-12-25T03:29:57.408358Z",
     "shell.execute_reply": "2023-12-25T03:29:57.407491Z"
    },
    "papermill": {
     "duration": 0.018791,
     "end_time": "2023-12-25T03:29:57.410855",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.392064",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def psnr_max_value(y, y_pred, max_value=inverse_radiation_length[-1]):\n",
    "    print(y_pred.shape, y.shape, max_value)\n",
    "    return 20 * tf.math.log(max_value / tf.sqrt(tf.math.reduce_mean(tf.square(y_pred - y)))) / tf.math.log(10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794ad91b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:29:57.431547Z",
     "iopub.status.busy": "2023-12-25T03:29:57.431272Z",
     "iopub.status.idle": "2023-12-25T03:33:39.857052Z",
     "shell.execute_reply": "2023-12-25T03:33:39.856099Z"
    },
    "papermill": {
     "duration": 222.437285,
     "end_time": "2023-12-25T03:33:39.859091",
     "exception": false,
     "start_time": "2023-12-25T03:29:57.421806",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dosage = 1024\n",
    "print(f\"Training on {dosage}...\")\n",
    "\n",
    "with strategy.scope():\n",
    "    model = Agg3D(\n",
    "        **{\n",
    "            'point_size': 1,\n",
    "            'downward_convs': [1, 2, 3, 4, 5],\n",
    "            'downward_filters': [8, 16, 32, 64, 128],\n",
    "            'upward_convs': [4, 3, 2, 1],\n",
    "            'upward_filters': [64, 32, 16, 8],\n",
    "            'resolution': 64,\n",
    "            'threshold': 1e-8\n",
    "        }\n",
    "    )\n",
    "    model.compile(\n",
    "        optimizer=tf.keras.optimizers.AdamW(learning_rate=2e-3), \n",
    "        loss=\"mse\", metrics=[\"mse\", \"mae\", psnr_max_value]\n",
    "    )\n",
    "\n",
    "    print(model(tf.random.normal((8, dosage, 14))).shape)\n",
    "\n",
    "    model.summary()\n",
    "\n",
    "# splitting dataset into train and validation\n",
    "train_ds = construct_ds(dosage, dataset=\"train\")\n",
    "val_ds = construct_ds(dosage, dataset=\"val\")\n",
    "test_ds = construct_ds(dosage, dataset=\"test\")\n",
    "\n",
    "# training the model\n",
    "model.fit(train_ds, validation_data=val_ds, epochs=15)\n",
    "model.save_weights(f\"model_{dosage}.h5\")\n",
    "\n",
    "# evaluate how good the model really is\n",
    "model.evaluate(test_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ff27ab2",
   "metadata": {
    "papermill": {
     "duration": 0.029328,
     "end_time": "2023-12-25T03:33:39.918521",
     "exception": false,
     "start_time": "2023-12-25T03:33:39.889193",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Looking at some sample reconstructions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model.evaluate(test_ds)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "41726f3fb9824e8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52706186",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:33:39.979912Z",
     "iopub.status.busy": "2023-12-25T03:33:39.979270Z",
     "iopub.status.idle": "2023-12-25T03:34:28.214435Z",
     "shell.execute_reply": "2023-12-25T03:34:28.213416Z"
    },
    "papermill": {
     "duration": 48.293146,
     "end_time": "2023-12-25T03:34:28.241572",
     "exception": false,
     "start_time": "2023-12-25T03:33:39.948426",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "for x, y in test_ds.skip(3).take(1): break\n",
    "y_pred = model.predict(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d50870ad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:34:28.305567Z",
     "iopub.status.busy": "2023-12-25T03:34:28.305232Z",
     "iopub.status.idle": "2023-12-25T03:34:31.803413Z",
     "shell.execute_reply": "2023-12-25T03:34:31.802475Z"
    },
    "papermill": {
     "duration": 3.533241,
     "end_time": "2023-12-25T03:34:31.806527",
     "exception": false,
     "start_time": "2023-12-25T03:34:28.273286",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(4, 4, figsize=(20,20))\n",
    "\n",
    "for i in range(4):\n",
    "    ax[0, i].imshow(y[i, :, :, 32])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[1, i].imshow(y_pred[i, :, :, 32])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[2, i].imshow(y[4+i, :, :, 32])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[3, i].imshow(y_pred[4+i, :, :, 32])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eaf71b2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-25T03:34:31.914937Z",
     "iopub.status.busy": "2023-12-25T03:34:31.914565Z",
     "iopub.status.idle": "2023-12-25T03:34:35.021740Z",
     "shell.execute_reply": "2023-12-25T03:34:35.020868Z"
    },
    "papermill": {
     "duration": 3.142578,
     "end_time": "2023-12-25T03:34:35.024716",
     "exception": false,
     "start_time": "2023-12-25T03:34:31.882138",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(4, 4, figsize=(20,20))\n",
    "\n",
    "for i in range(4):\n",
    "    ax[0, i].imshow(y[i, :, 32, :])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[1, i].imshow(y_pred[i, :, 32, :])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[2, i].imshow(y[4+i, :, 32, :])\n",
    "    \n",
    "for i in range(4):\n",
    "    ax[3, i].imshow(y_pred[4+i, :, 32, :])"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "datasetId": 3523040,
     "sourceId": 7267207,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 4201158,
     "sourceId": 7274831,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30627,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 339.843392,
   "end_time": "2023-12-25T03:34:40.432988",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-12-25T03:29:00.589596",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}