project5.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "\n",
    "os.environ['NEURITE_BACKEND'] = 'pytorch'\n",
    "os.environ['VXM_BACKEND'] = 'pytorch'\n",
    "\n",
    "# some third party very useful libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import nibabel as nib\n",
    "\n",
    "# our libraries\n",
    "import voxelmorph as vxm\n",
    "import neurite as ne"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = [f + '/slice_norm.nii.gz' for f in os.listdir('./data') if f.startswith('OASIS_OAS1_')]\n",
    "vols = [nib.load('./data/'+f).get_fdata() for f in files]\n",
    "x_vols = np.stack(vols, 0)\n",
    "vol_shape = x_vols.shape[1:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions.normal import Normal\n",
    "\n",
    "from voxelmorph.voxelmorph import default_unet_features\n",
    "from voxelmorph.voxelmorph.torch import layers\n",
    "from voxelmorph.voxelmorph.torch.modelio import LoadableModel, store_config_args\n",
    "\n",
    "\n",
    "class Unet(nn.Module):\n",
    "    \"\"\"\n",
    "    A unet architecture. Layer features can be specified directly as a list of encoder and decoder\n",
    "    features or as a single integer along with a number of unet levels. The default network features\n",
    "    per layer (when no options are specified) are:\n",
    "\n",
    "        encoder: [16, 32, 32, 32]\n",
    "        decoder: [32, 32, 32, 32, 32, 16, 16]\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 inshape=None,\n",
    "                 infeats=None,\n",
    "                 nb_features=None,\n",
    "                 nb_levels=None,\n",
    "                 max_pool=2,\n",
    "                 feat_mult=1,\n",
    "                 nb_conv_per_level=1,\n",
    "                 half_res=False):\n",
    "        \"\"\"\n",
    "        Parameters:\n",
    "            inshape: Input shape. e.g. (192, 192, 192)\n",
    "            infeats: Number of input features.\n",
    "            nb_features: Unet convolutional features. Can be specified via a list of lists with\n",
    "                the form [[encoder feats], [decoder feats]], or as a single integer. \n",
    "                If None (default), the unet features are defined by the default config described in \n",
    "                the class documentation.\n",
    "            nb_levels: Number of levels in unet. Only used when nb_features is an integer. \n",
    "                Default is None.\n",
    "            feat_mult: Per-level feature multiplier. Only used when nb_features is an integer. \n",
    "                Default is 1.\n",
    "            nb_conv_per_level: Number of convolutions per unet level. Default is 1.\n",
    "            half_res: Skip the last decoder upsampling. Default is False.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        # ensure correct dimensionality\n",
    "        ndims = len(inshape)\n",
    "        assert ndims in [1, 2, 3], 'ndims should be one of 1, 2, or 3. found: %d' % ndims\n",
    "\n",
    "        # cache some parameters\n",
    "        self.half_res = half_res\n",
    "\n",
    "        # default encoder and decoder layer features if nothing provided\n",
    "        if nb_features is None:\n",
    "            nb_features = default_unet_features()\n",
    "\n",
    "        # build feature list automatically\n",
    "        if isinstance(nb_features, int):\n",
    "            if nb_levels is None:\n",
    "                raise ValueError('must provide unet nb_levels if nb_features is an integer')\n",
    "            feats = np.round(nb_features * feat_mult ** np.arange(nb_levels)).astype(int)\n",
    "            nb_features = [\n",
    "                np.repeat(feats[:-1], nb_conv_per_level),\n",
    "                np.repeat(np.flip(feats), nb_conv_per_level)\n",
    "            ]\n",
    "        elif nb_levels is not None:\n",
    "            raise ValueError('cannot use nb_levels if nb_features is not an integer')\n",
    "\n",
    "        # extract any surplus (full resolution) decoder convolutions\n",
    "        enc_nf, dec_nf = nb_features\n",
    "        nb_dec_convs = len(enc_nf)\n",
    "        final_convs = dec_nf[nb_dec_convs:]\n",
    "        dec_nf = dec_nf[:nb_dec_convs]\n",
    "        self.nb_levels = int(nb_dec_convs / nb_conv_per_level) + 1\n",
    "\n",
    "        if isinstance(max_pool, int):\n",
    "            max_pool = [max_pool] * self.nb_levels\n",
    "\n",
    "        # cache downsampling / upsampling operations\n",
    "        MaxPooling = getattr(nn, 'MaxPool%dd' % ndims)\n",
    "        self.pooling = [MaxPooling(s) for s in max_pool]\n",
    "        self.upsampling = [nn.Upsample(scale_factor=s, mode='nearest') for s in max_pool]\n",
    "\n",
    "        # configure encoder (down-sampling path)\n",
    "        prev_nf = infeats\n",
    "        encoder_nfs = [prev_nf]\n",
    "        self.encoder = nn.ModuleList()\n",
    "        for level in range(self.nb_levels - 1):\n",
    "            convs = nn.ModuleList()\n",
    "            for conv in range(nb_conv_per_level):\n",
    "                nf = enc_nf[level * nb_conv_per_level + conv]\n",
    "                convs.append(ConvBlock(ndims, prev_nf, nf))\n",
    "                prev_nf = nf\n",
    "            self.encoder.append(convs)\n",
    "            encoder_nfs.append(prev_nf)\n",
    "\n",
    "        # configure decoder (up-sampling path)\n",
    "        encoder_nfs = np.flip(encoder_nfs)\n",
    "        self.decoder = nn.ModuleList()\n",
    "        for level in range(self.nb_levels - 1):\n",
    "            convs = nn.ModuleList()\n",
    "            for conv in range(nb_conv_per_level):\n",
    "                nf = dec_nf[level * nb_conv_per_level + conv]\n",
    "                convs.append(ConvBlock(ndims, prev_nf, nf))\n",
    "                prev_nf = nf\n",
    "            self.decoder.append(convs)\n",
    "            if not half_res or level < (self.nb_levels - 2):\n",
    "                prev_nf += encoder_nfs[level]\n",
    "\n",
    "        # now we take care of any remaining convolutions\n",
    "        self.remaining = nn.ModuleList()\n",
    "        for num, nf in enumerate(final_convs):\n",
    "            self.remaining.append(ConvBlock(ndims, prev_nf, nf))\n",
    "            prev_nf = nf\n",
    "\n",
    "        # cache final number of features\n",
    "        self.final_nf = prev_nf\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        # encoder forward pass\n",
    "        x_history = [x]\n",
    "        for level, convs in enumerate(self.encoder):\n",
    "            for conv in convs:\n",
    "                x = conv(x)\n",
    "            x_history.append(x)\n",
    "            x = self.pooling[level](x)\n",
    "\n",
    "        # decoder forward pass with upsampling and concatenation\n",
    "        for level, convs in enumerate(self.decoder):\n",
    "            for conv in convs:\n",
    "                x = conv(x)\n",
    "            if not self.half_res or level < (self.nb_levels - 2):\n",
    "                x = self.upsampling[level](x)\n",
    "                x = torch.cat([x, x_history.pop()], dim=1)\n",
    "\n",
    "        # remaining convs at full resolution\n",
    "        for conv in self.remaining:\n",
    "            x = conv(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class VxmDense(LoadableModel):\n",
    "    \"\"\"\n",
    "    VoxelMorph network for (unsupervised) nonlinear registration between two images.\n",
    "    \"\"\"\n",
    "\n",
    "    @store_config_args\n",
    "    def __init__(self,\n",
    "                 inshape,\n",
    "                 nb_unet_features=None,\n",
    "                 nb_unet_levels=None,\n",
    "                 unet_feat_mult=1,\n",
    "                 nb_unet_conv_per_level=1,\n",
    "                 int_steps=7,\n",
    "                 int_downsize=1,\n",
    "                 bidir=False,\n",
    "                 use_probs=False,\n",
    "                 src_feats=1,\n",
    "                 trg_feats=1,\n",
    "                 unet_half_res=False):\n",
    "        \"\"\" \n",
    "        Parameters:\n",
    "            inshape: Input shape. e.g. (192, 192, 192)\n",
    "            nb_unet_features: Unet convolutional features. Can be specified via a list of lists with\n",
    "                the form [[encoder feats], [decoder feats]], or as a single integer. \n",
    "                If None (default), the unet features are defined by the default config described in \n",
    "                the unet class documentation.\n",
    "            nb_unet_levels: Number of levels in unet. Only used when nb_features is an integer. \n",
    "                Default is None.\n",
    "            unet_feat_mult: Per-level feature multiplier. Only used when nb_features is an integer. \n",
    "                Default is 1.\n",
    "            nb_unet_conv_per_level: Number of convolutions per unet level. Default is 1.\n",
    "            int_steps: Number of flow integration steps. The warp is non-diffeomorphic when this \n",
    "                value is 0.\n",
    "            int_downsize: Integer specifying the flow downsample factor for vector integration. \n",
    "                The flow field is not downsampled when this value is 1.\n",
    "            bidir: Enable bidirectional cost function. Default is False.\n",
    "            use_probs: Use probabilities in flow field. Default is False.\n",
    "            src_feats: Number of source image features. Default is 1.\n",
    "            trg_feats: Number of target image features. Default is 1.\n",
    "            unet_half_res: Skip the last unet decoder upsampling. Requires that int_downsize=2. \n",
    "                Default is False.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        # internal flag indicating whether to return flow or integrated warp during inference\n",
    "        self.training = True\n",
    "\n",
    "        # ensure correct dimensionality\n",
    "        ndims = len(inshape)\n",
    "        assert ndims in [1, 2, 3], 'ndims should be one of 1, 2, or 3. found: %d' % ndims\n",
    "\n",
    "        # configure core unet model\n",
    "        self.unet_model = Unet(\n",
    "            inshape,\n",
    "            infeats=(src_feats + trg_feats),\n",
    "            nb_features=nb_unet_features,\n",
    "            nb_levels=nb_unet_levels,\n",
    "            feat_mult=unet_feat_mult,\n",
    "            nb_conv_per_level=nb_unet_conv_per_level,\n",
    "            half_res=unet_half_res,\n",
    "        )\n",
    "\n",
    "        # configure unet to flow field layer\n",
    "        Conv = getattr(nn, 'Conv%dd' % ndims)\n",
    "        self.flow = Conv(self.unet_model.final_nf, ndims, kernel_size=3, padding=1)\n",
    "\n",
    "        # init flow layer with small weights and bias\n",
    "        self.flow.weight = nn.Parameter(Normal(0, 1e-5).sample(self.flow.weight.shape))\n",
    "        self.flow.bias = nn.Parameter(torch.zeros(self.flow.bias.shape))\n",
    "\n",
    "        # probabilities are not supported in pytorch\n",
    "        if use_probs:\n",
    "            raise NotImplementedError(\n",
    "                'Flow variance has not been implemented in pytorch - set use_probs to False')\n",
    "\n",
    "        # configure optional resize layers (downsize)\n",
    "        if not unet_half_res and int_steps > 0 and int_downsize > 1:\n",
    "            self.resize = layers.ResizeTransform(int_downsize, ndims)\n",
    "        else:\n",
    "            self.resize = None\n",
    "\n",
    "        # resize to full res\n",
    "        if int_steps > 0 and int_downsize > 1:\n",
    "            self.fullsize = layers.ResizeTransform(1 / int_downsize, ndims)\n",
    "        else:\n",
    "            self.fullsize = None\n",
    "\n",
    "        # configure bidirectional training\n",
    "        self.bidir = bidir\n",
    "\n",
    "        # configure optional integration layer for diffeomorphic warp\n",
    "        down_shape = [int(dim / int_downsize) for dim in inshape]\n",
    "        self.integrate = layers.VecInt(down_shape, int_steps) if int_steps > 0 else None\n",
    "\n",
    "        # configure transformer\n",
    "        self.transformer = layers.SpatialTransformer(inshape)\n",
    "\n",
    "    def forward(self, source, target, registration=False):\n",
    "        '''\n",
    "        Parameters:\n",
    "            source: Source image tensor.\n",
    "            target: Target image tensor.\n",
    "            registration: Return transformed image and flow. Default is False.\n",
    "        '''\n",
    "\n",
    "        # concatenate inputs and propagate unet\n",
    "        x = torch.cat([source, target], dim=1)\n",
    "\n",
    "        x = self.unet_model(x)\n",
    "\n",
    "\n",
    "        # transform into flow field\n",
    "        flow_field = self.flow(x)\n",
    "\n",
    "        # resize flow for integration\n",
    "        pos_flow = flow_field\n",
    "        if self.resize:\n",
    "            pos_flow = self.resize(pos_flow)\n",
    "\n",
    "        preint_flow = pos_flow\n",
    "\n",
    "        # negate flow for bidirectional model\n",
    "        neg_flow = -pos_flow if self.bidir else None\n",
    "\n",
    "        # integrate to produce diffeomorphic warp\n",
    "        if self.integrate:\n",
    "            pos_flow = self.integrate(pos_flow)\n",
    "            neg_flow = self.integrate(neg_flow) if self.bidir else None\n",
    "\n",
    "            # resize to final resolution\n",
    "            if self.fullsize:\n",
    "                pos_flow = self.fullsize(pos_flow)\n",
    "                neg_flow = self.fullsize(neg_flow) if self.bidir else None\n",
    "\n",
    "        # warp image with flow field\n",
    "        y_source = self.transformer(source, pos_flow)\n",
    "        y_target = self.transformer(target, neg_flow) if self.bidir else None\n",
    "\n",
    "        # return non-integrated flow field if training\n",
    "        if not registration:\n",
    "            return (y_source, y_target, preint_flow) if self.bidir else (y_source, preint_flow)\n",
    "        else:\n",
    "            return y_source, pos_flow\n",
    "\n",
    "\n",
    "class ConvBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    Specific convolutional block followed by leakyrelu for unet.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, ndims, in_channels, out_channels, stride=1):\n",
    "        super().__init__()\n",
    "\n",
    "        Conv = getattr(nn, 'Conv%dd' % ndims)\n",
    "        self.main = Conv(in_channels, out_channels, 3, stride, 1)\n",
    "        self.activation = nn.LeakyReLU(0.2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.main(x)\n",
    "        out = self.activation(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MeanStream(nn.Module):\n",
    "    \"\"\"\n",
    "    Maintain stream of data mean.\n",
    "\n",
    "    cap refers to maintaining an approximation of up to that number of subjects -- that is,\n",
    "    any incoming data point will have at least 1/cap weight.\n",
    "\n",
    "    If you find this class useful, please cite the original paper this was written for:\n",
    "        A.V. Dalca, M. Rakic, J. Guttag, M.R. Sabuncu.\n",
    "        Learning Conditional Deformable Templates with Convolutional Networks\n",
    "        NeurIPS: Advances in Neural Information Processing Systems. pp 804-816, 2019.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, cap=100, **kwargs):\n",
    "        super(MeanStream, self).__init__(**kwargs)\n",
    "        self.cap = float(cap)\n",
    "        self.mean = nn.Parameter(torch.zeros(1, 2, 160, 192), requires_grad=False)\n",
    "        self.count = nn.Parameter(torch.zeros(1), requires_grad=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Get batch size\n",
    "        this_bs_int = x.size(0)\n",
    "\n",
    "        # If calling in inference mode, use moving stats\n",
    "        if not self.training:\n",
    "            return torch.min(torch.tensor(1.), self.count / self.cap) * (torch.ones((this_bs_int, 1, 160, 192)) * self.mean)\n",
    "\n",
    "        # Get new mean and count\n",
    "        new_mean, new_count = self._mean_update(x)\n",
    "\n",
    "        # Update mean and count\n",
    "        self.count.data = new_count\n",
    "        self.mean.data = new_mean\n",
    "\n",
    "        # The first few 1000 should not matter that much towards this cost\n",
    "        return torch.min(torch.tensor(1.), new_count / self.cap) * (torch.ones((this_bs_int, 1, 160, 192)).cuda() * new_mean)\n",
    "\n",
    "    def _mean_update(self, x):\n",
    "        # Convert to float for calculations\n",
    "        x_float = x.float()\n",
    "\n",
    "        # Compute sum and count\n",
    "        sum_x = torch.sum(x_float, dim=0, keepdim=True)\n",
    "        count_x = torch.tensor(x.size(0), dtype=torch.float)\n",
    "\n",
    "        # Compute new mean and count\n",
    "        new_mean = sum_x / count_x\n",
    "        new_count = self.count.data + count_x\n",
    "\n",
    "        return new_mean, new_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReferenceContainer:\n",
    "    def __init__(self):\n",
    "        self.atlas_layer = None\n",
    "        self.vxm_model = None\n",
    "        self.pos_flow = None\n",
    "        self.neg_flow = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TemplateCreation(nn.Module):\n",
    "    \"\"\"\n",
    "    VoxelMorph network to generate an unconditional template image.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, inshape, nb_unet_features=None, mean_cap=100, atlas_feats=1, src_feats=1, mean_temp = None, load_weights = False, **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters:\n",
    "            inshape: Input shape. e.g. (192, 192, 192)\n",
    "            nb_unet_features: Unet convolutional features.\n",
    "                See VxmDense documentation for more information.\n",
    "            mean_cap: Cap for mean stream. Default is 100.\n",
    "            atlas_feats: Number of atlas/template features. Default is 1.\n",
    "            src_feats: Number of source image features. Default is 1.\n",
    "            kwargs: Forwarded to the internal VxmDense model.\n",
    "        \"\"\"\n",
    "        super(TemplateCreation, self).__init__()\n",
    "\n",
    "        # configure inputs\n",
    "        self.src_feats = src_feats\n",
    "\n",
    "        # pre-warp (atlas) model\n",
    "        if mean_temp.any() != None:\n",
    "            mean = torch.tensor(mean_temp,dtype = torch.float32)#.permute(1,2,0).unsqueeze(0)\n",
    "            self.atlas_layer = nn.Parameter(mean, requires_grad=True)\n",
    "        else:\n",
    "            self.atlas_layer = nn.Parameter(torch.randn(*inshape, atlas_feats) * 1e-7, requires_grad=True)\n",
    "        # warp model\n",
    "        self.vxm_model = VxmDense(inshape, nb_unet_features=nb_unet_features, bidir=True, **kwargs)\n",
    "        if load_weights == True:\n",
    "            checkpoint = torch.load('./models/1000.pt')\n",
    "            self.vxm_model.load_state_dict(checkpoint['model_state'], strict = False)\n",
    "            for name, param in self.vxm_model.named_parameters():\n",
    "                if 'flow' not in name:\n",
    "                    param.requires_grad = False\n",
    "        \n",
    "\n",
    "        # mean stream\n",
    "        self.mean_stream = MeanStream(mean_cap)\n",
    "        self.mean_stream.cuda()\n",
    "\n",
    "        # cache references\n",
    "        self.references = ReferenceContainer()\n",
    "        self.references.atlas_layer = self.atlas_layer\n",
    "        self.references.vxm_model = self.vxm_model\n",
    "\n",
    "    def forward(self, source_input):\n",
    "        # atlas tensor\n",
    "        atlas_tensor = self.atlas_layer.expand(source_input.shape[0], *self.atlas_layer.shape).permute(0,3,1,2)\n",
    "        \n",
    "        \n",
    "\n",
    "        # warp model\n",
    "        y_source, y_target, preint_flow = self.vxm_model(atlas_tensor, source_input)\n",
    "        #print(len(y_pred))\n",
    "        \n",
    "\n",
    "        # get mean stream of negative flow\n",
    "        mean_stream = self.mean_stream(-1*preint_flow)\n",
    "        \n",
    "\n",
    "        return y_source, y_target, mean_stream, preint_flow\n",
    "\n",
    "    def set_atlas(self, atlas):\n",
    "        \"\"\"\n",
    "        Sets the atlas weights.\n",
    "        \"\"\"\n",
    "        if len(atlas.shape) > len(self.atlas_layer.shape):\n",
    "            atlas = np.reshape(atlas, atlas.shape[1:])\n",
    "        self.atlas_layer.data = torch.from_numpy(atlas)\n",
    "\n",
    "    def get_atlas(self):\n",
    "        \"\"\"\n",
    "        Gets the atlas weights.\n",
    "        \"\"\"\n",
    "        return self.atlas_layer.data.squeeze().cpu().numpy()\n",
    "\n",
    "    def get_registration_model(self):\n",
    "        \"\"\"\n",
    "        Returns a reconfigured model to predict only the final transform.\n",
    "        \"\"\"\n",
    "        return lambda src, trg: self.vxm_model(src, trg, registration = True)#[-1]\n",
    "\n",
    "    def register(self, src, trg):\n",
    "        \"\"\"\n",
    "        Predicts the transform from src to trg tensors.\n",
    "        \"\"\"\n",
    "        return self.get_registration_model()(src, trg)\n",
    "\n",
    "    def apply_transform(self, src, trg, img, interp_method='linear', fill_value=None):\n",
    "        \"\"\"\n",
    "        Predicts the transform from src to trg and applies it to the img tensor.\n",
    "        \"\"\"\n",
    "        warp_model = self.get_registration_model()\n",
    "        return spatial_transformer(img, warp_model(src, trg), interp_method, fill_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mean_temp(invol):\n",
    "    mean = np.mean(invol, axis = 0)\n",
    "    stdev = np.std(invol,axis= (0,1,2))\n",
    "    return mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_temp = get_mean_temp(x_vols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_nf = [16, 32, 32, 32]\n",
    "dec_nf = [32, 32, 32, 32, 32, 16, 16]\n",
    "model = TemplateCreation(vol_shape, nb_unet_features=[enc_nf, dec_nf], mean_temp = mean_temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def volgen(\n",
    "    vol_names,\n",
    "    batch_size=1,\n",
    "    segs=None,\n",
    "    np_var='vol',\n",
    "    pad_shape=None,\n",
    "    resize_factor=1,\n",
    "    add_feat_axis=True\n",
    "):\n",
    "    \"\"\"\n",
    "    Base generator for random volume loading. Volumes can be passed as a path to\n",
    "    the parent directory, a glob pattern, a list of file paths, or a list of\n",
    "    preloaded volumes. Corresponding segmentations are additionally loaded if\n",
    "    `segs` is provided as a list (of file paths or preloaded segmentations) or set\n",
    "    to True. If `segs` is True, npz files with variable names 'vol' and 'seg' are\n",
    "    expected. Passing in preloaded volumes (with optional preloaded segmentations)\n",
    "    allows volumes preloaded in memory to be passed to a generator.\n",
    "\n",
    "    Parameters:\n",
    "        vol_names: Path, glob pattern, list of volume files to load, or list of\n",
    "            preloaded volumes.\n",
    "        batch_size: Batch size. Default is 1.\n",
    "        segs: Loads corresponding segmentations. Default is None.\n",
    "        np_var: Name of the volume variable if loading npz files. Default is 'vol'.\n",
    "        pad_shape: Zero-pads loaded volumes to a given shape. Default is None.\n",
    "        resize_factor: Volume resize factor. Default is 1.\n",
    "        add_feat_axis: Load volume arrays with added feature axis. Default is True.\n",
    "    \"\"\"\n",
    "\n",
    "    # convert glob path to filenames\n",
    "    if isinstance(vol_names, str):\n",
    "        if os.path.isdir(vol_names):\n",
    "            vol_names = os.path.join(vol_names, '*')\n",
    "        vol_names = glob.glob(vol_names)\n",
    "\n",
    "    if isinstance(segs, list) and len(segs) != len(vol_names):\n",
    "        raise ValueError('Number of image files must match number of seg files.')\n",
    "\n",
    "    while True:\n",
    "        # generate [batchsize] random image indices\n",
    "        indices = np.random.randint(len(vol_names), size=batch_size)\n",
    "\n",
    "        # load volumes and concatenate\n",
    "        load_params = dict(np_var=np_var, add_batch_axis=True, add_feat_axis=add_feat_axis,\n",
    "                           pad_shape=pad_shape, resize_factor=resize_factor)\n",
    "        imgs = [py.utils.load_volfile(vol_names[i], **load_params) for i in indices]\n",
    "        vols = [np.concatenate(imgs, axis=0)]\n",
    "\n",
    "        # optionally load segmentations and concatenate\n",
    "        if segs is True:\n",
    "            # assume inputs are npz files with 'seg' key\n",
    "            load_params['np_var'] = 'seg'  # be sure to load seg\n",
    "            s = [py.utils.load_volfile(vol_names[i], **load_params) for i in indices]\n",
    "            vols.append(np.concatenate(s, axis=0))\n",
    "        elif isinstance(segs, list):\n",
    "            # assume segs is a corresponding list of files or preloaded volumes\n",
    "            s = [py.utils.load_volfile(segs[i], **load_params) for i in indices]\n",
    "            vols.append(np.concatenate(s, axis=0))\n",
    "\n",
    "        yield tuple(vols)\n",
    "\n",
    "def template_creation(vol_names, bidir=False, batch_size=1, **kwargs):\n",
    "    \"\"\"\n",
    "    Generator for unconditional template creation.\n",
    "\n",
    "    Parameters:\n",
    "        vol_names: List of volume files to load, or list of preloaded volumes.\n",
    "        bidir: Yield input image as output for bidirectional models. Default is False.\n",
    "        batch_size: Batch size. Default is 1.\n",
    "        kwargs: Forwarded to the internal volgen generator.\n",
    "    \"\"\"\n",
    "    zeros = None\n",
    "    gen = volgen(vol_names, batch_size=batch_size, **kwargs)\n",
    "    while True:\n",
    "        scan = next(gen)[0]\n",
    "\n",
    "        # cache zeros\n",
    "        if zeros is None:\n",
    "            shape = scan.shape[1:-1]\n",
    "            zeros = np.zeros((1, *shape, len(shape)))\n",
    "\n",
    "        invols = [scan]\n",
    "        outvols = [scan, zeros, zeros, zeros] if bidir else [scan, zeros, zeros]\n",
    "        yield (invols, outvols)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = template_creation([\"./data/\"+ i for i in files], bidir=True, batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from voxelmorph.voxelmorph import losses as lss\n",
    "\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "\n",
    "# Define optimizer\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "# Function to compute total loss\n",
    "def compute_total_loss_(outputs, targets):\n",
    "\n",
    "    ncc_l1 = lss.MSE().loss(outputs[0], targets[0])\n",
    "    #print(ncc_l1)\n",
    "    ncc_l2 = lss.MSE().loss(model.references.atlas_layer.permute(2,0,1).unsqueeze(0), outputs[1][0].unsqueeze(0))\n",
    "    #print(ncc_l2)\n",
    "    mse_l = lss.MSE().loss(outputs[2], targets[2])\n",
    "    #print(mse_l)\n",
    "    lgrad = lss.Grad('l2', loss_mult=2).loss(outputs[3], targets[3])\n",
    "    #print(lgrad)\n",
    "    total_loss = ncc_l1+ ncc_l2+0.01*mse_l+lgrad\n",
    "    return total_loss\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def dice_loss(pred, target):\n",
    "    \"\"\"This definition generalize to real valued pred and target vector.\n",
    "This should be differentiable.\n",
    "    pred: tensor with first dimension as batch\n",
    "    target: tensor with first dimension as batch\n",
    "    \"\"\"\n",
    "\n",
    "    smooth = 1.\n",
    "\n",
    "    # have to use contiguous since they may from a torch.view op\n",
    "    iflat = pred.contiguous().view(-1)\n",
    "    tflat = target.contiguous().view(-1)\n",
    "    intersection = (iflat * tflat).sum()\n",
    "\n",
    "    A_sum = torch.sum(tflat * iflat)\n",
    "    B_sum = torch.sum(tflat * tflat)\n",
    "    \n",
    "    return ((2. * intersection + smooth) / (A_sum + B_sum + smooth) )\n",
    "    \n",
    "def test_dice():\n",
    "    dice = dice_loss#lss.Dice().loss\n",
    "    selected_samples_src = random.choices(x_vols,k = 100)\n",
    "    selected_samples_src = torch.tensor(selected_samples_src,dtype = torch.float32).permute(0,3,1,2)\n",
    "    selected_samples_trg = random.choices(x_vols, k= 100)\n",
    "    selected_samples_trg = torch.tensor(selected_samples_trg,dtype = torch.float32).permute(0,3,1,2)\n",
    "    template = torch.tensor(model.get_atlas(),dtype = torch.float32).unsqueeze(0).unsqueeze(0).cuda()\n",
    "    dice_score = 0\n",
    "    for i in range(100):\n",
    "        with torch.no_grad():\n",
    "            #print(template.shape)\n",
    "            #print(selected_samples_src[i].unsqueeze(0).shape)\n",
    "            src_to_temp = model.register(selected_samples_src[i].unsqueeze(0).cuda(), template)\n",
    "            #print(\"gets\")\n",
    "            #print(len(src_to_temp))\n",
    "            temp_to_trg = model.register(src_to_temp[0][0].unsqueeze(0),selected_samples_trg[i].unsqueeze(0).cuda())\n",
    "            #print(\"here\")\n",
    "            #print(dice(temp_to_trg[0], selected_samples_trg[i].unsqueeze(0).cuda()))\n",
    "            dice_score += dice(temp_to_trg[0][0], selected_samples_trg[i].cuda())\n",
    "            #print(dice(temp_to_trg[0][0], selected_samples_trg[i].cuda()))\n",
    "            #print(src_to_temp.shape)\n",
    "            #print(selected_samples_src[i].unsqueeze(0).shape)\n",
    "            \"\"\"plt.figure(figsize=(12, 3))\n",
    "            title = [\"Source\",\"Mean Atlas\", \"Source to Target\", \"Target\"]\n",
    "            # Iterate through each image\n",
    "            for i, img in enumerate([np.rot90(selected_samples_src[i][0].cpu().numpy(),-1),\n",
    "                                     np.rot90(model.get_atlas(),-1),\n",
    "                                     np.rot90(temp_to_trg[0][0][0].cpu().numpy(), -1),\n",
    "                                     np.rot90(selected_samples_trg[i][0].cpu().numpy(),-1)]):\n",
    "            \n",
    "                # Plot each image\n",
    "                plt.subplot(1, 4, i + 1)\n",
    "                plt.imshow(img, cmap=\"gray\")\n",
    "                plt.title(title[i])\n",
    "                plt.axis('off')\"\"\"\n",
    "            \n",
    "            # Show the plot\n",
    "            #plt.title(\"Atlas with PreLoaded RegNet Weights\")\n",
    "            #plt.show()\n",
    "            \n",
    "            \"\"\"plt.imshow(selected_samples_src[i][0].cpu().numpy(),cmap=\"grey\")\n",
    "            plt.title(\"Source\")\n",
    "            plt.show()\n",
    "            plt.imshow(model.get_atlas(),cmap = \"grey\")\n",
    "            plt.title(\"Atlas\")\n",
    "            plt.show()\n",
    "            #plt.imshow(src_to_temp[0][0][0].cpu().numpy(),cmap=\"grey\")\n",
    "            #plt.show()\n",
    "            plt.imshow(temp_to_trg[0][0][0].cpu().numpy(),cmap=\"grey\")\n",
    "            plt.title(\"Transformed Source\")\n",
    "            plt.show()\n",
    "            plt.imshow(selected_samples_trg[i][0].cpu().numpy(),cmap=\"grey\")\n",
    "            plt.title(\"Target\")\n",
    "            plt.show()\"\"\"\n",
    "            #asd\n",
    "    return dice_score / 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "epochs = 1000\n",
    "steps_per_epoch = 25\n",
    "model.cuda()\n",
    "total_dice_loss = []\n",
    "loss_list = []\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    total_loss = 0.0\n",
    "    for i, batch in enumerate(gen):\n",
    "        # Extract inputs and targets from the batch\n",
    "        inputs, targets = batch\n",
    "        #inputs = torch.FloatTensor(inputs)\n",
    "        #targets = torch.FloatTensor(targets)\n",
    "        \n",
    "        # Perform a training step\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        \n",
    "        inputs = torch.tensor(inputs[0], dtype = torch.float32).permute(0,3,1,2).cuda()\n",
    "\n",
    "        targets = [torch.from_numpy(d).cuda().float().permute(0, 3, 1, 2) for d in targets]#.permute(0,3,1,2)\n",
    "        outputs = model(inputs)\n",
    "\n",
    "        loss = compute_total_loss_(outputs, targets)\n",
    "        #print(loss)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss\n",
    "        \n",
    "        if i >= steps_per_epoch - 1:\n",
    "            break  \n",
    "    \n",
    "\n",
    "\n",
    "    loss_list.append(total_loss)\n",
    "    if epoch % 50 == 0:\n",
    "        print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss / steps_per_epoch}')\n",
    "    if epoch % 50 == 0:\n",
    "        dice_score = test_dice()\n",
    "        total_dice_loss.append(dice_score)\n",
    "        print(dice_score)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}