{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\nPractical Deep Learning for Image Registration\n==============================================\n\nCredit: A Grigis\n\nLoad the data\n-------------\n\nLoad some data.\nYou may need to change the 'outdir' parameter.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import os\nimport sys\nif \"CI_MODE\" in os.environ:\n sys.exit()\nimport logging\nimport numpy as np\nfrom pynet import NetParameters\nfrom pynet.datasets import DataManager, fetch_registration\nfrom pynet.utils import setup_logging\nfrom pynet.interfaces import (\n VoxelMorphNetRegister, ADDNetRegister, VTNetRegister, RCNetRegister)\nimport pynet\nfrom pynet.models.voxelmorphnet import FlowRegularizer\nfrom pynet.models.vtnet import ADDNetRegularizer\nfrom torch.optim import lr_scheduler\nfrom pynet.plotting import plot_history\nfrom pynet.history import History\nfrom pynet.losses import MSELoss, NCCLoss, RCNetLoss, PCCLoss\nfrom pynet.plotting import Board, update_board\nimport matplotlib.pyplot as plt\n\nsetup_logging(level=\"debug\")\nlogger = logging.getLogger(\"pynet\")\nlosses = pynet.get_tools(tool_name=\"losses\")\n\noutdir = \"/neurospin/nsap/tmp/registration\"\ndata = fetch_registration(\n datasetdir=outdir)\nmanager = DataManager(\n input_path=data.input_path,\n metadata_path=data.metadata_path,\n number_of_folds=2,\n batch_size=8,\n sampler=\"random\",\n stratify_label=\"studies\",\n projection_labels={\"studies\": [\"abide\"]},\n test_size=0.1,\n add_input=True,\n sample_size=0.1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Training\n--------\n\nFrom the available models load the VoxelMorphRegister, VTNetRegister or\nADDNet and start the training.\nNote that the two first estimate a non linear deformation and require\nthe input data to be afinely registered. The ADDNet estimate an affine\ntransform. We will see in the next section how to combine them in an\nefficient way.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "base_network = \"rcnet\" # \"vtnet\" # \"addnet\"\n\nif base_network == \"rcnet\":\n rcnet_params = NetParameters(\n input_shape=(128, 128, 128),\n in_channels=2,\n base_network=\"VTNet\",\n n_cascades=1,\n rep=1)\n net = RCNetRegister(\n rcnet_params,\n optimizer_name=\"Adam\",\n learning_rate=1e-4,\n loss=losses[\"RCNetLoss\"](),\n use_cuda=True)\nelif base_network == \"addnet\":\n addnet_params = NetParameters(\n input_shape=(128, 128, 128),\n in_channels=2,\n kernel_size=3,\n padding=1,\n flow_multiplier=1.)\n net = ADDNetRegister(\n addnet_params,\n optimizer_name=\"Adam\",\n learning_rate=1e-4,\n loss=losses[\"PCCLoss\"](concat=True),\n use_cuda=True)\n regularizer = ADDNetRegularizer(k1=0.1, k2=0.1)\n net.add_observer(\"regularizer\", regularizer)\nelif base_network == \"vtnet\":\n vtnet_params = NetParameters(\n input_shape=(128, 128, 128),\n in_channels=2,\n kernel_size=3,\n padding=1,\n flow_multiplier=1.,\n nb_channels=16)\n net = VTNetRegister(\n vtnet_params,\n optimizer_name=\"Adam\",\n learning_rate=1e-4,\n loss=losses[\"PCCLoss\"](concat=True), # MSELoss(concat=True),\n use_cuda=True)\n flow_regularizer = FlowRegularizer(k1=1.)\n net.add_observer(\"regularizer\", flow_regularizer)\nelse:\n vmnet_params = NetParameters(\n vol_size=(128, 128, 128),\n enc_nf=[16, 32, 32, 32],\n dec_nf=[32, 32, 32, 32, 32, 16, 16],\n full_size=True)\n net = VoxelMorphNetRegister(\n vmnet_params,\n optimizer_name=\"Adam\",\n learning_rate=1e-4,\n # weight_decay=1e-5,\n loss=losses[\"MSELoss\"](concat=True), # NCCLoss,\n use_cuda=False)\n flow_regularizer = FlowRegularizer(k1=0.01)\n net.add_observer(\"regularizer\", flow_regularizer)\nprint(net.model)\ndef prepare_pred(y_pred):\n moving = y_pred[0, :, :, :, 64]\n validation_dataset = manager[\"validation\"][0]\n corresponding_index = validation_dataset.indices[0]\n reference = validation_dataset.inputs[corresponding_index, 1:, :, : , 64]\n orginal = validation_dataset.inputs[corresponding_index, :1, :, : , 64]\n moving = np.expand_dims(moving, axis=1)\n reference = np.expand_dims(reference, axis=1)\n orginal = np.expand_dims(orginal, axis=1)\n moving = (moving / moving.max())\n moving = moving * 255\n reference = (reference / reference.max())\n reference = reference * 255\n orginal = (orginal / reference.max())\n orginal = orginal * 255\n return np.concatenate((moving, orginal, reference), axis=0)\nnet.board = Board(port=8097, host=\"http://localhost\",\n env=base_network, display_pred=True,\n prepare_pred=prepare_pred)\nnet.add_observer(\"after_epoch\", update_board)\n\nscheduler = lr_scheduler.ReduceLROnPlateau(\n optimizer=net.optimizer,\n mode=\"min\",\n factor=0.5,\n patience=4,\n verbose=True,\n min_lr=1e-7)\ntrain_history, valid_history = net.training(\n manager=manager,\n nb_epochs=(1 if \"CI_MODE\" in os.environ else 150000),\n checkpointdir=None, # outdir,\n fold_index=0,\n scheduler=scheduler,\n with_validation=True)\nprint(train_history)\nprint(valid_history)\nplot_history(train_history)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Testing\n-------\n\nFinaly use the testing set and check the results.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "y_pred, X, y_true, loss, values = net.testing(\n manager=manager,\n with_logit=False,\n predict=False,\n concat_layer_outputs=[\"flow\"])\nprint(y_pred.shape, X.shape, y_true.shape)\n# y_pred = np.expand_dims(y_pred, axis=1)\n# data = np.concatenate((y_pred, y_true, X), axis=1)\n# plot_data(data, nb_samples=5)\n\nif \"CI_MODE\" not in os.environ:\n plt.show()" ] } ], "metadata": { "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.6.12" } }, "nbformat": 4, "nbformat_minor": 0 }