{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Variational Autoencoder for OSIC","metadata":{}},{"cell_type":"markdown","source":"Adapted some of Carlos' work on the 3DCNN autoencoder to do the VAE version (https://www.kaggle.com/carlossouza/osic-autoencoder-training). I've done something similar for the google landmarks recognition challenge so wasn't too much work to adapt it.","metadata":{}},{"cell_type":"markdown","source":"## What is a variational autoencoder?\n\nNormal autoencoders take an input and map it one-to-one with an output. However in the variational autoencoder setup, the encoder maps the input to a distribution (we assume that the input was 'generated' by this distribution). Usually this distribution is a multivariate gaussian so each of our latent dimensions is one dimension of the distribution.\n\nMapping the inputs to a gaussian is a kind of regularization. It is also what allows us to use VAEs to generate new images.","metadata":{}},{"cell_type":"code","source":"from tqdm.notebook import tqdm\nimport cv2\nimport copy\nimport matplotlib.pyplot as plt\nplt.style.use('ggplot')\nimport pydicom.pixel_data_handlers.gdcm_handler as gdcm_handler \nimport scipy\nimport cv2\nimport pydicom\nimport os\nfrom matplotlib import cm\nimport imageio\nfrom pathlib import Path\nfrom skimage.segmentation import clear_border\nfrom skimage.morphology import ball, disk, dilation, binary_erosion, remove_small_objects, erosion, closing, reconstruction, binary_closing\nfrom skimage.measure import label, regionprops\nfrom skimage.segmentation import clear_border\nfrom skimage.filters import roberts, sobel\nfrom scipy import ndimage as ndi\nfrom skimage import measure, morphology\nfrom scipy.stats import kurtosis\nimport seaborn as sns\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport torch\nimport torch.nn as nn\nfrom torch import optim\nfrom sklearn.model_selection import KFold\nimport random\nimport copy\nfrom torchvision import models\nimport torch.multiprocessing as mp\nfrom torch.nn import functional as F\nfrom torch.autograd import Variable\nfrom torch.utils.data import DataLoader, Dataset\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Seed for reproducibility of experiments","metadata":{}},{"cell_type":"code","source":"def seed_all(seed=42):\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    \nseed_all()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model and Loss function\n\nThe VAE loss function has two parts. \nThe familiar mean-squared-error: the difference between the input and output. This encourages the decoder to faithfully reconstruct the input.\nThe 'kullback-liebler divergence': a measure of the difference between two distributions. This term is what forces the encoding to be close to a gaussian.\n\nThe main difference between the autoencoder and variational autoencoder model in code is the encoder having both a mean - mu and variance (here log variance) - logvar output. In addition we have a 'reparameterize' function. This takes the encoded distribution defined by the mean and variance, generates a sample from the distribution, and uses that as the input to the decoder. \n\nThe rest of the model is pretty much the same as Carlos' autoencoder but I use 224x224 image size due to some experiments using resnet.","metadata":{}},{"cell_type":"code","source":"def vae_loss_function(recon_x, x, mu, logvar,KL_weight=1):\n    MSE = F.mse_loss(recon_x, x, reduction='none')\n    MSE = torch.div(MSE,torch.numel(x))\n    MSE = torch.sum(MSE)\n    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n    KLD = torch.div(KLD,x.shape[0])\n    return MSE, KLD\n\n\nclass VariationalAutoEncoder(nn.Module):\n    def __init__(self, latent_features=10):\n        super(VariationalAutoEncoder, self).__init__()\n        # Encoder\n        self.conv1 = nn.Conv3d(1, 16, 3)\n        self.conv2 = nn.Conv3d(16, 32, 3)\n        self.conv3 = nn.Conv3d(32, 96, 2)\n        self.conv4 = nn.Conv3d(96, 1, 1)\n        self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        self.pool2 = nn.MaxPool3d(kernel_size=3, stride=3, return_indices=True)\n        self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        self.pool4 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        self.fc1_mu = nn.Linear(8*8, latent_features)\n        self.fc1_logvar = nn.Linear(8*8, latent_features)\n        # Decoder\n        self.fc2 = nn.Linear(latent_features, 8*8)\n        self.deconv0 = nn.ConvTranspose3d(1, 96, 1)\n        self.deconv1 = nn.ConvTranspose3d(96, 32, 2)\n        self.deconv2 = nn.ConvTranspose3d(32, 16, 3)\n        self.deconv3 = nn.ConvTranspose3d(16, 1, 3)\n        self.unpool0 = nn.MaxUnpool3d(kernel_size=2, stride=2)\n        self.unpool1 = nn.MaxUnpool3d(kernel_size=2, stride=2)\n        self.unpool2 = nn.MaxUnpool3d(kernel_size=3, stride=3)\n        self.unpool3 = nn.MaxUnpool3d(kernel_size=2, stride=2)\n\n    def encode(self, x, return_partials=True):\n        # Encoder\n        x = self.conv1(x)\n        up3out_shape = x.shape\n        x, i1 = self.pool1(x)\n\n        x = self.conv2(x)\n        up2out_shape = x.shape\n        x, i2 = self.pool2(x)\n\n        x = self.conv3(x)\n        up1out_shape = x.shape\n        x, i3 = self.pool3(x)\n\n        x = self.conv4(x)\n        up0out_shape = x.shape\n        x, i4 = self.pool4(x)\n\n        x = x.view(-1, 8*8)\n        mu = F.relu(self.fc1_mu(x))\n        logvar = F.relu(self.fc1_logvar(x))\n\n        if return_partials:\n            return mu,logvar, up3out_shape, i1, up2out_shape, i2, up1out_shape, i3, \\\n                   up0out_shape, i4\n\n        else:\n            return x\n    \n    def reparameterize(self, mu, logvar):\n        if self.training:\n            std = logvar.mul(0.5).exp_()\n            eps = Variable(std.data.new(std.size()).normal_())\n            return eps.mul(std).add_(mu)\n        else:\n            return mu\n\n    def decode(self, z, up3out_shape, i1, up2out_shape, i2, \\\n        up1out_shape, i3, up0out_shape, i4):\n        # Decoder\n        x = F.relu(self.fc2(z))\n        x = x.view(-1, 1, 1, 8, 8)\n        x = self.unpool0(x, output_size=up0out_shape, indices=i4)\n        x = self.deconv0(x)\n        x = self.unpool1(x, output_size=up1out_shape, indices=i3)\n        x = self.deconv1(x)\n        x = self.unpool2(x, output_size=up2out_shape, indices=i2)\n        x = self.deconv2(x)\n        x = self.unpool3(x, output_size=up3out_shape, indices=i1)\n        x = self.deconv3(x)\n        return x\n        \n    def forward(self, x):\n        mu,logvar, up3out_shape, i1, up2out_shape, i2, \\\n        up1out_shape, i3, up0out_shape, i4 = self.encode(x)\n        z = self.reparameterize(mu, logvar)\n        x_reconst = self.decode(z,up3out_shape, i1, up2out_shape, i2, \\\n        up1out_shape, i3, up0out_shape, i4)\n        return x_reconst, z, mu, logvar","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_training_loss(train, val,title='loss'):\n    plt.figure()\n    plt.plot(train, label='Train')\n    plt.plot(val, label='Val')\n    if title=='loss':\n        plt.title('Model Training Loss')\n    else:\n        plt.title('Model Metric Loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.yscale('log')\n    plt.legend()\n    plt.savefig('training_loss')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Helper functions for preprocessing image data","metadata":{}},{"cell_type":"code","source":"# Load the scans in given folder path\ndef load_scan(path):\n\n    #slices = [pydicom.read_file(path / s) for s in os.listdir(path)]\n    slices = [pydicom.read_file(path / s) for s in os.listdir(path)]\n    slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))\n    try:\n        slice_thickness = np.abs(slices[0].ImagePositionPatient[2] - slices[1].ImagePositionPatient[2])\n    except:\n        slice_thickness = np.abs(slices[0].SliceLocation - slices[1].SliceLocation)\n    if slice_thickness==0:\n        slice_thickness=slices[0].SliceThickness\n    for s in slices:\n        s.SliceThickness = slice_thickness\n        \n    return slices\n\ndef get_pixels_hu(slices):\n    image = np.stack([np.array(s.pixel_array,dtype=np.int16) for s in slices])\n    # Convert to int16 (from sometimes int16), \n    # should be possible as values should always be low enough (<32k)\n    image = image.astype(np.int16)\n\n    # Set outside-of-scan pixels to 0\n    # The intercept is usually -1024, so air is approximately 0\n    image[image == -2000] = 0\n    \n    # Convert to Hounsfield units (HU)\n    for slice_number in range(len(slices)):\n        \n        intercept = slices[slice_number].RescaleIntercept\n        slope = slices[slice_number].RescaleSlope\n        \n        if slope != 1:\n            image[slice_number] = slope * image[slice_number].astype(np.float64)\n            image[slice_number] = image[slice_number].astype(np.int16)\n            \n        image[slice_number] += np.int16(intercept)\n    \n    return np.array(image, dtype=np.int16)\n\ndef resample(image, scan, new_spacing=[1,1,1]):\n    # Determine current pixel spacing\n    #spacing = np.array([scan[0].SliceThickness] + scan[0].PixelSpacing, dtype=np.float32)\n    spacing = np.array([scan[0].SliceThickness] + list(scan[0].PixelSpacing), dtype=np.float32)\n    resize_factor = spacing / new_spacing\n    new_real_shape = image.shape * resize_factor\n    new_shape = np.round(new_real_shape)\n    real_resize_factor = new_shape / image.shape\n    new_spacing = spacing / real_resize_factor\n    \n    image = scipy.ndimage.interpolation.zoom(image, real_resize_factor, mode='nearest')\n    return image, new_spacing\n\ndef get_segmented_lungs(im, plot=False):\n    \n    '''\n    This funtion segments the lungs from the given 2D slice.\n    '''\n    if plot == True:\n        f, plots = plt.subplots(8, 1, figsize=(5, 40))\n    '''\n    Step 1: Convert into a binary image. \n    '''\n    binary = im < -200\n    if plot == True:\n        plots[0].axis('off')\n        plots[0].imshow(binary, cmap=plt.cm.bone) \n    '''\n    Step 2: Remove the blobs connected to the border of the image.\n    '''\n    cleared = clear_border(binary)\n    if plot == True:\n        plots[1].axis('off')\n        plots[1].imshow(cleared, cmap=plt.cm.bone) \n    '''\n    Step 3: Label the image.\n    '''\n    label_image = label(cleared)\n    if plot == True:\n        plots[2].axis('off')\n        plots[2].imshow(label_image, cmap=plt.cm.bone) \n    '''\n    Step 4: Keep the labels with 2 largest areas.\n    '''\n    areas = [r.area for r in regionprops(label_image)]\n    areas.sort()\n    if len(areas) > 2:\n        for region in regionprops(label_image):\n            if region.area < areas[-2]:\n                for coordinates in region.coords:                \n                       label_image[coordinates[0], coordinates[1]] = 0\n    binary = label_image > 0\n    if plot == True:\n        plots[3].axis('off')\n        plots[3].imshow(binary, cmap=plt.cm.bone) \n    '''\n    Step 5: Erosion operation with a disk of radius 2. This operation is \n    seperate the lung nodules attached to the blood vessels.\n    '''\n    selem = disk(2)\n    binary = binary_erosion(binary, selem)\n    if plot == True:\n        plots[4].axis('off')\n        plots[4].imshow(binary, cmap=plt.cm.bone) \n    '''\n    Step 6: Closure operation with a disk of radius 10. This operation is \n    to keep nodules attached to the lung wall.\n    '''\n    selem = disk(10)\n    binary = binary_closing(binary, selem)\n    if plot == True:\n        plots[5].axis('off')\n        plots[5].imshow(binary, cmap=plt.cm.bone) \n    '''\n    Step 7: Fill in the small holes inside the binary mask of lungs.\n    '''\n    edges = roberts(binary)\n    binary = ndi.binary_fill_holes(edges)\n    if plot == True:\n        plots[6].axis('off')\n        plots[6].imshow(binary, cmap=plt.cm.bone) \n    '''\n    Step 8: Superimpose the binary mask on the input image.\n    '''\n    get_high_vals = binary == 0\n    im[get_high_vals] = 0\n    if plot == True:\n        plots[7].axis('off')\n        plots[7].imshow(im, cmap=plt.cm.bone) \n        \n    return im\n\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Helper functions for loading and preprocessing data","metadata":{}},{"cell_type":"code","source":"def load_and_prepare_data(add_pixel_stats=True):\n    train=pd.read_csv('../input/osic-pulmonary-fibrosis-progression/train.csv')\n    test=pd.read_csv('../input/osic-pulmonary-fibrosis-progression/test.csv')\n    submission=pd.read_csv('../input/osic-pulmonary-fibrosis-progression/sample_submission.csv')\n    \n    #Prepare Train Data\n    train['base_Weeks']=train.groupby(['Patient'])['Weeks'].transform('min')\n    base=train[train.Weeks==train.base_Weeks]\n    base = base.rename(columns={'FVC': 'base_FVC','Percent': 'base_Percent'})\n    base.drop_duplicates(subset=['Patient', 'Weeks'], keep='first',inplace=True)\n    train=train.merge(base[['Patient','base_FVC','base_Percent']],on='Patient',how='left')\n    train['Week_passed'] = train['Weeks'] - train['base_Weeks']\n    \n    test = test.rename(columns={'Weeks': 'base_Weeks', 'FVC': 'base_FVC','Percent': 'base_Percent'})\n    # Adding Sample Submission\n    submission = pd.read_csv(\"../input/osic-pulmonary-fibrosis-progression/sample_submission.csv\")\n    # In submisison file, format: ID_'week', using lambda to split the ID\n    submission['Patient'] = submission['Patient_Week'].apply(lambda x:x.split('_')[0])\n    # In submisison file, format: ID_'week', using lambda to split the Week\n    submission['Weeks'] = submission['Patient_Week'].apply(lambda x:x.split('_')[1]).astype(int)\n    test = submission.drop(columns = [\"FVC\", \"Confidence\"]).merge(test, on = 'Patient')\n    test['Week_passed'] = test['Weeks'] - test['base_Weeks']\n    test=test[train.columns.drop(['FVC','Percent'])]\n    \n    if add_pixel_stats:\n        pixel_stats=pd.read_csv('../input/osic-histogram-features/train_pixel_stats.csv')\n        train=train.merge(pixel_stats[['Patient','kurtosis','std','mean','median']],how='left',on='Patient')\n        test_ids=test.Patient.unique()\n        root_dir = Path('/kaggle/input/osic-pulmonary-fibrosis-progression')\n        ct_scans_dir=root_dir/'test'\n        pixel_stats_test=test.copy()\n        pixel_stats_test.drop_duplicates(subset=['Patient'],inplace=True)\n        k,s,m,me=get_kurtosis_stats(test_ids,ct_scans_dir)\n        pixel_stats_test['kurtosis']=np.array(k)\n        pixel_stats_test['std']=np.array(s)\n        pixel_stats_test['mean']=np.array(m)\n        pixel_stats_test['median']=np.array(me)\n        test=test.merge(pixel_stats_test[['Patient','kurtosis','std','mean','median']],how='left',on='Patient')\n    return train, test\n\ndef OH_encode(train,test):\n    #OH Encoding of categorical variables (https://www.kaggle.com/ulrich07/osic-keras-starter-with-custom-metrics)\n    COLS = ['Sex','SmokingStatus']\n    for col in COLS:\n        for mod in train[col].unique():\n            train[mod] = (train[col] == mod).astype(int)\n            test[mod] = (test[col] == mod).astype(int)\n        train.drop(col,axis=1,inplace=True)\n        test.drop(col,axis=1,inplace=True)\n    return train, test\n\ndef Scale(train):\n    from sklearn import preprocessing\n    robust_scaler = preprocessing.RobustScaler()\n    train.loc[:,train.columns.difference(['Patient','FVC','Percent','Weeks','base_Weeks'])]=robust_scaler.fit_transform(train.loc[:,train.columns.difference(['Patient','FVC','Percent','Weeks','base_Weeks'])])\n    return robust_scaler   ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Pytorch Dataset class\nNeatened up since some of my other notebooks!","metadata":{}},{"cell_type":"code","source":"class OSIC(Dataset):\n    def __init__(self,patient_ids,df,scaler=None,train=True,add_pixel_stats=True,impute_vals=None):\n        root_dir = Path('/kaggle/input/osic-pulmonary-fibrosis-progression')\n        self.df=df.copy()\n        self.df=self.df.loc[self.df.Patient.isin(patient_ids),:]\n        if not train:\n            ct_scans_dir=root_dir/'test'\n        else:\n            ctscans_dir=root_dir/'train'\n        self.df.loc[:,self.df.columns.difference(['Patient','FVC','Percent','Weeks','base_Weeks'])]=scaler.transform(self.df.loc[:,self.df.columns.difference(['Patient','FVC','Percent','Weeks','base_Weeks'])])\n        self.data=self.df.loc[:,self.df.columns.difference(['FVC','Patient','Percent'])].values\n        if train:\n            self.impute_vals=np.nanmean(self.data, axis=0)\n        else:\n            self.impute_vals=impute_vals\n        inds = np.where(np.isnan(self.data))\n        self.data[inds] = np.take(self.impute_vals, inds[1])\n        self.patients=self.df['Patient'].unique()\n        self.train=train\n        if self.train:\n            self.fvc=self.df['FVC'].values\n    \n    def __len__(self):\n        return len(self.patients)\n\n    def __getitem__(self, idx):\n        patient_id = self.patients[idx]\n        if os.path.isfile('/kaggle/input/processed-osic/processed/'+patient_id+'.npy'):\n            image = np.load('/kaggle/input/processed-osic/processed/'+patient_id+'.npy') \n        else:\n            try:\n                patient_path= ctscans_dir / patient_id\n                scan = load_scan(patient_path)\n                image=get_pixels_hu(scan)\n                image, new_spacing = resample(image, scan, new_spacing=[1,1,1])\n                image = segment_lung_mask(image, False)\n            except:\n                image=np.zeros((1,224,224))\n       \n        image=image.astype(np.uint8)\n        padded_image=np.zeros((1,50,224,224))\n        max_ind_i=min(50,image.shape[0])\n        max_ind_j=min(224,image.shape[1])\n        max_ind_k=min(224,image.shape[2])\n        padded_image[:,:max_ind_i,:max_ind_j,:max_ind_k] = image[:max_ind_i,:max_ind_j,:max_ind_k]\n        all_fvc=np.zeros(146)\n        patient_weeks=self.df['Weeks'][self.df.Patient==patient_id].values+12\n        all_fvc[patient_weeks]=self.fvc[self.df.Patient==patient_id]\n        base_fvc=self.fvc[self.df.Patient==patient_id][0]\n        if self.train:\n            data = {'data': self.data[idx],\n                    'image': padded_image,\n                    'allfvc': all_fvc,\n                    'base_fvc': base_fvc}\n        else:\n            \n            data = {'data': self.data[idx]}\n        return data","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Functions to train and validate a single epoch (pass over the data)","metadata":{}},{"cell_type":"code","source":"def train_epoch(model,optimizer,train_loader,KL_weight):\n    criterion = nn.CrossEntropyLoss()\n    epoch_losses=[]\n    model.train()\n    epoch_loss=0\n    epoch_kld=0\n    epoch_mse=0\n    for i, data in enumerate(train_loader): \n        batch_size,_, _, _, _ = data['image'].shape\n        optimizer.zero_grad()\n        img = data['image'].float().to(device)\n        X_reconst, z, mu, logvar = model(img)  # VAE\n        #NOTE THE WARM-UP on the VAE LOSS function\n        MSE, KLD = vae_loss_function(X_reconst, img, mu, logvar,KL_weight)\n        loss = KL_weight*KLD+MSE\n        epoch_loss+=loss.item()\n        epoch_mse+=MSE.item()\n        epoch_kld+=KLD.item()\n        loss.backward()\n        optimizer.step()\n        \n    return model,optimizer,epoch_loss,epoch_mse,epoch_kld\n\ndef val_epoch(model,val_loader):\n    criterion = nn.CrossEntropyLoss()\n    all_y,all_z, all_mu, all_logvar = [], [], [], []\n    model.eval()\n    epoch_loss=0\n    epoch_kld=0\n    epoch_mse=0\n    for i, data in enumerate(val_loader): \n        batch_size, _, _, _, _ = data['image'].shape\n        img = data['image'].float().to(device)\n        X_reconst, z, mu, logvar = model(img)  # VAE\n        MSE, KLD = vae_loss_function(X_reconst, img, mu, logvar)\n        loss = KLD+MSE\n        epoch_loss+=loss.item()\n        epoch_mse+=MSE.item()\n        epoch_kld+=KLD.item()\n        y=data['base_fvc']\n        all_y.extend(y.data.cpu().numpy())\n        all_z.extend(z.data.cpu().numpy())\n        all_mu.extend(mu.data.cpu().numpy())\n        all_logvar.extend(logvar.data.cpu().numpy())\n    return epoch_loss,epoch_mse,epoch_kld, all_y,all_z, all_mu, all_logvar","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Outer loop to train a model","metadata":{}},{"cell_type":"code","source":"def train_model(ids,train,max_epochs=100, patience=40,batch_size=1, plot_losses=True,add_pixel_stats=True):\n    \n    np.random.shuffle(ids)\n    train_ids,val_ids=np.split(ids, [int(round(0.9 * len(ids), 0))])\n    \n    scaler=Scale(train.loc[train.Patient.isin(train_ids),:].copy())\n    train_dataset = OSIC(train_ids,train,scaler=scaler,add_pixel_stats=add_pixel_stats)  \n    train_dataloader = DataLoader(train_dataset, batch_size=batch_size,shuffle=True,num_workers=3)\n    val_dataset = OSIC(val_ids,train,scaler=scaler,impute_vals=train_dataset.impute_vals,add_pixel_stats=add_pixel_stats)  \n    val_dataloader = DataLoader(val_dataset, batch_size=batch_size,shuffle=True,num_workers=3)\n    model = VariationalAutoEncoder().to(device)\n    print('Number of parameters:')\n    print(sum(p.numel() for p in model.parameters() if p.requires_grad))\n    \n    epoch_train_metric=[]\n    epoch_val_metric=[]\n    epoch_train_loss=[]\n    epoch_val_loss=[]\n    \n    min_val_loss = 1e+100\n    min_val_metric = 1e+100\n    early_stop = False\n    epoch=0\n    optimizer = optim.Adam(model.parameters(),lr=5e-4)\n    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.05)\n    while epoch<max_epochs and not early_stop:\n        epoch+=1\n        if epoch<5:\n            KL_weight=0\n        else:\n            KL_weight=epoch/max_epochs\n        model, optimizer, train_loss,train_mse,train_kld = train_epoch(model,optimizer,train_dataloader,KL_weight)\n        scheduler.step()\n        print('====> Epoch: {} Average train loss: {:.4f}'.format(\n                        epoch, train_loss/(batch_size*len(train_dataloader))))\n        print('====> Epoch: {} Average train mse: {:.4f}'.format(\n                        epoch, train_mse/(batch_size*len(train_dataloader))))\n        print('====> Epoch: {} Average train kld: {:.4f}'.format(\n                        epoch, train_kld/(batch_size*len(train_dataloader))))\n\n        val_loss,val_mse,val_kld, all_y,all_z, all_mu, all_logvar = val_epoch(model,val_dataloader)\n        print('====> Epoch: {} Average val loss: {:.4f}'.format(\n                        epoch, val_loss/(batch_size*len(val_dataloader))))\n        print('====> Epoch: {} Average val mse: {:.4f}'.format(\n                        epoch, val_mse/(batch_size*len(val_dataloader))))\n        print('====> Epoch: {} Average val kld: {:.4f}'.format(\n                        epoch, val_kld/(batch_size*len(val_dataloader))))\n        \n        \n        epoch_train_loss.append(train_loss)\n        epoch_val_loss.append(val_loss)\n        \n    if plot_losses:\n        plot_training_loss(epoch_train_loss, epoch_val_loss)\n        plot_training_loss(epoch_train_metric, epoch_val_metric,title='metric')\n    return model, scaler, train_dataset.impute_vals, train_ids, val_ids, all_y,all_z, all_mu, all_logvar","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train,test=load_and_prepare_data(add_pixel_stats=False)\ntrain,test=OH_encode(train,test)\nids=train.Patient.unique()\nmodel, scaler, impute_vals, train_ids, val_ids,all_y,all_z, all_mu, all_logvar=train_model(ids,train,max_epochs=100, patience=40,batch_size=4, plot_losses=True,add_pixel_stats=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inspection\nWhen we wish to inspect the output of the model, we can take just the mean of the encoded distribution as the input to the decoder.","metadata":{}},{"cell_type":"code","source":"slc = 0.5\nval_dataset = OSIC(val_ids,train,scaler=scaler,impute_vals=impute_vals)\nsample_id = np.random.randint(len(val_dataset))\nprint(f'Inspecting CT Scan {val_dataset.patients[sample_id]}')\n\nfig, axs = plt.subplots(1, 2, figsize=(10, 7))\n\nsample = val_dataset[sample_id]['image'].squeeze(0)\naxs[0].imshow(sample[int(40 * slc), :, :], cmap=cm.bone)\naxs[0].axis('off')\nimageio.mimsave(\"sample_input.gif\", sample, duration=0.0001)\n\nwith torch.no_grad():\n    img = torch.tensor(val_dataset[sample_id]['image']).unsqueeze(0).float().to(device)\n    latent_features = model.encode(img, return_partials=False)\\\n        .squeeze().cpu().numpy().tolist()\n    outputs = model(img)[0].squeeze().cpu().numpy()\n\naxs[1].imshow(outputs[int(40 * slc), :, :], cmap=cm.bone)\naxs[1].axis('off')\n\nimageio.mimsave(\"sample_output.gif\", outputs, duration=0.0001)\n\nrmse = ((sample - outputs)**2).mean()\nplt.show()\nprint(f'Latent features: {latent_features} \\nLoss: {rmse}')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"l=((outputs[int(40 * slc), :, :]-sample[int(40 * slc), :, :])**2)\n(l/l.size).sum()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from IPython.display import HTML\nHTML('<br/><img src=\"https://i.ibb.co/gFxgRq6/sample-input.gif\" style=\"float: left; width: 30%; margin-right: 1%; margin-bottom: 0.5em;\">'\n     '<img src=\"https://i.ibb.co/Jm57fWw/sample-output.gif\" style=\"float: left; width: 30%; margin-right: 1%; margin-bottom: 0.5em;\">'\n     '<p style=\"clear: both;\">')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## What else might we try?\n\n\n### Identify clusters of patients in the latent space?\n","metadata":{}},{"cell_type":"code","source":"from sklearn.manifold import TSNE\n\nz_embed = TSNE(n_components=2, n_iter=12000).fit_transform(all_z)\n\nfig = plt.figure(figsize=(12, 10))\nplots = []\nplt.scatter(z_embed[:, 0], z_embed[:, 1], c=all_y, s=8)\n\nplt.axis('off')\nplt.title('t-SNE: 2-dim')\n#plt.savefig(\"./ResNetVAE_{}_embedded_plot.png\".format(exp), bbox_inches='tight', dpi=600)\nplt.show()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Identify anomalous images \n\nIf a particular patient is outside of the usual distribution for lungs, then the reconstruction error is likely to be higher. We could use this to identify patients with an unusual PF trajectory.\n\nOr with a 2D version of the VAE on either the full lung or patches of the lung, we could identify regions that are unusual and perhaps contain symtpoms of PF.\n\nTBC\n","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}