first commit without models

README.md

@@ -1,8 +1,8 @@
 ---
-title: TEDM Demo
-emoji: 🔥
-colorFrom: indigo
-colorTo: blue
+title: TEDM
+emoji: 🐨
+colorFrom: purple
+colorTo: yellow
 sdk: gradio
 sdk_version: 3.35.2
 app_file: app.py
@@ -10,4 +10,64 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Timestep ensembling diffusion models for semi-supervised image segmentation
+
+Results
+
+| Training data size | 1 (1\%)                                        | 3 (2\%)                 | 6 (3\%)                 | 12 (96\%)                | 197 (100\%)             |
+|:------------------|:----------------------------------------------:|:-----------------------:|:-----------------------:|:-----------------------:|:-----------------------:|
+|                     |JSRT (labelled in-domain)  |
+| Baseline           | 84.4 $\pm$ 5.4                                 | 91.7 $\pm$ 3.7          | 93.3 $\pm$ 2.9          | 95.3 $\pm$ 2.3          | 97.3 $\pm$ 1.2          |
+| LEDM               | 90.8 $\pm$ 3.5                                 | 94.1 $\pm$ 1.6          | 95.5 $\pm$ 1.4          | 96.4 $\pm$ 1.4          | 97.0 $\pm$ 1.3          |
+| LEDMe              | **93.7 $\pm$ 2.6**                        | **95.5 $\pm$ 1.5** | **96.7 $\pm$ 1.5** | **97.0 $\pm$ 1.1** | **97.6 $\pm$ 1.2** |
+| Ours               | **93.1 $\pm$ 3.4**                        | 94.8 $\pm$ 1.4          | 95.8 $\pm$ 1.2          | 96.6 $\pm$ 1.1          | 97.3 $\pm$ 1.2          |
+|                    |NIH (unlabelled in-domain) |
+| Baseline           | 68.5 $\pm$ 12.8                                | 71.2 $\pm$ 15.1         | 71.4 $\pm$ 15.9         | 77.8 $\pm$ 14.0         | 81.5 $\pm$ 12.7         |
+| LEDM               | 63.3 $\pm$ 12.2                                | 78.0 $\pm$ 10.1         | 81.2 $\pm$ 9.3          | 85.9 $\pm$ 7.4          | 88.9 $\pm$ 5.9          |
+| LEDMe              | 70.3 $\pm$ 11.4                                | 78.3 $\pm$ 9.8          | 83.0 $\pm$ 8.6          | 84.4 $\pm$ 8.1          | 90.1 $\pm$ 5.3          |
+| Ours               | **80.3 $\pm$ 9.0**                        | **86.4 $\pm$ 6.2** | **89.2 $\pm$ 5.5** | **91.3 $\pm$ 4.1** | **92.9 $\pm$ 3.2** |
+|                    | Montgomery (out-of-domain) |
+| Baseline           | 77.1 $\pm$ 12.0                                | 83.0 $\pm$ 12.2         | 80.9 $\pm$ 14.7         | 83.8 $\pm$ 14.9         | 94.1 $\pm$ 6.6          |
+| LEDM               | 79.3 $\pm$ 8.1                                 | 85.9 $\pm$ 7.4          | 89.4 $\pm$ 6.7          | 92.3 $\pm$ 7.2          | 94.4 $\pm$ 7.2          |
+| LEDMe              | 80.7 $\pm$ 6.6                                 | 86.3 $\pm$ 6.5          | 89.5 $\pm$ 5.9          | 91.2 $\pm$ 5.6          | **95.3 $\pm$ 4.0** |
+| Ours               | **90.5 $\pm$ 5.3**                        | **91.4 $\pm$ 6.1** | **93.3 $\pm$ 6.0** | **94.6 $\pm$ 6.0** | 95.1 $\pm$ 6.9          |
+
+## Training
+
+- training the backbone
+
+```python train.py --dataset CXR14 --data_dir <PATH TO CXR14 DATASET>```
+
+- our method
+
+```python train.py --experiment TEDM --data_dir <PATH TO JSRT DATASET> --n_labelled_images <TRAINING SET SIZE>```
+
+- LEDM method
+
+```python train.py --experiment LEDM --data_dir <PATH TO JSRT DATASET> --n_labelled_images <TRAINING SET SIZE>```
+
+- LEDMe method
+
+```python train.py --experiment LEDMe --data_dir <PATH TO JSRT DATASET> --n_labelled_images <TRAINING SET SIZE>```
+
+- baseline method
+
+```python train.py --experiment JSRT_baseline --data_dir <PATH TO JSRT DATASET> --n_labelled_images <TRAINING SET SIZE>```
+
+## Testing
+
+- update 
+    - `DATADIR` in paths `dataloaders/JSRT.py`, `dataloaders/NIH.py` and `dataloaders/Montgomery.py`
+    - `NIHPATH`, `NIHFILE`, `MONPATH` and `MONFILE` in paths `auxiliary/postprocessing/run_tests.py` and `auxiliary/postprocessing/testing_shared_weights.py`
+
+- for baseline and LEDM methods, run
+
+```python auxiliary/postprocessing/run_tests.py --experiment <PATH TO LOG FOLDER>```
+
+- for our method, run
+
+```python auxiliary/postprocessing/testing_shared_weights.py --experiment <PATH TO LOG FOLDER>```
+
+## Figures and reporting
+
+VS Code notebooks can be found in `auxiliary/notebooks_and_reporting`.

app.py

@@ -0,0 +1,192 @@
+import numpy as np
+import gradio as gr
+from PIL import Image
+import torch
+from torch import nn
+from einops.layers.torch import Rearrange
+from torchvision import transforms
+from models.unet_model import Unet
+from models.datasetDM_model import DatasetDM
+from skimage import measure, segmentation
+import cv2
+from tqdm import tqdm
+from einops import repeat
+
+img_size = 128
+font = cv2.FONT_HERSHEY_SIMPLEX
+
+
+## %%
+def load_img(img_file):
+        # assert type of input
+    if isinstance(img_file, np.ndarray):
+        img = torch.Tensor(img_file).float()
+        # make sure img is between 0 and 1
+        if img.max() > 1:
+            img /= 255
+        # resize 
+        img = transforms.Resize(img_size)(img)
+    elif isinstance(img_file, str):
+        img = Image.open(img_file).convert('L').resize((img_size, img_size))
+        img = transforms.ToTensor()(img).float()
+    elif isinstance(img_file, Image.Image):
+        img = img_file.convert('L').resize((img_size, img_size))
+        img = transforms.ToTensor()(img).float()
+    else:
+        raise TypeError("Input must be a numpy array, PIL image, or filepath")
+    if len(img.shape) == 2:
+        img = img[None, None]
+    elif len(img.shape) == 3:
+        img = img[None]
+    else:
+        raise ValueError("Input must be a 2D or 3D array")
+    return img
+
+def predict_baseline(img, checkpoint_path):
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    config = checkpoint["config"]
+    baseline = Unet(**vars(config))
+    baseline.load_state_dict(checkpoint["model_state_dict"])
+    baseline.eval()
+    return (torch.sigmoid(baseline(img)) > .5).float().squeeze().numpy()
+
+def predict_LEDM(img, checkpoint_path):
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    config = checkpoint["config"]
+    config.verbose = False
+    LEDM = DatasetDM(config)
+    LEDM.load_state_dict(checkpoint["model_state_dict"])
+    LEDM.eval()
+    return (torch.sigmoid(LEDM(img)) > .5).float().squeeze().numpy()
+
+def predict_TEDM(img, checkpoint_path):
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    config = checkpoint["config"]
+    config.verbose = False
+    TEDM = DatasetDM(config)
+    TEDM.classifier = nn.Sequential(
+        Rearrange('b (step act) h w -> (b step) act h w', step=len(TEDM.steps)),
+        nn.Conv2d(960, 128, 1),
+        nn.ReLU(),
+        nn.BatchNorm2d(128),
+        nn.Conv2d(128, 32, 1),
+        nn.ReLU(),
+        nn.BatchNorm2d(32),
+        nn.Conv2d(32, 1, config.out_channels)
+        )
+    TEDM.load_state_dict(checkpoint["model_state_dict"])
+    TEDM.eval()
+    return (torch.sigmoid(TEDM(img)).mean(0) > .5).float().squeeze().numpy()
+
+predictors = {'Baseline': predict_baseline, 
+              'Global CL': predict_baseline,
+              'Global & Local CL': predict_baseline,
+              'LEDM': predict_LEDM, 
+              'LEDMe': predict_LEDM,
+              'TEDM': predict_TEDM}
+model_folders = {
+    'Baseline': 'baseline',
+    'Global CL': 'global_finetune',
+    'Global & Local CL': 'glob_loc_finetune',
+    'LEDM': 'LEDM',
+    'LEDMe': 'LEDMe',
+    'TEDM': 'TEDM'
+}
+
+
+def postprocess(pred, img):
+    all_labels = measure.label(pred, background=0)
+    _, cn = np.unique(all_labels, return_counts=True)
+    # find the two largest connected components that are not the background
+    if len(cn) >= 3:
+        lungs = np.argsort(cn[1:])[-2:] + 1
+        all_labels[(all_labels!=lungs[0]) & (all_labels!=lungs[1])] = 0
+        all_labels[(all_labels==lungs[0]) | (all_labels==lungs[1])] = 1
+    # put all_labels into a cv2 object
+    if len(cn) > 1:
+        img = segmentation.mark_boundaries(img, all_labels, color=(1,0,0), mode='outer', background_label=0)
+    else:
+        img = repeat(img, 'h w -> h w c', c=3)
+    return img
+
+
+
+def predict(img_file, models:list, training_sizes:list, seg_img=False, progress=gr.Progress()):
+    max_progress = len(models) * len(training_sizes)
+    n_progress = 0
+    progress((n_progress, max_progress), desc="Starting")
+    img = load_img(img_file)
+    print(img.shape)
+    preds = []
+    # sorting models so that they show as  baseline - LEDM - LEDMe - TEDM
+    models = sorted(models, key=lambda x: 0 if x == 'Baseline' else 1 if x == 'Global CL' else 2 if x == 'Global & Local CL' else 3 if x == 'LEDM' else 4 if x == 'LEDMe' else 5)
+    
+    for model in models:
+        print(model)
+        model_preds = []
+        for training_size in sorted(training_sizes):
+            #if n_progress < max_progress:
+            progress((n_progress, max_progress) , desc=f"Predicting {model} {training_size}")
+            n_progress += 1
+            print(training_size)
+            out = predictors[model](img, f"logs/{model_folders[model]}/{training_size}/best_model.pt")
+            writing_colour = (.5,.5,.5)
+            if seg_img:
+                out = postprocess(out, img.squeeze().numpy())
+                writing_colour = (1,1,1)
+            out = cv2.putText(np.array(out),f"{model} {training_size}",(5,125), font, .5, writing_colour,1, cv2.LINE_AA)
+            #ImageDraw.Draw(out).text((0,128), f"{model} {training_size}", fill=(255,0,0))
+            model_preds.append(np.asarray(out))
+        preds.append(np.concatenate(model_preds, axis=1))
+    prediction = np.concatenate(preds, axis=0)
+    if (prediction.shape[1] <=128*2):
+        pad = (330 - prediction.shape[1])//2
+        if len(prediction.shape) == 2:
+            prediction = np.pad(prediction, ((0,0), (pad, pad)), 'constant', constant_values=1)
+        else:
+            prediction = np.pad(prediction, ((0,0), (pad, pad), (0,0)), 'constant', constant_values=1)
+    return prediction
+
+
+## %%
+input = gr.Image( label="Chest X-ray", shape=(img_size, img_size), type="pil")
+output = gr.Image(label="Segmentation", shape=(img_size, img_size))
+## %%
+demo = gr.Interface(
+    fn=predict,
+    inputs=[input, 
+            gr.CheckboxGroup(["Baseline", "Global CL", "Global & Local CL", "LEDM", "LEDMe", "TEDM"], label="Model", value=["Baseline", "LEDM", "LEDMe", "TEDM"]),
+            gr.CheckboxGroup([1,3,6,12,197], label="Training size", value=[1,3,6,12,197]),
+            gr.Checkbox(label="Show masked image (otherwise show binary segmentation)", value=True),],
+
+    outputs=output,
+    examples = [
+    ['img_examples/NIH_0006.png'], 
+    ['img_examples/NIH_0076.png'], 
+    ["img_examples/00016568_041.png"], 
+    ['img_examples/NIH_0024.png'], 
+    ['img_examples/00015548_000.png'], 
+    ['img_examples/NIH_0019.png'], 
+    ['img_examples/NIH_0094.png'],
+    ['img_examples/NIH_0051.png'], 
+    ['img_examples/NIH_0012.png'], 
+    ['img_examples/NIH_0014.png'], 
+    ['img_examples/NIH_0055.png'], 
+    ['img_examples/NIH_0035.png'], 
+                ],
+    title="Chest X-ray Segmentation with TEDM.",
+    description="""<img src="file/img_examples/TEDM-model-visualisation.png"
+     alt="Markdown Monster icon"
+     style="margin-right: 10px;" />"""+
+    "\nMedical image segmentation is a challenging task, made more difficult by many datasets' limited size and annotations. Denoising diffusion probabilistic models (DDPM) have recently shown promise in modelling " + 
+    "the distribution of natural images and were successfully applied to various medical imaging tasks. This work focuses on semi-supervised image segmentation using diffusion models, particularly addressing domain " + 
+    "generalisation. Firstly, we demonstrate that smaller diffusion steps generate latent representations that are more robust for downstream tasks than larger steps. Secondly, we use this insight to propose an improved " + 
+    "esembling scheme that leverages information-dense small steps and the regularising effect of larger steps to generate predictions. Our model shows significantly better performance in domain-shifted settings while " +  
+    "retaining competitive performance in-domain. Overall, this work highlights the potential of DDPMs for semi-supervised medical image segmentation and provides insights into optimising their performance under domain shift."+ 
+    "\n\n\n When choosing 'Show masked image', we post-process the segmentation by choosing up to two largest connected components and drawing their outline. "+
+    "\nNote that each model takes 10-35 seconds to run on CPU. Choosing all models and all training sizes will take some time. "+
+    "We noticed that gradio sometimes fails on the first try. If it doesn't work, try again.",
+    cache_examples=False,
+)
+demo.queue().launch(debug=True)
+#demo.queue().launch(share=True)

auxiliary/notebooks_and_reporting/generate_figures.py

@@ -0,0 +1,175 @@
+# %%
+import numpy as np
+import torch
+from pathlib import Path
+import os
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+HEAD = Path(os.getcwd()).parent.parent
+
+if __name__=="__main__":
+    # load baseline and LEDM data
+    metrics = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[]}
+    files_needed = ["JSRT_val_predictions.pt", "JSRT_test_predictions.pt",  "NIH_predictions.pt", "Montgomery_predictions.pt",]
+    head = HEAD / 'logs'
+    for exp in ['baseline', 'LEDM']:
+         for datasize in [1, 3, 6, 12, 24, 49, 98, 197]:
+            if len(set(files_needed) - set(os.listdir(head / exp / str(datasize)))) == 0:
+                print(f"Experiment {exp} {datasize}")
+                output = torch.load(head / exp / str(datasize) / "JSRT_val_predictions.pt")
+                print(f"{output['dice'].mean()}\t{output['dice'].std()}")
+                for file in files_needed[1:]:
+                    output = torch.load(head / exp / str(datasize)  / file)
+                    metrics_datasize = 197 if datasize == "None" else int(datasize)
+                    metrics["dice"].append(output["dice"].numpy())
+                    metrics["precision"].append(output["precision"].numpy())
+                    metrics["recall"].append(output["recall"].numpy())
+                    metrics["exp"].append(np.array([exp] * len(output["dice"])))
+                    metrics["datasize"].append(np.array([int(datasize)] * len(output["dice"])))
+                    metrics["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+            else:
+                    print(f"Experiment {exp} is missing files")
+
+    for key in metrics:
+        metrics[key] = np.concatenate([el.squeeze() for el in metrics[key]])
+    df = pd.DataFrame(metrics)
+    df.head()
+
+
+    # %% Load step data
+    metrics2 = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[], 'timestep':[]}
+    for timestep in [1, 10, 25, 50, 500, 950]:
+        exp = f"Step_{timestep}"
+        for datasize in [197, 98, 49, 24, 12, 6, 3, 1]:
+            if os.path.isdir(head / exp / str(datasize)):
+                    if len(set(files_needed) - set(os.listdir(head / exp / str(datasize)))) == 0:
+                        print(f"Experiment {datasize} {timestep}")
+                        output = torch.load(head / exp / str(datasize)/  "JSRT_val_predictions.pt")
+                        print(f"{output['dice'].mean()}\t{output['dice'].std()}")
+                        for file in files_needed[1:]:
+                            output = torch.load(head / exp / str(datasize) / file)
+                            metrics_datasize = datasize if datasize is not None else 197
+                            metrics2["dice"].append(output["dice"].numpy())
+                            metrics2["precision"].append(output["precision"].numpy())
+                            metrics2["recall"].append(output["recall"].numpy())
+                            metrics2["exp"].append(np.array([exp] * len(output["dice"])))
+                            metrics2["datasize"].append(np.array([metrics_datasize] * len(output["dice"])))
+                            metrics2["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+                            metrics2["timestep"].append(np.array([timestep] * len(output["dice"])))
+                    else:
+                            print(f"Experiment {datasize} is missing files")
+
+
+    for key in metrics2:
+        metrics2[key] = np.concatenate(metrics2[key]).squeeze()
+        print(key, metrics2[key].shape)
+    df2 = pd.DataFrame(metrics2)
+
+    # %%  figure with line for baseline and datasetDM and boxplots for the rest
+    #  separating dice from precision and recall
+    font = 16
+    x = [1, 1, 3, 3, 6, 6, 12, 12, 24, 24, 49, 49, 197, 197]
+    plot_x = np.concatenate([np.array([-.4, .4]) + i for i in range(len(x)//2)]).flatten()
+    fig, axs = plt.subplots(3, 1, figsize=[12, 10])
+    sns.set_style("whitegrid")
+    m = 'dice'
+    for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+        ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'baseline') & (df.datasize == _x), m].to_numpy() for _x in x])
+        ys_std = np.quantile(ys, (.25, .75), axis=1, )
+        axs[i ].fill_between(plot_x, ys_std[0], ys_std[1], alpha=.2, zorder=0, color='C6')
+        ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'LEDM') & (df.datasize == _x), m].to_numpy() for _x in x])
+        ys_std = np.quantile(ys, (.25, .75), axis=1, )
+        axs[i ].fill_between(plot_x, ys_std[0], ys_std[1], alpha=.2, zorder=0, color='C8')
+        ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'baseline') & (df.datasize == _x), m].to_numpy() for _x in x])
+        ys_mean = np.quantile(ys, .5, axis=1)
+        axs[i ].plot(plot_x, ys_mean, label="baseline", c='C6', zorder=0)
+        ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'LEDM') & (df.datasize == _x), m].to_numpy() for _x in x])
+        ys_mean = np.quantile(ys, .5, axis=1)
+        axs[i ].plot(plot_x, ys_mean, label="LEDM" , c='C7', zorder=0)
+
+
+    for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+        temp_df = df2[(df2.dataset == dataset) & (df2.datasize != 98)]
+        out = sns.boxplot(data=temp_df, x="datasize", y=m, hue="timestep", ax=axs[i ],  showfliers=False, saturation=1,)
+        axs[i ].set_title(f"{dataset}", fontsize=font)
+        axs[i ].set_xlabel("" )
+        y_min, _ = axs[i ].get_ylim()
+        axs[i ].set_ylim(y_min, 1)
+        h, l = axs[i].get_legend_handles_labels()
+        axs[i].get_legend().remove()
+        axs[i].set_ylabel("Dice", fontsize=font)
+    sns.despine(ax=axs[0 ], offset=10, trim=True, bottom=True)
+    sns.despine(ax=axs[1 ], offset=10, trim=True, bottom=True)
+    sns.despine(ax=axs[2 ], offset=10, trim=True)
+    axs[0].set_xticks([])
+    axs[1].set_xticks([])
+    axs[-1 ].set_xlabel("Training dataset size", fontsize=font)
+    # Shrink current axis by 20%
+    for i, ax in enumerate(axs):
+        box = ax.get_position()
+        ax.tick_params(axis='both', labelsize=font)
+        ax.set_position([box.x0, box.y0, box.width , box.height])
+
+    # Put a legend to the right of the current axis
+
+    fig.legend(h, ['baseline', 'LEDM'] + ['step ' + _l for _l in l[2:]], title="", ncol=4, 
+            loc='center left', bbox_to_anchor=(0.2, -0.03), fontsize=font)
+    plt.tight_layout()
+    #plt.savefig("results_per_timestep.png")
+    plt.savefig("results_per_timestep_dice.pdf", bbox_inches='tight')
+    plt.show()
+    # %%
+    x = [1, 1, 3, 3, 6, 6, 12, 12, 24, 24, 49, 49, 197, 197]
+    plot_x = np.concatenate([np.array([-.4, .4]) + i for i in range(len(x)//2)]).flatten()
+    fig, axs = plt.subplots(3, 2, figsize=[15, 15])
+    sns.set_style("whitegrid")
+    for j, m in enumerate(["precision", "recall"]):
+        for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+            ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'baseline') & (df.datasize == _x), m].to_numpy() for _x in x])
+            ys_std = np.quantile(ys, (.25, .75), axis=1, )
+            axs[i, j].fill_between(plot_x, ys_std[0], ys_std[1], alpha=.2, zorder=0, color='C6')
+            ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'LEDM') & (df.datasize == _x), m].to_numpy() for _x in x])
+            ys_std = np.quantile(ys, (.25, .75), axis=1, )
+            axs[i, j].fill_between(plot_x, ys_std[0], ys_std[1], alpha=.2, zorder=0, color='C8')
+            ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'baseline') & (df.datasize == _x), m].to_numpy() for _x in x])
+            ys_mean = np.quantile(ys, .5, axis=1)
+            axs[i, j].plot(plot_x, ys_mean, label="baseline", c='C6', zorder=0)
+            ys = np.stack([df.loc[(df.dataset == dataset)& (df.exp == 'LEDM') & (df.datasize == _x), m].to_numpy() for _x in x])
+            ys_mean = np.quantile(ys, .5, axis=1)
+            axs[i, j].plot(plot_x, ys_mean, label="LEDM" , c='C7', zorder=0)
+
+
+            ##
+            temp_df = df2[(df2.dataset == dataset) & (df2.datasize != 98)]
+            out = sns.boxplot(data=temp_df, x="datasize", y=m, hue="timestep", ax=axs[i,j],  showfliers=False, saturation=1)
+            axs[i,j].set_title(f"{dataset}", fontsize=font)
+            y_min, _ = axs[i,j].get_ylim()
+            axs[i,j].set_ylim(y_min, 1)
+            sns.despine(ax=axs[i,j], offset=10, trim=True)
+            h, l = axs[i,j].get_legend_handles_labels()
+            axs[i,j].get_legend().remove()
+            axs[i, 0].set_ylabel("Precison", fontsize=font)
+            axs[i, 1].set_ylabel("Recall", fontsize=font)
+            axs[i,j].set_xlabel("")
+
+    for ax in axs.flatten():
+        ax.tick_params(axis='both', labelsize=font)
+    for ax in [axs[:, 0], axs[:, 1]]:
+        sns.despine(ax=ax[0 ], offset=10, trim=True, bottom=True)
+        sns.despine(ax=ax[1 ], offset=10, trim=True, bottom=True)
+        sns.despine(ax=ax[2 ], offset=10, trim=True)
+        ax[0].set_xticks([])
+        ax[1].set_xticks([])
+        ax[-1 ].set_xlabel("Training dataset size", fontsize=font)
+    # Put a legend to the right of the current axis
+
+
+    fig.legend(h, ['baseline', 'LEDM'] + ['step ' + _l for _l in l[2:]], title="", ncol=4, 
+            loc='center left', bbox_to_anchor=(0.25, -0.03), fontsize=font)
+    plt.tight_layout()
+    #plt.savefig("results_per_timestep.png")
+    plt.savefig("results_per_timestep_prec_recall.pdf", bbox_inches='tight')
+    plt.show()
+
+# %%

auxiliary/notebooks_and_reporting/print_tests_shared_weights.py

@@ -0,0 +1,222 @@
+# %%
+import numpy as np
+import torch
+from pathlib import Path
+import os
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+HEAD = Path(os.getcwd()).parent.parent
+
+if __name__=="__main__":
+    # load baseline and LEDM data
+    metrics = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[]}
+    files_needed = ["JSRT_val_predictions.pt", "JSRT_test_predictions.pt",  "NIH_predictions.pt", "Montgomery_predictions.pt",]
+    head = HEAD / 'logs'
+    for exp in ['baseline', 'LEDM']:
+         for datasize in [1, 3, 6, 12, 24, 49, 98, 197]:
+            if len(set(files_needed) - set(os.listdir(head / exp / str(datasize)))) == 0:
+                print(f"Experiment {exp} {datasize}")
+                output = torch.load(head / exp / str(datasize) / "JSRT_val_predictions.pt")
+                print(f"{output['dice'].mean()}\t{output['dice'].std()}")
+                for file in files_needed[1:]:
+                    output = torch.load(head / exp / str(datasize)  / file)
+                    metrics_datasize = 197 if datasize == "None" else int(datasize)
+                    metrics["dice"].append(output["dice"].numpy())
+                    metrics["precision"].append(output["precision"].numpy())
+                    metrics["recall"].append(output["recall"].numpy())
+                    metrics["exp"].append(np.array([exp] * len(output["dice"])))
+                    metrics["datasize"].append(np.array([int(datasize)] * len(output["dice"])))
+                    metrics["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+            else:
+                    print(f"Experiment {exp} is missing files")
+
+    for key in metrics:
+        metrics[key] = np.concatenate([el.squeeze() for el in metrics[key]])
+    df = pd.DataFrame(metrics)
+    df.head()
+
+    # %% load TEDM data
+    metrics3 = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[], }
+    exp = "TEDM"
+    for datasize in [1, 3, 6, 12, 24, 49, 98, 197]:
+                    if len(set(files_needed) - set(os.listdir(head / exp / str(datasize) ))) == 0:
+                        print(f"Experiment {datasize}")
+                        output = torch.load(head / exp / str(datasize)/ "JSRT_val_predictions.pt")
+                        print(f"{output['dice'].mean()}\t{output['dice'].std()}")
+                        for file in files_needed[1:]:
+                            output = torch.load(head / exp / str(datasize) /  file)
+                        
+                            metrics_datasize = datasize if datasize is not None else 197
+                            metrics3["dice"].append(output["dice"].numpy())
+                            metrics3["precision"].append(output["precision"].numpy())
+                            metrics3["recall"].append(output["recall"].numpy())
+                            metrics3["exp"].append(np.array(['TEDM'] * len(output["dice"])))
+                            metrics3["datasize"].append(np.array([metrics_datasize] * len(output["dice"])))
+                            metrics3["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+                            
+                    else:
+                            print(f"Experiment {datasize} is missing files")
+
+    for key in metrics3:
+        metrics3[key] = np.concatenate(metrics3[key]).squeeze()
+        print(key, metrics3[key].shape)
+    df3 = pd.DataFrame(metrics3)
+    # %% Boxplot of TEDM vs LEDM and baseline
+    df4 = pd.concat([df, df3])
+    df4.datasize = df4.datasize.astype(int)
+    m='dice'
+    dataset="JSRT"
+    fig, axs = plt.subplots(3, 3, figsize=(20, 20))
+    for j, m in enumerate(["dice", "precision", "recall"]):
+        #axs[0,j].set_ylim(0.8, 1)
+        #axs[0,j].set_ylim(0.6, 1)
+        #axs[0,j].set_ylim(0.7, 1)
+        for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+            temp_df = df4[(df4.dataset == dataset)]
+            #sns.lineplot(data=df[df.dataset == dataset], x="datasize", y=m, hue="exp", ax=axs[i,j])
+            sns.boxplot(data=temp_df, x="datasize", y=m,  ax=axs[i,j], hue="exp", showfliers=False, saturation=1,
+                        hue_order=['baseline', 'LEDM', 'TEDM'])
+            axs[i,j].set_title(f"{dataset} {m}")
+            axs[i,j].set_xlabel("Training dataset size")
+            h, l = axs[i,j].get_legend_handles_labels()
+            axs[i,j].legend(h, ['Baseline', 'LEDM', 'TEDM (ours)'], title="", loc='lower right')
+    plt.tight_layout()
+    plt.savefig("results_shared_weights.pdf")
+    plt.show()
+    # %% Load LEDMe and Step 1
+    metrics2 = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[], }
+    for exp in ["LEDMe", 'Step_1']:
+        for datasize in [1, 3, 6, 12, 24, 49, 98, 197]:
+            if len(set(files_needed) - set(os.listdir(head / exp / str(datasize) ))) == 0:
+                print(f"Experiment {exp} {datasize}")
+                output = torch.load(head / exp / str(datasize)/ "JSRT_val_predictions.pt")
+                print(f"{output['dice'].mean()}\t{output['dice'].std()}")
+                for file in files_needed[1:]:
+                    output = torch.load(head / exp / str(datasize) / file)
+                    #print(f"{output['dice'].mean()*100:.3}\t{output['dice'].std()*100:.3}\t{output['precision'].mean()*100:.3}\t{output['precision'].std()*100:.3}\t{output['recall'].mean()*100:.3}\t{output['recall'].std()*100:.3}",
+                    #    end="\n\n\n\n")
+                    metrics_datasize = 197 if datasize == "None" else datasize
+                    metrics2["dice"].append(output["dice"].numpy())
+                    metrics2["precision"].append(output["precision"].numpy())
+                    metrics2["recall"].append(output["recall"].numpy())
+                    metrics2["exp"].append(np.array([exp] * len(output["dice"])))
+                    metrics2["datasize"].append(np.array([int(metrics_datasize)] * len(output["dice"])))
+                    metrics2["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+            else:
+                    print(f"Experiment {exp} is missing files")
+
+    for key in metrics2:
+        metrics2[key] = np.concatenate(metrics2[key]).squeeze()
+        print(key, metrics2[key].shape)
+    df2 = pd.DataFrame(metrics2)
+    # %% Boxplot of TEDM vs LEDM and baseline, Step 1 and LEDMe
+    df4 = pd.concat([df, df3, df2])
+    df4.datasize = df4.datasize.astype(int)
+
+
+    m='dice'
+    dataset="JSRT"
+    fig, axs = plt.subplots(3, 3, figsize=(20, 20))
+    for j, m in enumerate(["dice", "precision", "recall"]):
+
+        for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+            temp_df = df4[(df4.dataset == dataset)]
+            #sns.lineplot(data=df[df.dataset == dataset], x="datasize", y=m, hue="exp", ax=axs[i,j])
+            sns.boxplot(data=temp_df, x="datasize", y=m,  ax=axs[i,j], hue="exp", showfliers=False, saturation=1,
+                        hue_order=['baseline', 'LEDM', 'Step_1', 'LEDMe', 'TEDM', ])
+            axs[i,j].set_title(f"{dataset} {m}")
+            axs[i,j].set_xlabel("Training dataset size")
+            h, l = axs[i,j].get_legend_handles_labels()
+            axs[i,j].legend(h, ['Baseline', 'LEDM',   'Step 1', 'LEDMe', 'TEDM'], title="", loc='lower right')
+    plt.tight_layout()
+    plt.savefig("results_shared_weights.pdf")
+    plt.show()
+    # %% Load TEDM ablation studies
+    metrics4 = {"dice": [], "precision": [], "recall": [], "exp": [], "datasize": [], "dataset":[], }
+    exp = "TEDM"
+    for datasize in [1, 3, 6, 12, 24, 49, 98, 197]:
+                if len(set(files_needed) - set(os.listdir(head / exp / str(datasize)))) == 0:
+                    print(f"Experiment {datasize} ")
+                    for step in [1,10,25]:
+                        for file in files_needed[1:]:
+                            output = torch.load(head / exp / str(datasize) /  file.replace("predictions", f"timestep{step}_predictions"))
+                            #print(f"{output['dice'].mean()*100:.3}\t{output['dice'].std()*100:.3}\t{output['precision'].mean()*100:.3}\t{output['precision'].std()*100:.3}\t{output['recall'].mean()*100:.3}\t{output['recall'].std()*100:.3}",
+                            #    end="\n\n\n\n")
+                            metrics_datasize = datasize if datasize is not None else 197
+                            metrics4["dice"].append(output["dice"].numpy())
+                            metrics4["precision"].append(output["precision"].numpy())
+                            metrics4["recall"].append(output["recall"].numpy())
+                            metrics4["exp"].append(np.array([f'Step {step} (MLP)'] * len(output["dice"])))
+                            metrics4["datasize"].append(np.array([metrics_datasize] * len(output["dice"])))
+                            metrics4["dataset"].append(np.array([file.split("_")[0]]*len(output["dice"])))
+                        #metrics3["timestep"].append(np.array(timestep * len(output["dice"])))
+                else:
+                        print(f"Experiment {datasize} is missing files")
+
+    for key in metrics3:
+        metrics4[key] = np.concatenate(metrics4[key]).squeeze()
+        print(key, metrics4[key].shape)
+    df4 = pd.DataFrame(metrics4)
+    # %% Print inputs to paper table
+    df_all = pd.concat([df, df3, df2, df4])
+    df_all.datasize = df_all.datasize.astype(int)
+    for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+        temp_df = df_all.loc[(df_all.dataset == dataset) & (df_all.datasize.isin([1, 3, 6, 12, 197])), ["exp", "datasize", "dice"]]
+        print(dataset)
+        mean = temp_df.groupby(["exp", "datasize"]).mean().unstack() * 100
+        std = temp_df.groupby(["exp", "datasize"]).std().unstack() * 100
+        for exp, exp_name in zip(['baseline', 'LEDM','Step_1', 'Step 1 (MLP)',
+                                'Step 10 (MLP)','Step 25 (MLP)', 'LEDMe', 'TEDM'],
+        ['Baseline', 'DatasetDDPM', 'Step 1 (linear)','Step 1 (MLP)', 'Step 10 (MLP)','Step 25 (MLP)','DatasetDDPMe', 'Ours', ]):
+        
+            print(exp_name, end='&\t')
+            print(f"{round(mean.loc[exp, ('dice', 1)],2):.3} $\pm$ {round(std.loc[exp, ('dice', 1)],1)}", end='&\t')
+            print(f"{round(mean.loc[exp, ('dice', 3)], 2):.3} $\pm$ {round(std.loc[exp, ('dice', 3)],1)}", end='&\t')
+            print(f"{round(mean.loc[exp, ('dice', 6)], 2):.3} $\pm$ {round(std.loc[exp, ('dice', 6)],1)}", end='&\t')
+            print(f"{round(mean.loc[exp, ('dice', 12)], 2):.3} $\pm$ {round(std.loc[exp, ('dice', 12)],1)}", end='&\t')
+            print(f"{round(mean.loc[exp, ('dice', 197)], 2):.3} $\pm$ {round(std.loc[exp, ('dice', 197)],1)}", end="""\\\\""")
+            
+            print()
+        
+    # %% Print inputs to paper appendix table
+    for i, dataset in enumerate(["JSRT", "NIH", "Montgomery"]):
+        print("\n" + dataset)
+        for m in ["precision", "recall"]:
+            temp_df = df_all.loc[(df_all.dataset == dataset) & (df_all.datasize.isin([1, 3, 6, 12, 24, 49, 98, 197])), ["exp", "datasize", m]]
+            print("\n"+m)
+            mean = temp_df.groupby(["exp", "datasize"]).mean().unstack() * 100
+            std = temp_df.groupby(["exp", "datasize"]).std().unstack() * 100
+            for exp, exp_name in zip(['baseline', 'LEDM','Step_1', 'LEDMe', 'TEDM'],
+            ['Baseline', 'LEDM', 'Step 1 (linear)','LEDMe', 'TEDM (ours)',]):
+            
+                print(exp_name, end='&\t')
+                print(f"{round(mean.loc[exp, (m, 1)],2):.3} $\pm$ {round(std.loc[exp, (m, 1)],1)}", end='&\t')
+                print(f"{round(mean.loc[exp, (m, 3)],2):.3} $\pm$ {round(std.loc[exp, (m, 3)],1)}", end='&\t')
+                print(f"{round(mean.loc[exp, (m, 6)],2):.3} $\pm$ {round(std.loc[exp, (m, 6)],1)}", end='&\t')
+                print(f"{round(mean.loc[exp, (m, 12)],2):.3} $\pm$ {round(std.loc[exp, (m, 12)],1)}", end='&\t')
+                print(f"{round(mean.loc[exp, (m, 197)],2):.3} $\pm$ {round(std.loc[exp, (m, 197)],1)}", end='\\\\')
+                
+                
+                print()  
+
+    # %% Wilcoxon tests - to use interactively
+    from scipy.stats import wilcoxon
+    m ="precision"
+    m='recall'
+    dataset ="Montgomery"
+    dssize =12
+
+    exp = "baseline"
+    exp = 'Step_1'
+    exp = "LEDM"
+    exp="TEDM"
+    exp_2= 'LEDMe'
+
+    x = df_all.loc[(df_all.dataset == dataset) & (df_all.exp == exp_2) & (df_all.datasize == dssize), m].to_numpy()
+    y = df_all.loc[(df_all.dataset == dataset) & (df_all.exp == exp)& (df_all.datasize == dssize), m].to_numpy()
+    print(f"{m} - {dataset} - {dssize} - {exp_2}: {x.mean():.4}+/-{x.std():.3} ")
+    print(f"{m} - {dataset} - {dssize} - {exp}: {y.mean():.4}+/-{y.std():.3} ")
+    print(f"{m} - {dataset} - {dssize}: {wilcoxon(x, y=y, zero_method='wilcox', correction=False, alternative='two-sided',).pvalue:.3} obs given equal   ")
+    print(f"{m} - {dataset} - {dssize}: {wilcoxon(x, y=y, zero_method='wilcox', correction=False, alternative='greater',).pvalue:.3} obs given {exp_2} < {exp} ")
+    print(f"{m} - {dataset} - {dssize}: {wilcoxon(x, y=y, zero_method='wilcox', correction=False, alternative='less',).pvalue:.3} obs given {exp_2} > {exp} ")

auxiliary/notebooks_and_reporting/visualisations.py

@@ -0,0 +1,162 @@
+# %%
+import numpy as np
+import torch
+from pathlib import Path
+import os, sys
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+HEAD = Path(os.getcwd()).parent.parent
+head = HEAD / 'logs'
+sys.path.append(HEAD)
+from dataloaders.JSRT import JSRTDataset
+from dataloaders.NIH import NIHDataset
+from dataloaders.Montgomery import MonDataset
+NIHPATH = "<PATH_TO_DATA>/NIH/"
+NIHFILE = "correspondence_with_chestXray8.csv"
+MONPATH = "<PATH_TO_DATA>/MontgomerySet/"
+MONFILE = "patient_data.csv"
+JSRTPATH = "<PATH_TO_DATA>/JSRT"
+
+if __name__=="__main__":
+    predictions = {'baseline':{'JSRT':{}, 'NIH':{}, 'Montgomery':{}},
+                'LEDM':{'JSRT':{}, 'NIH':{}, 'Montgomery':{}},
+                'TEDM':{'JSRT':{}, 'NIH':{}, 'Montgomery':{}},}
+    files_needed = ["JSRT_val_predictions.pt", "JSRT_test_predictions.pt",  "NIH_predictions.pt", "Montgomery_predictions.pt",]
+    for exp in ['baseline', 'LEDM', "TEDM"]:
+        for datasize in [1,3,6,12,24,49,98,197]:
+            if len(set(files_needed) - set(os.listdir(head / exp / str(datasize) ))) == 0:
+                for file in files_needed[1:]:
+                    output = torch.load(head / exp / str(datasize) / file)
+                    metrics_datasize = 197 if datasize == "None" else int(datasize)
+                    predictions[exp][file.rsplit("_")[0]][metrics_datasize]= output['y_hat']
+            else:
+                    print(f"Experiment {exp} is missing files")
+    # %%
+
+    img_size = 128
+    NIH_dataset = NIHDataset(NIHPATH, NIHPATH, NIHFILE, img_size)
+    JSRT_dataset = JSRTDataset(JSRTPATH, HEAD/ "data/", "JSRT_test_split.csv", img_size)
+    MON_dataset = MonDataset(MONPATH, MONPATH, MONFILE, img_size)
+    
+    # %%
+    loaders = {'JSRT': JSRT_dataset, 'NIH': NIH_dataset, 'Montgomery': MON_dataset}
+    m ="dice"
+    sz=4
+    ftsize= 40
+    fig, all_axs = plt.subplots(6, 21, figsize=(21*sz, 6*sz))
+    all_patients = [17, 13, 0, 1, 72, 78]
+
+    # JSRT
+    dataset ="JSRT"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[0]
+    print("JSRT1 - ", patient)
+    out = loaders[dataset][patient]
+    axs = all_axs[:3, :7]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+    axs[0, 0].set_title("JSRT - Image", fontsize=ftsize)
+    axs[0, 1].set_title("JSRT - GT", fontsize=ftsize)
+    axs[0, 2].set_title("1 (1%)"  , fontsize=ftsize)
+    axs[0, 3].set_title("3 (2%)", fontsize=ftsize)
+    axs[0, 4].set_title("6 (3%)", fontsize=ftsize)
+    axs[0, 5].set_title("12 (6%)", fontsize=ftsize)
+    axs[0, 6].set_title("197 (100%)", fontsize=ftsize)
+    axs[0,0].set_ylabel("Baseline", fontsize=ftsize)
+    axs[1,0].set_ylabel("LEDM", fontsize=ftsize)
+    axs[2,0].set_ylabel("TEDM", fontsize=ftsize)
+    #
+    axs = all_axs[3:, :7]
+    dataset ="JSRT"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[1]
+    print("JSRT2 - ", patient)
+    out = loaders[dataset][patient]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+    axs[0,0].set_ylabel("Baseline", fontsize=ftsize)
+    axs[1,0].set_ylabel("LEDM", fontsize=ftsize)
+    axs[2,0].set_ylabel("TEDM", fontsize=ftsize)
+    #
+    axs = all_axs[:3, 7:14]
+    dataset ="NIH"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[2]
+    print("NIH1 - ", patient)
+    out = loaders[dataset][patient]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+    axs[0, 0].set_title("NIH - Image", fontsize=ftsize)
+    axs[0, 1].set_title("NIH - GT", fontsize=ftsize)
+    axs[0, 2].set_title("1 (1%)"  , fontsize=ftsize)
+    axs[0, 3].set_title("3 (2%)", fontsize=ftsize)
+    axs[0, 4].set_title("6 (3%)", fontsize=ftsize)
+    axs[0, 5].set_title("12 (6%)", fontsize=ftsize)
+    axs[0, 6].set_title("197 (100%)", fontsize=ftsize)
+    #
+    #
+    axs = all_axs[3:, 7:14]
+    dataset ="NIH"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[3]
+    print("NIH2 - ", patient)
+    out = loaders[dataset][patient]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+    #
+    #
+    axs = all_axs[:3, 14:]
+    dataset ="Montgomery"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[4]
+    print("MON1 - ",patient)
+    out = loaders[dataset][patient]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+    axs[0, 0].set_title("Mont. - Image", fontsize=ftsize)
+    axs[0, 1].set_title("Mont. - GT", fontsize=ftsize)
+    axs[0, 2].set_title("1 (1%)", fontsize=ftsize)
+    axs[0, 3].set_title("3 (2%)", fontsize=ftsize)
+    axs[0, 4].set_title("6 (3%)", fontsize=ftsize)
+    axs[0, 5].set_title("12 (6%)", fontsize=ftsize)
+    axs[0, 6].set_title("197 (100%)", fontsize=ftsize)
+    #
+    axs = all_axs[3:, 14:]
+    dataset ="Montgomery"
+    patient = np.random.randint(0, len(loaders[dataset]))
+    patient = all_patients[5]
+    print("MON2 - ",patient)
+    out = loaders[dataset][patient]
+    for rowax, exp in zip(axs, ['baseline', 'LEDM', 'TEDM']):
+        rowax[0].imshow(out[0][0].numpy(), cmap='gray')
+        rowax[1].imshow(out[1][0].numpy(), interpolation='none', cmap='gray')
+        for ax, dssize in zip(rowax[2:], [1, 3, 6, 12, 197]):
+            ax.imshow(predictions[exp][dataset][dssize][patient].numpy()[0]>.5, interpolation='none')
+
+
+    # remove ticks
+    for ax in all_axs.flatten():
+        ax.set_xticks([])
+        ax.set_yticks([])
+        sns.despine(ax=ax, left=True, bottom=True)
+    plt.subplots_adjust(wspace=0.00, 
+                        hspace=0.00)
+    plt.tight_layout()
+    plt.savefig("visualisations2.pdf", bbox_inches='tight')
+    plt.show()

auxiliary/postprocessing/run_tests.py

@@ -0,0 +1,162 @@
+import argparse
+from pathlib import Path
+import os
+import torch
+from tqdm.auto import tqdm
+from torch import autocast
+from torch.utils.data import DataLoader
+import sys
+HEAD = Path(os.getcwd()).parent.parent
+sys.path.append("/vol/biomedic3/mmr12/projects/TEDM/")
+from models.diffusion_model import DiffusionModel
+from models.unet_model import Unet
+from models.datasetDM_model import DatasetDM
+from trainers.datasetDM_per_step import ModDatasetDM
+from trainers.train_baseline import dice, precision, recall
+from dataloaders.JSRT import build_dataloaders
+from dataloaders.NIH import NIHDataset
+from dataloaders.Montgomery import MonDataset
+
+
+NIHPATH = "/vol/biodata/data/chest_xray/NIH/"
+NIHFILE = "correspondence_with_chestXray8.csv"
+MONPATH = "/vol/biodata/data/chest_xray/NLM/MontgomerySet/"
+MONFILE = "patient_data.csv"
+
+
+if __name__ == "__main__":
+    # load config file and parse arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--experiment', "-e", type=str, help='Experiment path', default="logs/JSRT_conditional/20230213_171633")
+    parser.add_argument('--rerun', "-r", help='Run the test again', default=False, action="store_true")
+    args = parser.parse_args()
+
+    if os.path.isdir(args.experiment):
+        print("Experiment path identified as a directory")
+    else:
+        raise ValueError("Experiment path is not a directory")
+    files = os.listdir(args.experiment)
+    torch_file = None
+    if {'JSRT_val_predictions.pt', 'JSRT_test_predictions.pt', 'NIH_predictions.pt', 'Montgomery_predictions.pt'} <= set(files) and not args.rerun:
+        print("Experiment already tested")
+        for file in ['JSRT_val_predictions.pt', 'JSRT_test_predictions.pt', 'NIH_predictions.pt', 'Montgomery_predictions.pt']:
+            output = torch.load(Path(args.experiment) / file)
+            dataset_key = file.split("_")[0]
+            print(f"{dataset_key} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+            print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+            print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+            #torch.save(output, Path(args.experiment) / f'{dataset_key}_predictions.pt')
+        exit(0)
+
+    for f in files:
+        if "model" in f:
+            torch_file = f
+            break
+    if torch_file is None:
+        raise ValueError("No checkpoint file found in experiment directory")
+    
+    print(f"Loading experiment from {torch_file}")
+    data = torch.load(Path(args.experiment) / torch_file)
+    config = data["config"]
+
+    # pick model
+    if config.experiment in ["baseline", "global_finetune", "glob_loc_finetune"]:
+        model = Unet(**vars(config))
+    elif config.experiment == "datasetDM":
+        model = DatasetDM(config)
+    elif config.experiment == "simple_datasetDM":
+        model = ModDatasetDM(config)
+    else:
+        raise ValueError(f"Experiment {config.experiment} not recognized")
+    model.load_state_dict(data['model_state_dict'])
+
+    # Gather model output
+    model.eval().to(config.device)
+
+    # Load data
+    dataloaders = build_dataloaders(
+        config.data_dir,
+        config.img_size,
+        config.batch_size,
+        config.num_workers,
+    )
+    datasets_to_test = {
+        "JSRT_val": dataloaders["val"],
+        "JSRT_test": dataloaders["test"],
+        "NIH": DataLoader(NIHDataset(NIHPATH, NIHPATH, NIHFILE, config.img_size), 
+                          config.batch_size, num_workers=config.num_workers),
+        "Montgomery": DataLoader(MonDataset(MONPATH, MONPATH, MONFILE, config.img_size),
+                                 config.batch_size, num_workers=config.num_workers)
+
+    }
+    if config.experiment == "simple_datasetDM":
+        # re-calculate mean and var as they were not saved in the model dict
+        train_dl = dataloaders["train"]
+        for x, _ in tqdm(train_dl, desc="Calculating mean and variance"):
+            x = x.to(config.device)
+            features = model.extract_features(x)
+            model.mean += features.sum(dim=0)
+            model.mean_squared += (features ** 2).sum(dim=0)
+        model.mean = model.mean / len(train_dl.dataset)
+        model.std = (model.mean_squared / len(train_dl.dataset) - model.mean ** 2).sqrt() + 1e-6
+
+        model.mean = model.mean.to(config.device)
+        model.std = model.std.to(config.device)
+    
+    for dataset_key in datasets_to_test:
+        if f"{dataset_key}_predictions.pt" in files and not args.rerun:
+            print(f"{dataset_key} already tested")
+            output = torch.load(Path(args.experiment) / f'{dataset_key}_predictions.pt')
+            print(f"{dataset_key} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+            print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+            print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+            continue
+
+        print(f"Testing {dataset_key} set")
+        y_hat = []
+        y_star = []
+        for i, (x, y) in tqdm(enumerate(datasets_to_test[dataset_key]), desc='Validating'):
+            x = x.to(config.device)
+
+            if config.experiment == "conditional":
+                # sample n = 5 different segmetations
+                y_hats = []
+                for _ in range(5):
+                    img = torch.randn(x.shape, device=config.device)
+                    for t in tqdm(range(0, config.timesteps)[::-1]):
+                        # sample next timestep image (x_{t-1})
+                        with autocast(device_type=config.device, enabled=config.mixed_precision):
+                            with torch.no_grad():
+                                img = model.sample_timestep(img, t=t, cond=x)
+                    y_hats.append(img.detach().cpu() / 2 + .5)
+                # take the average over the 5 samples
+                y_hats = torch.stack(y_hats, -1).mean(-1)   
+
+                # record
+                y_hat.append(y_hats)
+                y_star.append(y)
+
+            elif config.experiment in ["baseline", "datasetDM",  "simple_datasetDM", "global_finetune", "glob_loc_finetune"] :
+                with autocast(device_type=config.device, enabled=config.mixed_precision):
+                    with torch.no_grad():
+                        pred = torch.sigmoid(model(x))
+                y_hat.append(pred.detach().cpu())
+                y_star.append(y)
+
+            else:
+                raise ValueError(f"Experiment {config.experiment} not recognized")
+        
+        # save predictions
+        y_hat = torch.cat(y_hat, 0)
+        y_star = torch.cat(y_star, 0)
+        output = {
+            'y_hat': y_hat, 
+            'y_star': y_star,
+            'dice':dice(y_hat>.5, y_star),
+            'precision':precision(y_hat>.5, y_star),
+            'recall':recall(y_hat>.5, y_star),}
+        
+        print(f"{dataset_key} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+        print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+        print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+        torch.save(output, Path(args.experiment) / f'{dataset_key}_predictions.pt')

auxiliary/postprocessing/testing_shared_weights.py

@@ -0,0 +1,145 @@
+import argparse
+from pathlib import Path
+import os
+import numpy as np
+import pandas as pd
+import torch
+import seaborn as sns
+import matplotlib.pyplot as plt
+from torch import nn
+from tqdm.auto import tqdm
+from torch import autocast
+from torch.utils.data import DataLoader
+from einops.layers.torch import Rearrange
+from einops import rearrange
+import sys
+HEAD = Path(os.getcwd()).parent.parent
+sys.path.append(HEAD)
+from models.datasetDM_model import DatasetDM
+from trainers.train_baseline import dice, precision, recall
+from dataloaders.JSRT import build_dataloaders
+from dataloaders.NIH import NIHDataset
+from dataloaders.Montgomery import MonDataset
+
+NIHPATH = "<PATH_TO_DATA>/NIH/"
+NIHFILE = "correspondence_with_chestXray8.csv" # saved in data
+MONPATH = "<PATH_TO_DATA>/NLM/MontgomerySet/"
+MONFILE = "patient_data.csv"
+
+
+if __name__ == "__main__":
+    # load config file and parse arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--experiment', "-e", type=str, help='Experiment path', default="logs/JSRT_conditional/20230213_171633")
+    parser.add_argument('--rerun', "-r", help='Run the test again', default=False, action="store_true")
+    args = parser.parse_args()
+
+    if os.path.isdir(args.experiment):
+        print("Experiment path identified as a directory")
+    else:
+        raise ValueError("Experiment path is not a directory")
+    files = os.listdir(args.experiment)
+    torch_file = None
+    if {'JSRT_val_predictions.pt', 'JSRT_test_predictions.pt', 'NIH_predictions.pt', 'Montgomery_predictions.pt'} <= set(files) and not args.rerun:
+        print("Experiment already tested")
+        sys.exit(0)
+
+    for f in files:
+        if "model" in f:
+            torch_file = f
+            break
+    if torch_file is None:
+        raise ValueError("No checkpoint file found in experiment directory")
+    
+    print(f"Loading experiment from {torch_file}")
+    data = torch.load(Path(args.experiment) / torch_file)
+    config = data["config"]
+
+    # pick model
+    if config.experiment == "datasetDM":
+        model = DatasetDM(config)
+        model.classifier = nn.Sequential(
+            Rearrange('b (step act) h w -> (b step) act h w', step=len(model.steps)),
+            nn.Conv2d(960, 128, 1),
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 32, 1),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.Conv2d(32, 1, config.out_channels)
+            )
+    else:
+        raise ValueError(f"Experiment {config.experiment} not recognized")
+    model.load_state_dict(data['model_state_dict'])
+
+    # Gather model output
+    model.eval().to(config.device)
+
+    # Load data
+    dataloaders = build_dataloaders(
+        config.data_dir,
+        config.img_size,
+        config.batch_size,
+        config.num_workers,
+    )
+    datasets_to_test = {
+        "JSRT_val": dataloaders["val"],
+        "JSRT_test": dataloaders["test"],
+        "NIH": DataLoader(NIHDataset(NIHPATH, NIHPATH, NIHFILE, config.img_size), 
+                          config.batch_size, num_workers=config.num_workers),
+        "Montgomery": DataLoader(MonDataset(MONPATH, MONPATH, MONFILE, config.img_size),
+                                 config.batch_size, num_workers=config.num_workers)
+
+    }
+
+    for dataset_key in datasets_to_test:
+        if f"{dataset_key}_predictions.pt" in files and not args.rerun:
+            print(f"{dataset_key} already tested")
+            output = torch.load(Path(args.experiment) / f'{dataset_key}_predictions.pt')
+            print(f"{dataset_key} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+            print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+            print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+            continue
+
+        print(f"Testing {dataset_key} set")
+        y_hats = []
+        y_star = []
+        for i, (x, y) in tqdm(enumerate(datasets_to_test[dataset_key]), desc='Validating'):
+            x = x.to(config.device)
+
+            with autocast(device_type=config.device, enabled=config.mixed_precision):
+                with torch.no_grad():
+                    # all depths
+                    pred = torch.sigmoid(model(x))
+            y_hats.append(pred.detach().cpu())
+            y_star.append(y)
+        
+        # save predictions
+        y_star = torch.cat(y_star, 0)
+        y_hats = torch.cat(y_hats, 0)
+        y_hats = rearrange(y_hats, '(b step) 1 h w -> step b 1 h w', step=len(model.steps))
+        for i, y_hat in enumerate(y_hats):
+            output = {
+                'y_hat': y_hat, 
+                'y_star': y_star,
+                'dice':dice(y_hat>.5, y_star),
+                'precision':precision(y_hat>.5, y_star),
+                'recall':recall(y_hat>.5, y_star),}
+            
+            print(f"{dataset_key} {model.steps[i]} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+            print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+            print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+            torch.save(output, Path(args.experiment) / f'{dataset_key}_timestep{model.steps[i]}_predictions.pt')
+        y_hat = y_hats.mean(0)
+        output = {
+                'y_hat': y_hat, 
+                'y_star': y_star,
+                'dice':dice(y_hat>.5, y_star),
+                'precision':precision(y_hat>.5, y_star),
+                'recall':recall(y_hat>.5, y_star),}
+            
+        print(f"{dataset_key} metrics: \n\tdice:      {output['dice'].mean():.3}+/-{output['dice'].std():.3}")
+        print(f"\tprecision: {output['precision'].mean():.3}+/-{output['precision'].std():.3}")
+        print(f"\trecall:    {output['recall'].mean():.3}+/-{output['recall'].std():.3}")
+        torch.save(output, Path(args.experiment) / f'{dataset_key}_predictions.pt')
+    

auxiliary/preprocessing/CXR14_preprocessing_separate_data.py

@@ -0,0 +1,31 @@
+# %%
+import pandas as pd
+from pathlib import Path
+import numpy as np
+import os
+
+CWDIR = Path(os.getcwd()).parent.parent
+DATADIR = Path("<PATH_TO_DATA>/ChestXray-NIHCC")
+if not os.path.isdir(DATADIR):
+    print(f"Data directory {DATADIR} not found")
+
+
+df = pd.concat([pd.read_csv(DATADIR / "train_val_list.csv"),pd.read_csv(DATADIR / "test_list.csv")])
+df.reset_index(inplace=True)
+# %%
+from tqdm import tqdm
+items = []
+for el in tqdm(df["Image Index"]):
+    items.append(os.path.isfile(DATADIR / "images"/ el))
+# %% Shuffle and remove 20% for test and val
+idx = np.arange(len(df))
+np.random.shuffle(idx)
+n1 = int(len(df)*.8)
+n2 = int(len(df)*.9)
+idxs = [idx[:n1], idx[n1:n2], idx[n2:]]
+for i in range(3):
+    print(len(df.loc[idxs[i]]))
+# %%
+df.loc[idxs[0]].to_csv(CWDIR / 'data' / 'train_split.csv', index=False)
+df.loc[idxs[1]].to_csv(CWDIR / 'data' / 'val_split.csv', index=False)
+df.loc[idxs[2]].to_csv(CWDIR / 'data' / 'test_split.csv', index=False)

auxiliary/preprocessing/JSRT_preprocessing_separate_data.py

@@ -0,0 +1,26 @@
+# %%
+import pandas as pd
+from pathlib import Path
+import numpy as np
+import os
+
+CWDIR = Path(os.getcwd()).parent.parent
+
+
+head = Path("<PATH_TO_DATA>/JSRT")
+
+df = pd.read_csv(head / "jsrt_metadata_with_masks.csv")
+df.reset_index(inplace=True)
+
+# %% Shuffle and remove 20% for test and val
+idx = np.arange(len(df))
+np.random.shuffle(idx)
+n1 = int(len(df)*.8)
+n2 = int(len(df)*.9)
+idxs = [idx[:n1], idx[n1:n2], idx[n2:]]
+for i in range(3):
+    print(len(df.loc[idxs[i]]))
+# %%
+df.loc[idxs[0]].to_csv(CWDIR / 'data' / 'JSRT_train_split.csv', index=False)
+df.loc[idxs[1]].to_csv(CWDIR / 'data' / 'JSRT_val_split.csv', index=False)
+df.loc[idxs[2]].to_csv(CWDIR / 'data' / 'JSRT_test_split.csv', index=False)

config.py

@@ -0,0 +1,84 @@
+import argparse
+import os
+
+from datetime import datetime
+from pathlib import Path
+import torch
+
+
+this_dir = os.path.dirname(os.path.realpath(__file__))
+default_logdir = os.path.join(this_dir, 'logs', datetime.now().strftime('%Y%m%d_%H%M%S'))
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--mixed_precision', type=bool, default=False, help='Use mixed precision')
+parser.add_argument('--resume_path', type=str, default=None, help='Path to checkpoint to resume from')
+
+# Experiment parameters
+parser.add_argument('--experiment', type=str, default="img_only",choices=[
+    "PDDM",
+    "baseline", 
+    "LEDM", 
+    "LEDMe", 
+    "TEDM",
+    "global_cl",
+    "local_cl",
+    "global_finetune",
+    "glob_loc_finetune"
+    ], help='Whether to generate only images or images and segmentations')
+parser.add_argument('--dataset', type=str, default="JSRT",choices=["JSRT", "CXR14"], help='Dataset to use')
+
+# Data parameters
+parser.add_argument('--img_size', type=int, default=128, help='Height / width of the input image to the network')
+parser.add_argument('--data_dir', type=str, help='Path to the dataset')
+parser.add_argument('--num_workers', type=int, default=4, help='Number of subprocesses to use for data loading')
+
+# Model parameters
+parser.add_argument('--dim', type=int, default=64, help='Width of the U-Net')
+parser.add_argument('--dim_mults', nargs='+', type=int, default=(1, 2, 4, 8), help='Dimension multipliers for U-Net levels')
+# SegDiff model parameters
+parser.add_argument('--seg_out_dim', type=int, default=1, help='Dimension of segmentation embedding')
+parser.add_argument('--img_out_dim', type=int, default=4, help='Dimension of image embedding')
+parser.add_argument('--img_inter_dim', type=int, default=32, help='Width of image embedding')
+
+# Diffusion parameters
+parser.add_argument('--timesteps', type=int, default=1000, help='Number of diffusion timesteps')
+parser.add_argument('--beta_schedule', type=str, default='cosine', choices=['linear', 'cosine'])
+parser.add_argument('--objective', type=str, default='pred_noise', help='Model output', choices=['pred_noise', 'pred_x_0'])
+
+# CL parameters
+parser.add_argument('--tau', type=float, default=0.1, help='Temperature parameter for contrastive loss')
+parser.add_argument('--global_model_path', type=str, default=None, help='Path to global model checkpoint')
+parser.add_argument('--glob_loc_model_path', type=str, default=None, help='Path to global & local CL model checkpoint')
+parser.add_argument('--unfreeze_weights_at_step', type=int, default=0, help='Step at which to unfreeze pretrained weights. If 0, weights are not frozen')
+parser.add_argument('--augment_at_finetuning', default=False, action='store_true', help='Whether to augment images during finetuning')
+
+# Training parameters
+parser.add_argument('--batch_size', type=int, default=16, help='Input batch size')
+parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
+parser.add_argument('--weight_decay', type=float, default=0, help='Weight decay')
+# parser.add_argument('--adam_betas', nargs=2, type=float, default=(0.9, 0.99), help='Betas for the Adam optimizer')
+parser.add_argument('--max_steps', type=int, default=500000, help='Number of training steps to perform')
+parser.add_argument('--p2_loss_weight_gamma', type=float, default=0., help='p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended')
+parser.add_argument('--p2_loss_weight_k', type=float, default=1.)
+parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device to use')
+parser.add_argument('--seed', type=int, default=0, help='Random seed')
+
+# Logging parameters
+parser.add_argument('--log_freq', type=int, default=100, help='Frequency of logging')
+parser.add_argument('--val_freq', type=int, default=100, help='Frequency of validation')
+parser.add_argument('--val_steps', type=int, default=250, help='Number of timestep to use for validation')
+parser.add_argument('--log_dir', type=str, default=default_logdir, help='Logging directory')
+parser.add_argument('--n_sampled_imgs', type=int, default=8, help='Number of images to sample during logging')
+parser.add_argument('--max_val_steps', type=int, default=-1, help='Number of validation steps to perform')
+
+# datasetGAN like segmentation model parameters
+parser.add_argument("--saved_diffusion_model", type=str, help='Path to checkpoint of trained diffusion model', default="logs/20230127_164150/best_model.pt")
+parser.add_argument("--t_steps_to_save", type=int, nargs='*', choices=range(1000), help='Diffusion steps to be used as features', default=[50, 200, 400, 600, 800])
+parser.add_argument("--n_labelled_images", type=int, help='Number of labelled images to use for semi-supervised training', default=None, 
+                    choices=[197, 98, 49, 24, 12, 6, 3, 1])
+
+# other experiments I played with
+parser.add_argument("--shared_weights_over_timesteps", help='In datasetDM, only use last timestep to predict, and intermediate timesteps to train', default=False, action='store_true')
+parser.add_argument("--early_stop", help='In baseline, if validation loss increases by more than 50%, stop', default=False, action='store_true')

data/JSRT_test_split.csv

@@ -0,0 +1,26 @@
+id,path
+JPCNN017,JSRT/PNG_data/JPCNN017.png
+JPCLN151,JSRT/PNG_data/JPCLN151.png
+JPCNN007,JSRT/PNG_data/JPCNN007.png
+JPCNN089,JSRT/PNG_data/JPCNN089.png
+JPCLN153,JSRT/PNG_data/JPCLN153.png
+JPCNN020,JSRT/PNG_data/JPCNN020.png
+JPCNN093,JSRT/PNG_data/JPCNN093.png
+JPCLN118,JSRT/PNG_data/JPCLN118.png
+JPCLN143,JSRT/PNG_data/JPCLN143.png
+JPCLN073,JSRT/PNG_data/JPCLN073.png
+JPCLN018,JSRT/PNG_data/JPCLN018.png
+JPCLN109,JSRT/PNG_data/JPCLN109.png
+JPCLN095,JSRT/PNG_data/JPCLN095.png
+JPCNN055,JSRT/PNG_data/JPCNN055.png
+JPCLN131,JSRT/PNG_data/JPCLN131.png
+JPCLN130,JSRT/PNG_data/JPCLN130.png
+JPCLN053,JSRT/PNG_data/JPCLN053.png
+JPCLN107,JSRT/PNG_data/JPCLN107.png
+JPCNN081,JSRT/PNG_data/JPCNN081.png
+JPCLN146,JSRT/PNG_data/JPCLN146.png
+JPCLN058,JSRT/PNG_data/JPCLN058.png
+JPCLN010,JSRT/PNG_data/JPCLN010.png
+JPCLN137,JSRT/PNG_data/JPCLN137.png
+JPCLN086,JSRT/PNG_data/JPCLN086.png
+JPCLN114,JSRT/PNG_data/JPCLN114.png

data/JSRT_train_split.csv

@@ -0,0 +1,198 @@
+id,path
+JPCLN001,JSRT/PNG_data/JPCLN001.png
+JPCLN002,JSRT/PNG_data/JPCLN002.png
+JPCLN003,JSRT/PNG_data/JPCLN003.png
+JPCLN004,JSRT/PNG_data/JPCLN004.png
+JPCLN005,JSRT/PNG_data/JPCLN005.png
+JPCLN006,JSRT/PNG_data/JPCLN006.png
+JPCLN007,JSRT/PNG_data/JPCLN007.png
+JPCLN008,JSRT/PNG_data/JPCLN008.png
+JPCLN009,JSRT/PNG_data/JPCLN009.png
+JPCLN011,JSRT/PNG_data/JPCLN011.png
+JPCLN012,JSRT/PNG_data/JPCLN012.png
+JPCLN013,JSRT/PNG_data/JPCLN013.png
+JPCLN014,JSRT/PNG_data/JPCLN014.png
+JPCLN015,JSRT/PNG_data/JPCLN015.png
+JPCLN016,JSRT/PNG_data/JPCLN016.png
+JPCLN017,JSRT/PNG_data/JPCLN017.png
+JPCLN019,JSRT/PNG_data/JPCLN019.png
+JPCLN020,JSRT/PNG_data/JPCLN020.png
+JPCLN021,JSRT/PNG_data/JPCLN021.png
+JPCLN022,JSRT/PNG_data/JPCLN022.png
+JPCLN023,JSRT/PNG_data/JPCLN023.png
+JPCLN024,JSRT/PNG_data/JPCLN024.png
+JPCLN025,JSRT/PNG_data/JPCLN025.png
+JPCLN026,JSRT/PNG_data/JPCLN026.png
+JPCLN027,JSRT/PNG_data/JPCLN027.png
+JPCLN029,JSRT/PNG_data/JPCLN029.png
+JPCLN031,JSRT/PNG_data/JPCLN031.png
+JPCLN032,JSRT/PNG_data/JPCLN032.png
+JPCLN033,JSRT/PNG_data/JPCLN033.png
+JPCLN034,JSRT/PNG_data/JPCLN034.png
+JPCLN035,JSRT/PNG_data/JPCLN035.png
+JPCLN036,JSRT/PNG_data/JPCLN036.png
+JPCLN038,JSRT/PNG_data/JPCLN038.png
+JPCLN039,JSRT/PNG_data/JPCLN039.png
+JPCLN040,JSRT/PNG_data/JPCLN040.png
+JPCLN041,JSRT/PNG_data/JPCLN041.png
+JPCLN042,JSRT/PNG_data/JPCLN042.png
+JPCLN043,JSRT/PNG_data/JPCLN043.png
+JPCLN044,JSRT/PNG_data/JPCLN044.png
+JPCLN046,JSRT/PNG_data/JPCLN046.png
+JPCLN047,JSRT/PNG_data/JPCLN047.png
+JPCLN048,JSRT/PNG_data/JPCLN048.png
+JPCLN049,JSRT/PNG_data/JPCLN049.png
+JPCLN050,JSRT/PNG_data/JPCLN050.png
+JPCLN051,JSRT/PNG_data/JPCLN051.png
+JPCLN052,JSRT/PNG_data/JPCLN052.png
+JPCLN054,JSRT/PNG_data/JPCLN054.png
+JPCLN056,JSRT/PNG_data/JPCLN056.png
+JPCLN057,JSRT/PNG_data/JPCLN057.png
+JPCLN059,JSRT/PNG_data/JPCLN059.png
+JPCLN060,JSRT/PNG_data/JPCLN060.png
+JPCLN061,JSRT/PNG_data/JPCLN061.png
+JPCLN062,JSRT/PNG_data/JPCLN062.png
+JPCLN063,JSRT/PNG_data/JPCLN063.png
+JPCLN065,JSRT/PNG_data/JPCLN065.png
+JPCLN066,JSRT/PNG_data/JPCLN066.png
+JPCLN067,JSRT/PNG_data/JPCLN067.png
+JPCLN068,JSRT/PNG_data/JPCLN068.png
+JPCLN069,JSRT/PNG_data/JPCLN069.png
+JPCLN070,JSRT/PNG_data/JPCLN070.png
+JPCLN072,JSRT/PNG_data/JPCLN072.png
+JPCLN074,JSRT/PNG_data/JPCLN074.png
+JPCLN075,JSRT/PNG_data/JPCLN075.png
+JPCLN076,JSRT/PNG_data/JPCLN076.png
+JPCLN077,JSRT/PNG_data/JPCLN077.png
+JPCLN078,JSRT/PNG_data/JPCLN078.png
+JPCLN079,JSRT/PNG_data/JPCLN079.png
+JPCLN080,JSRT/PNG_data/JPCLN080.png
+JPCLN081,JSRT/PNG_data/JPCLN081.png
+JPCLN082,JSRT/PNG_data/JPCLN082.png
+JPCLN083,JSRT/PNG_data/JPCLN083.png
+JPCLN084,JSRT/PNG_data/JPCLN084.png
+JPCLN085,JSRT/PNG_data/JPCLN085.png
+JPCLN088,JSRT/PNG_data/JPCLN088.png
+JPCLN089,JSRT/PNG_data/JPCLN089.png
+JPCLN090,JSRT/PNG_data/JPCLN090.png
+JPCLN091,JSRT/PNG_data/JPCLN091.png
+JPCLN092,JSRT/PNG_data/JPCLN092.png
+JPCLN093,JSRT/PNG_data/JPCLN093.png
+JPCLN097,JSRT/PNG_data/JPCLN097.png
+JPCLN098,JSRT/PNG_data/JPCLN098.png
+JPCLN100,JSRT/PNG_data/JPCLN100.png
+JPCLN101,JSRT/PNG_data/JPCLN101.png
+JPCLN102,JSRT/PNG_data/JPCLN102.png
+JPCLN103,JSRT/PNG_data/JPCLN103.png
+JPCLN104,JSRT/PNG_data/JPCLN104.png
+JPCLN105,JSRT/PNG_data/JPCLN105.png
+JPCLN108,JSRT/PNG_data/JPCLN108.png
+JPCLN110,JSRT/PNG_data/JPCLN110.png
+JPCLN111,JSRT/PNG_data/JPCLN111.png
+JPCLN112,JSRT/PNG_data/JPCLN112.png
+JPCLN113,JSRT/PNG_data/JPCLN113.png
+JPCLN115,JSRT/PNG_data/JPCLN115.png
+JPCLN116,JSRT/PNG_data/JPCLN116.png
+JPCLN117,JSRT/PNG_data/JPCLN117.png
+JPCLN120,JSRT/PNG_data/JPCLN120.png
+JPCLN121,JSRT/PNG_data/JPCLN121.png
+JPCLN122,JSRT/PNG_data/JPCLN122.png
+JPCLN123,JSRT/PNG_data/JPCLN123.png
+JPCLN124,JSRT/PNG_data/JPCLN124.png
+JPCLN125,JSRT/PNG_data/JPCLN125.png
+JPCLN126,JSRT/PNG_data/JPCLN126.png
+JPCLN127,JSRT/PNG_data/JPCLN127.png
+JPCLN128,JSRT/PNG_data/JPCLN128.png
+JPCLN129,JSRT/PNG_data/JPCLN129.png
+JPCLN132,JSRT/PNG_data/JPCLN132.png
+JPCLN133,JSRT/PNG_data/JPCLN133.png
+JPCLN134,JSRT/PNG_data/JPCLN134.png
+JPCLN135,JSRT/PNG_data/JPCLN135.png
+JPCLN136,JSRT/PNG_data/JPCLN136.png
+JPCLN138,JSRT/PNG_data/JPCLN138.png
+JPCLN139,JSRT/PNG_data/JPCLN139.png
+JPCLN140,JSRT/PNG_data/JPCLN140.png
+JPCLN141,JSRT/PNG_data/JPCLN141.png
+JPCLN142,JSRT/PNG_data/JPCLN142.png
+JPCLN144,JSRT/PNG_data/JPCLN144.png
+JPCLN145,JSRT/PNG_data/JPCLN145.png
+JPCLN147,JSRT/PNG_data/JPCLN147.png
+JPCLN148,JSRT/PNG_data/JPCLN148.png
+JPCLN149,JSRT/PNG_data/JPCLN149.png
+JPCLN150,JSRT/PNG_data/JPCLN150.png
+JPCLN152,JSRT/PNG_data/JPCLN152.png
+JPCLN154,JSRT/PNG_data/JPCLN154.png
+JPCNN001,JSRT/PNG_data/JPCNN001.png
+JPCNN002,JSRT/PNG_data/JPCNN002.png
+JPCNN004,JSRT/PNG_data/JPCNN004.png
+JPCNN006,JSRT/PNG_data/JPCNN006.png
+JPCNN008,JSRT/PNG_data/JPCNN008.png
+JPCNN009,JSRT/PNG_data/JPCNN009.png
+JPCNN010,JSRT/PNG_data/JPCNN010.png
+JPCNN011,JSRT/PNG_data/JPCNN011.png
+JPCNN014,JSRT/PNG_data/JPCNN014.png
+JPCNN015,JSRT/PNG_data/JPCNN015.png
+JPCNN016,JSRT/PNG_data/JPCNN016.png
+JPCNN018,JSRT/PNG_data/JPCNN018.png
+JPCNN019,JSRT/PNG_data/JPCNN019.png
+JPCNN021,JSRT/PNG_data/JPCNN021.png
+JPCNN022,JSRT/PNG_data/JPCNN022.png
+JPCNN023,JSRT/PNG_data/JPCNN023.png
+JPCNN024,JSRT/PNG_data/JPCNN024.png
+JPCNN025,JSRT/PNG_data/JPCNN025.png
+JPCNN026,JSRT/PNG_data/JPCNN026.png
+JPCNN028,JSRT/PNG_data/JPCNN028.png
+JPCNN029,JSRT/PNG_data/JPCNN029.png
+JPCNN030,JSRT/PNG_data/JPCNN030.png
+JPCNN031,JSRT/PNG_data/JPCNN031.png
+JPCNN032,JSRT/PNG_data/JPCNN032.png
+JPCNN033,JSRT/PNG_data/JPCNN033.png
+JPCNN034,JSRT/PNG_data/JPCNN034.png
+JPCNN035,JSRT/PNG_data/JPCNN035.png
+JPCNN037,JSRT/PNG_data/JPCNN037.png
+JPCNN038,JSRT/PNG_data/JPCNN038.png
+JPCNN039,JSRT/PNG_data/JPCNN039.png
+JPCNN040,JSRT/PNG_data/JPCNN040.png
+JPCNN041,JSRT/PNG_data/JPCNN041.png
+JPCNN042,JSRT/PNG_data/JPCNN042.png
+JPCNN043,JSRT/PNG_data/JPCNN043.png
+JPCNN044,JSRT/PNG_data/JPCNN044.png
+JPCNN045,JSRT/PNG_data/JPCNN045.png
+JPCNN046,JSRT/PNG_data/JPCNN046.png
+JPCNN047,JSRT/PNG_data/JPCNN047.png
+JPCNN048,JSRT/PNG_data/JPCNN048.png
+JPCNN049,JSRT/PNG_data/JPCNN049.png
+JPCNN050,JSRT/PNG_data/JPCNN050.png
+JPCNN051,JSRT/PNG_data/JPCNN051.png
+JPCNN052,JSRT/PNG_data/JPCNN052.png
+JPCNN053,JSRT/PNG_data/JPCNN053.png
+JPCNN056,JSRT/PNG_data/JPCNN056.png
+JPCNN059,JSRT/PNG_data/JPCNN059.png
+JPCNN060,JSRT/PNG_data/JPCNN060.png
+JPCNN062,JSRT/PNG_data/JPCNN062.png
+JPCNN063,JSRT/PNG_data/JPCNN063.png
+JPCNN065,JSRT/PNG_data/JPCNN065.png
+JPCNN067,JSRT/PNG_data/JPCNN067.png
+JPCNN068,JSRT/PNG_data/JPCNN068.png
+JPCNN069,JSRT/PNG_data/JPCNN069.png
+JPCNN070,JSRT/PNG_data/JPCNN070.png
+JPCNN071,JSRT/PNG_data/JPCNN071.png
+JPCNN072,JSRT/PNG_data/JPCNN072.png
+JPCNN073,JSRT/PNG_data/JPCNN073.png
+JPCNN074,JSRT/PNG_data/JPCNN074.png
+JPCNN075,JSRT/PNG_data/JPCNN075.png
+JPCNN076,JSRT/PNG_data/JPCNN076.png
+JPCNN077,JSRT/PNG_data/JPCNN077.png
+JPCNN078,JSRT/PNG_data/JPCNN078.png
+JPCNN079,JSRT/PNG_data/JPCNN079.png
+JPCNN080,JSRT/PNG_data/JPCNN080.png
+JPCNN082,JSRT/PNG_data/JPCNN082.png
+JPCNN083,JSRT/PNG_data/JPCNN083.png
+JPCNN084,JSRT/PNG_data/JPCNN084.png
+JPCNN085,JSRT/PNG_data/JPCNN085.png
+JPCNN086,JSRT/PNG_data/JPCNN086.png
+JPCNN087,JSRT/PNG_data/JPCNN087.png
+JPCNN088,JSRT/PNG_data/JPCNN088.png
+JPCNN090,JSRT/PNG_data/JPCNN090.png
+JPCNN091,JSRT/PNG_data/JPCNN091.png
+JPCNN092,JSRT/PNG_data/JPCNN092.png

data/JSRT_val_split.csv

@@ -0,0 +1,26 @@
+id,path
+JPCLN028,JSRT/PNG_data/JPCLN028.png
+JPCNN005,JSRT/PNG_data/JPCNN005.png
+JPCNN013,JSRT/PNG_data/JPCNN013.png
+JPCLN064,JSRT/PNG_data/JPCLN064.png
+JPCLN055,JSRT/PNG_data/JPCLN055.png
+JPCNN054,JSRT/PNG_data/JPCNN054.png
+JPCLN096,JSRT/PNG_data/JPCLN096.png
+JPCLN099,JSRT/PNG_data/JPCLN099.png
+JPCNN064,JSRT/PNG_data/JPCNN064.png
+JPCLN030,JSRT/PNG_data/JPCLN030.png
+JPCNN057,JSRT/PNG_data/JPCNN057.png
+JPCLN094,JSRT/PNG_data/JPCLN094.png
+JPCLN087,JSRT/PNG_data/JPCLN087.png
+JPCNN012,JSRT/PNG_data/JPCNN012.png
+JPCNN061,JSRT/PNG_data/JPCNN061.png
+JPCLN071,JSRT/PNG_data/JPCLN071.png
+JPCLN119,JSRT/PNG_data/JPCLN119.png
+JPCNN027,JSRT/PNG_data/JPCNN027.png
+JPCLN037,JSRT/PNG_data/JPCLN037.png
+JPCLN045,JSRT/PNG_data/JPCLN045.png
+JPCNN066,JSRT/PNG_data/JPCNN066.png
+JPCNN003,JSRT/PNG_data/JPCNN003.png
+JPCNN058,JSRT/PNG_data/JPCNN058.png
+JPCNN036,JSRT/PNG_data/JPCNN036.png
+JPCLN106,JSRT/PNG_data/JPCLN106.png

data/correspondence_with_chestXray8.csv

@@ -0,0 +1,101 @@
+NIH,ChestX-ray14,scan,mask
+NIH_0010,00009863_008.png,images/NIH_0010.png,masks/NIH_0010_mask.png
+NIH_0017,00003028_078.png,images/NIH_0017.png,masks/NIH_0017_mask.png
+NIH_0065,00012045_005.png,images/NIH_0065.png,masks/NIH_0065_mask.png
+NIH_0019,00029481_008.png,images/NIH_0019.png,masks/NIH_0019_mask.png
+NIH_0062,00010805_040.png,images/NIH_0062.png,masks/NIH_0062_mask.png
+NIH_0057,00011950_019.png,images/NIH_0057.png,masks/NIH_0057_mask.png
+NIH_0086,00004648_001.png,images/NIH_0086.png,masks/NIH_0086_mask.png
+NIH_0050,00019499_006.png,images/NIH_0050.png,masks/NIH_0050_mask.png
+NIH_0081,00022572_048.png,images/NIH_0081.png,masks/NIH_0081_mask.png
+NIH_0022,00014398_023.png,images/NIH_0022.png,masks/NIH_0022_mask.png
+NIH_0025,00017753_003.png,images/NIH_0025.png,masks/NIH_0025_mask.png
+NIH_0059,00020289_004.png,images/NIH_0059.png,masks/NIH_0059_mask.png
+NIH_0088,00004832_025.png,images/NIH_0088.png,masks/NIH_0088_mask.png
+NIH_0034,00009863_038.png,images/NIH_0034.png,masks/NIH_0034_mask.png
+NIH_0099,00018840_034.png,images/NIH_0099.png,masks/NIH_0099_mask.png
+NIH_0048,00018984_001.png,images/NIH_0048.png,masks/NIH_0048_mask.png
+NIH_0033,00009863_020.png,images/NIH_0033.png,masks/NIH_0033_mask.png
+NIH_0090,00009218_026.png,images/NIH_0090.png,masks/NIH_0090_mask.png
+NIH_0041,00010315_006.png,images/NIH_0041.png,masks/NIH_0041_mask.png
+NIH_0097,00018610_041.png,images/NIH_0097.png,masks/NIH_0097_mask.png
+NIH_0046,00007444_004.png,images/NIH_0046.png,masks/NIH_0046_mask.png
+NIH_0073,00002227_003.png,images/NIH_0073.png,masks/NIH_0073_mask.png
+NIH_0074,00003538_000.png,images/NIH_0074.png,masks/NIH_0074_mask.png
+NIH_0008,00010007_157.png,images/NIH_0008.png,masks/NIH_0008_mask.png
+NIH_0006,00014004_048.png,images/NIH_0006.png,masks/NIH_0006_mask.png
+NIH_0001,00005502_010.png,images/NIH_0001.png,masks/NIH_0001_mask.png
+NIH_0024,00027441_015.png,images/NIH_0024.png,masks/NIH_0024_mask.png
+NIH_0100,00008943_002.png,images/NIH_0100.png,masks/NIH_0100_mask.png
+NIH_0058,00017470_004.png,images/NIH_0058.png,masks/NIH_0058_mask.png
+NIH_0089,00008237_002.png,images/NIH_0089.png,masks/NIH_0089_mask.png
+NIH_0023,00001483_013.png,images/NIH_0023.png,masks/NIH_0023_mask.png
+NIH_0051,00013922_028.png,images/NIH_0051.png,masks/NIH_0051_mask.png
+NIH_0080,00022034_002.png,images/NIH_0080.png,masks/NIH_0080_mask.png
+NIH_0056,00004832_031.png,images/NIH_0056.png,masks/NIH_0056_mask.png
+NIH_0087,00021449_002.png,images/NIH_0087.png,masks/NIH_0087_mask.png
+NIH_0063,00029855_001.png,images/NIH_0063.png,masks/NIH_0063_mask.png
+NIH_0064,00005593_011.png,images/NIH_0064.png,masks/NIH_0064_mask.png
+NIH_0018,00021154_005.png,images/NIH_0018.png,masks/NIH_0018_mask.png
+NIH_0016,00017214_015.png,images/NIH_0016.png,masks/NIH_0016_mask.png
+NIH_0011,00001248_026.png,images/NIH_0011.png,masks/NIH_0011_mask.png
+NIH_0007,00017110_010.png,images/NIH_0007.png,masks/NIH_0007_mask.png
+NIH_0075,00014177_017.png,images/NIH_0075.png,masks/NIH_0075_mask.png
+NIH_0009,00009479_002.png,images/NIH_0009.png,masks/NIH_0009_mask.png
+NIH_0072,00020408_067.png,images/NIH_0072.png,masks/NIH_0072_mask.png
+NIH_0096,00000997_004.png,images/NIH_0096.png,masks/NIH_0096_mask.png
+NIH_0047,00000643_003.png,images/NIH_0047.png,masks/NIH_0047_mask.png
+NIH_0091,00013249_056.png,images/NIH_0091.png,masks/NIH_0091_mask.png
+NIH_0040,00029154_000.png,images/NIH_0040.png,masks/NIH_0040_mask.png
+NIH_0032,00017801_003.png,images/NIH_0032.png,masks/NIH_0032_mask.png
+NIH_0035,00010352_021.png,images/NIH_0035.png,masks/NIH_0035_mask.png
+NIH_0098,00003510_006.png,images/NIH_0098.png,masks/NIH_0098_mask.png
+NIH_0049,00002239_007.png,images/NIH_0049.png,masks/NIH_0049_mask.png
+NIH_0078,00019176_098.png,images/NIH_0078.png,masks/NIH_0078_mask.png
+NIH_0004,00004342_053.png,images/NIH_0004.png,masks/NIH_0004_mask.png
+NIH_0003,00012364_037.png,images/NIH_0003.png,masks/NIH_0003_mask.png
+NIH_0071,00011702_012.png,images/NIH_0071.png,masks/NIH_0071_mask.png
+NIH_0076,00006481_023.png,images/NIH_0076.png,masks/NIH_0076_mask.png
+NIH_0092,00006322_001.png,images/NIH_0092.png,masks/NIH_0092_mask.png
+NIH_0043,00013760_000.png,images/NIH_0043.png,masks/NIH_0043_mask.png
+NIH_0038,00021341_012.png,images/NIH_0038.png,masks/NIH_0038_mask.png
+NIH_0095,00026908_003.png,images/NIH_0095.png,masks/NIH_0095_mask.png
+NIH_0044,00008037_002.png,images/NIH_0044.png,masks/NIH_0044_mask.png
+NIH_0036,00030772_002.png,images/NIH_0036.png,masks/NIH_0036_mask.png
+NIH_0031,00023073_000.png,images/NIH_0031.png,masks/NIH_0031_mask.png
+NIH_0020,00021420_000.png,images/NIH_0020.png,masks/NIH_0020_mask.png
+NIH_0027,00010684_007.png,images/NIH_0027.png,masks/NIH_0027_mask.png
+NIH_0029,00030573_003.png,images/NIH_0029.png,masks/NIH_0029_mask.png
+NIH_0055,00025513_001.png,images/NIH_0055.png,masks/NIH_0055_mask.png
+NIH_0084,00001684_025.png,images/NIH_0084.png,masks/NIH_0084_mask.png
+NIH_0052,00011543_017.png,images/NIH_0052.png,masks/NIH_0052_mask.png
+NIH_0083,00010773_008.png,images/NIH_0083.png,masks/NIH_0083_mask.png
+NIH_0067,00015443_017.png,images/NIH_0067.png,masks/NIH_0067_mask.png
+NIH_0060,00002395_015.png,images/NIH_0060.png,masks/NIH_0060_mask.png
+NIH_0012,00013613_025.png,images/NIH_0012.png,masks/NIH_0012_mask.png
+NIH_0069,00010680_001.png,images/NIH_0069.png,masks/NIH_0069_mask.png
+NIH_0015,00004156_004.png,images/NIH_0015.png,masks/NIH_0015_mask.png
+NIH_0030,00016508_004.png,images/NIH_0030.png,masks/NIH_0030_mask.png
+NIH_0037,00016291_038.png,images/NIH_0037.png,masks/NIH_0037_mask.png
+NIH_0039,00025822_005.png,images/NIH_0039.png,masks/NIH_0039_mask.png
+NIH_0094,00002386_003.png,images/NIH_0094.png,masks/NIH_0094_mask.png
+NIH_0045,00001317_001.png,images/NIH_0045.png,masks/NIH_0045_mask.png
+NIH_0093,00025954_030.png,images/NIH_0093.png,masks/NIH_0093_mask.png
+NIH_0042,00009945_006.png,images/NIH_0042.png,masks/NIH_0042_mask.png
+NIH_0077,00013993_068.png,images/NIH_0077.png,masks/NIH_0077_mask.png
+NIH_0070,00013527_000.png,images/NIH_0070.png,masks/NIH_0070_mask.png
+NIH_0002,00002350_021.png,images/NIH_0002.png,masks/NIH_0002_mask.png
+NIH_0079,00007576_043.png,images/NIH_0079.png,masks/NIH_0079_mask.png
+NIH_0005,00015443_014.png,images/NIH_0005.png,masks/NIH_0005_mask.png
+NIH_0068,00025839_008.png,images/NIH_0068.png,masks/NIH_0068_mask.png
+NIH_0014,00014626_026.png,images/NIH_0014.png,masks/NIH_0014_mask.png
+NIH_0013,00018080_006.png,images/NIH_0013.png,masks/NIH_0013_mask.png
+NIH_0061,00009114_009.png,images/NIH_0061.png,masks/NIH_0061_mask.png
+NIH_0066,00018610_038.png,images/NIH_0066.png,masks/NIH_0066_mask.png
+NIH_0053,00000063_000.png,images/NIH_0053.png,masks/NIH_0053_mask.png
+NIH_0082,00009138_028.png,images/NIH_0082.png,masks/NIH_0082_mask.png
+NIH_0028,00006498_003.png,images/NIH_0028.png,masks/NIH_0028_mask.png
+NIH_0054,00006620_000.png,images/NIH_0054.png,masks/NIH_0054_mask.png
+NIH_0085,00005288_013.png,images/NIH_0085.png,masks/NIH_0085_mask.png
+NIH_0026,00001836_069.png,images/NIH_0026.png,masks/NIH_0026_mask.png
+NIH_0021,00006679_018.png,images/NIH_0021.png,masks/NIH_0021_mask.png

dataloaders/CXR14.py

@@ -0,0 +1,74 @@
+import logging
+import os
+from typing import List, Tuple, TypeVar
+import pandas as pd
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+from torch import Tensor
+from torchvision import transforms
+from pathlib import Path
+
+PathLike = TypeVar("PathLike", str, Path, None)
+log = logging.getLogger(__name__)
+
+
+PROJECT_DIR = Path(os.path.realpath(__file__)).parent.parent
+DATADIR = Path("<PATH_TO_DATA>/ChestXray-NIHCC/images")
+# can be found at https://nihcc.app.box.com/v/ChestXray-NIHCC/folder/36938765345
+
+
+def build_dataloaders(
+        data_dir: str=DATADIR,
+        img_size: int=128,
+        batch_size: int=16,
+        num_workers: int=1,
+) -> Tuple[List, List, List]:
+    """
+    Build dataloaders for the CXR14 dataset.
+    """
+    train_ds = CXR14Dataset(data_dir, PROJECT_DIR / 'data' / 'train_split.csv', img_size)
+    val_ds = CXR14Dataset(data_dir, PROJECT_DIR / 'data' / 'train_split.csv', img_size)
+    test_ds = CXR14Dataset(data_dir, PROJECT_DIR / 'data' / 'train_split.csv', img_size)
+
+    dataloaders = {}
+    dataloaders['train'] = DataLoader(train_ds, batch_size=batch_size,
+                                      shuffle=True, num_workers=num_workers,
+                                      pin_memory=True)
+    dataloaders['val'] = DataLoader(val_ds, batch_size=batch_size,
+                                    shuffle=False, num_workers=num_workers,
+                                    pin_memory=True)
+    dataloaders['test'] = DataLoader(test_ds, batch_size=batch_size,
+                                     shuffle=False, num_workers=num_workers,
+                                     pin_memory=True)
+
+    return dataloaders
+
+
+
+class CXR14Dataset(Dataset):
+    def __init__(
+        self,
+        data_path: PathLike,
+        csv_path: PathLike,
+        img_size: int,
+    ) -> None:
+        super().__init__()
+        assert(os.path.isdir(data_path))
+        assert(os.path.isfile(csv_path))
+
+        self.data_path = Path(data_path)
+        self.df = pd.read_csv(csv_path)
+        self.img_size = img_size
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def load_image(self, fname: str) -> Tensor:
+        img = Image.open(self.data_path /fname).convert('L').resize((self.img_size, self.img_size))
+        img = transforms.ToTensor()(img).float()
+        return img
+
+    def __getitem__(self, index) -> Tuple[Tensor, Tensor]:
+        img = self.load_image(self.df.loc[index, "Image Index"])
+        return img

dataloaders/JSRT.py

@@ -0,0 +1,94 @@
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+from typing import List, Tuple, TypeVar, Optional
+import pandas as pd
+import os
+from PIL import Image
+from torch import Tensor
+from torchvision import transforms
+import torch
+
+PathLike = TypeVar("PathLike", str, Path, None)
+
+PROJECT_DIR = Path(os.path.realpath(__file__)).parent.parent
+DATADIR = Path("<PATH_TO_DATA>/JSRT")
+# can be found at http://db.jsrt.or.jp/eng.php
+
+def build_dataloaders(
+        data_dir: str=DATADIR,
+        img_size: int=128,
+        batch_size: int=16,
+        num_workers: int=1,
+        n_labelled_images: Optional[int] = None,
+        **kwargs
+) -> Tuple[List, List, List]:
+    """
+    Build dataloaders for the JSRT dataset.
+    """
+    train_ds = JSRTDataset(data_dir, PROJECT_DIR / "data", "JSRT_train_split.csv", img_size)
+    if n_labelled_images is not None:
+        train_ds = torch.utils.data.Subset(train_ds, range(n_labelled_images))
+        print(f"Using {n_labelled_images} labelled images")
+    val_ds = JSRTDataset(data_dir, PROJECT_DIR / "data", "JSRT_val_split.csv", img_size)
+    test_ds = JSRTDataset(data_dir, PROJECT_DIR / "data", "JSRT_test_split.csv", img_size)
+
+    dataloaders = {}
+    dataloaders['train'] = DataLoader(train_ds, batch_size=batch_size,
+                                      shuffle=True, num_workers=num_workers,
+                                      pin_memory=True)
+    dataloaders['val'] = DataLoader(val_ds, batch_size=batch_size,
+                                    shuffle=False, num_workers=num_workers,
+                                    pin_memory=True)
+    dataloaders['test'] = DataLoader(test_ds, batch_size=batch_size,
+                                     shuffle=False, num_workers=num_workers,
+                                     pin_memory=True)
+
+    return dataloaders
+
+
+class JSRTDataset(Dataset):
+    def __init__(self, base_path:PathLike, 
+                 csv_path:PathLike,  
+                 csv_name:str, 
+                 img_size:int=128,
+                 labels:List[str] =('right lung', 'left lung', ),
+                 **kwargs) -> None:
+        self.df = pd.read_csv(os.path.join(csv_path, csv_name))
+        self.base_path = Path(base_path)
+        self.labels = labels 
+        self.img_size = img_size
+
+
+    def load_image(self, fname: str) -> Tensor:
+        img = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+        img = transforms.ToTensor()(img).float()
+        return img
+
+    def load_labels(self, fnames: List[str]) -> Tensor:
+        labels = []
+        for fname in fnames:
+            label = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+            # convert to tensor
+            label = transforms.ToTensor()(label).float()
+            # make binary
+            label = (label > .5).float()
+            labels.append(label)
+        # append all labels and merge
+        label = torch.stack(labels).sum(0)
+        # lungs have no overlap (right?)
+        if (label > 1).sum()>0:
+            print("overlapping lungs!", fnames)
+            label = (label > .5)
+        return label
+
+    def __getitem__(self, index) -> Tuple[Tensor, Tensor]:
+        i = self.df.index[index]
+        img = self.load_image(self.df.loc[i, "path"])
+        
+        label_paths = ["SCR/masks/" + item + "/" + self.df.loc[i, 'id']+ ".gif" for item in self.labels]
+        labels = self.load_labels(label_paths)
+
+        return img, labels
+
+    def __len__(self):
+        return len(self.df)

dataloaders/Montgomery.py

@@ -0,0 +1,61 @@
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+from typing import List, Tuple, TypeVar, Optional
+import pandas as pd
+import os
+from PIL import Image
+from torch import Tensor
+from torchvision import transforms
+import torch
+
+PathLike = TypeVar("PathLike", str, Path, None)
+# can be found at https://data.lhncbc.nlm.nih.gov/public/Tuberculosis-Chest-X-ray-Datasets/Montgomery-County-CXR-Set/MontgomerySet/index.html
+
+
+class MonDataset(Dataset):
+    def __init__(self, base_path:PathLike, 
+                 csv_path:PathLike,  
+                 csv_name:str, 
+                 img_size:int=128,
+                 labels:List[str] =('right lung', 'left lung', ),
+                 **kwargs) -> None:
+        self.df = pd.read_csv(os.path.join(csv_path, csv_name))
+        self.base_path = Path(base_path)
+        self.labels = labels 
+        self.img_size = img_size
+
+
+    def load_image(self, fname: str) -> Tensor:
+        img = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+        img = transforms.ToTensor()(img).float()
+        return img
+
+
+    def load_labels(self, fnames: List[str]) -> Tensor:
+        labels = []
+        for fname in fnames:
+            label = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+            # convert to tensor
+            label = transforms.ToTensor()(label).float()
+            # make binary
+            label = (label > .5).float()
+            labels.append(label)
+        # append all labels and merge
+        label = torch.stack(labels).sum(0)
+        # lungs have no overlap (right?)
+        if (label > 1).sum()>0:
+            print("overlapping lungs!", fnames)
+            label = (label > .5)
+        return label
+
+    def __getitem__(self, index) -> Tuple[Tensor, Tensor]:
+        i = self.df.index[index]
+        img = self.load_image(self.df.loc[i, "scan"])
+        
+        fnames = [ self.df.loc[i, l] for l in self.labels]
+        labels = self.load_labels(fnames)
+
+        return img, labels
+
+    def __len__(self) -> int:
+        return len(self.df)

dataloaders/NIH.py

@@ -0,0 +1,50 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from typing import List, Tuple, TypeVar
+import pandas as pd
+import os
+from PIL import Image
+from torch import Tensor
+from torchvision import transforms
+
+
+PathLike = TypeVar("PathLike", str, Path, None)
+# can be found at https://www.kaggle.com/datasets/nih-chest-xrays/data
+
+class NIHDataset(Dataset):
+    def __init__(self, base_path:PathLike, 
+                 csv_path:PathLike,  
+                 csv_name:str, 
+                 img_size:int=128,
+                 labels:List[str] =('right lung', 'left lung', ),
+                 **kwargs) -> None:
+        self.df = pd.read_csv(os.path.join(csv_path, csv_name))
+        self.base_path = Path(base_path)
+        self.labels = labels 
+        self.img_size = img_size
+
+
+    def load_image(self, fname: str) -> Tensor:
+        img = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+        img = transforms.ToTensor()(img).float()
+        return img
+
+    def load_labels(self, fname: str) -> Tensor:
+
+        label = Image.open(self.base_path /fname).convert('L').resize((self.img_size, self.img_size))
+        # convert to tensor
+        label = transforms.ToTensor()(label).float()
+        # make binary
+        label = (label > .5).float()
+        
+        return label
+
+    def __getitem__(self, index) -> Tuple[Tensor, Tensor]:
+        i = self.df.index[index]
+        img = self.load_image(self.df.loc[i, "scan"])
+        labels = self.load_labels(self.df.loc[i, "mask"])
+
+        return img, labels
+
+    def __len__(self):
+        return len(self.df)

models/datasetDM_model.py

@@ -0,0 +1,88 @@
+import os
+import torch
+from torch import nn, Tensor
+from typing import Dict, Tuple, Optional
+from argparse import Namespace
+from einops import repeat
+from einops.layers.torch import Rearrange
+from functools import partial
+from models.diffusion_model import DiffusionModel
+from trainers.utils import compare_configs
+
+
+
+# Hooks code inspired by https://www.lyndonduong.com/saving-activations/
+# Accessed on 13Feb23
+def save_activations(
+        activations: Dict,
+        name: str,
+        module: nn.Module,
+        inp: Tuple,
+        out: torch.Tensor
+        ) -> None:
+    """PyTorch Forward hook to save outputs at each forward
+    pass. Mutates specified dict objects with each fwd pass.
+    """
+    #activations[name].append(out.detach().cpu())
+    activations[name] = out.detach().cpu()
+
+
+class DatasetDM(nn.Module):
+    def __init__(self, args: Namespace) -> None:
+        super().__init__()
+        # Load the model
+        if not os.path.isfile(args.saved_diffusion_model):
+            self.diffusion_model = DiffusionModel(args)
+            if args.verbose:
+                print(f'No model found at {args.saved_diffusion_model}. Please load model!')
+        else:
+            checkpoint = torch.load(args.saved_diffusion_model, map_location=torch.device(args.device))
+            old_config = checkpoint['config']
+            compare_configs(old_config, args)
+            self.diffusion_model = DiffusionModel(old_config)
+            self.diffusion_model.load_state_dict(checkpoint['model_state_dict'])
+        self.diffusion_model.eval()
+
+        # storage for saved activations
+        self._features = {}
+
+        # Note that this only works for the model in model.py
+        for i, (block1, block2, attn, upsample) in enumerate(self.diffusion_model.model.ups):
+            attn.register_forward_hook(
+                partial(save_activations, self._features, i)
+            )
+
+        self.steps = args.t_steps_to_save
+
+        self.classifier = nn.Sequential(
+            nn.Conv2d(960 * len(self.steps), 128, 1), 
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 32, 1),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.Conv2d(32, 1, 1))
+        
+
+    @torch.no_grad()
+    def extract_features(self, x_0: Tensor, noise: Optional[Tensor] = None) -> Dict[int, Tensor]:
+        if noise is not None:
+            assert(x_0.shape == noise.shape)
+        activations=[]
+        for t_step in self.steps:
+            # Add t_steps of noise to x_0 - forward process
+            t_step = torch.Tensor([t_step]).long().to(x_0.device)
+            t_step = repeat(t_step, '1 -> b', b=x_0.shape[0])
+            x_t, _ = self.diffusion_model.forward_diffusion_model(x_0=x_0, t=t_step, noise=noise)
+            # Remove one step of noise from x_t - backward process
+            _ = self.diffusion_model.model(x_t, t_step)
+            # Resize features so that they all live in the image space
+            for idx in self._features:
+                activations.append(nn.functional.interpolate(self._features[idx], size=[x_0.shape[-1]] * 2))
+            # Return activations
+        return torch.cat(activations, dim=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        features = self.extract_features(x).to(x.device)
+        out = self.classifier(features)
+        return out

models/diffusion_model.py

@@ -0,0 +1,301 @@
+"""Adapted from https://github.com/lucidrains/denoising-diffusion-pytorch"""
+from argparse import Namespace
+import math
+from typing import List, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+
+from einops import reduce, rearrange
+from torch import nn, Tensor
+
+from models.unet_model import Unet
+from trainers.utils import default, get_index_from_list, normalize_to_neg_one_to_one
+
+
+def linear_beta_schedule(
+    timesteps: int,
+    start: float = 0.0001,
+    end: float = 0.02
+) -> Tensor:
+    """
+    :param timesteps: Number of time steps
+
+    :return schedule: betas at every timestep, (timesteps,)
+    """
+    scale = 1000 / timesteps
+    beta_start = scale * start
+    beta_end = scale * end
+    return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float32)
+
+
+def cosine_beta_schedule(timesteps: int, s: float = 0.008) -> Tensor:
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+
+    :param timesteps: Number of time steps
+    :param s: scaling factor
+
+    :return schedule: betas at every timestep, (timesteps,)
+    """
+    steps = timesteps + 1
+    x = torch.linspace(0, timesteps, steps, dtype=torch.float32)
+    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clip(betas, 0, 0.999)
+
+
+class DiffusionModel(nn.Module):
+    def __init__(self, config: Namespace):
+        super().__init__()
+
+        # Default parameters
+        self.config = config
+        dim: int = self.default('dim', 64)
+        dim_mults: List[int] = self.default('dim_mults', [1, 2, 4, 8])
+        channels: int = self.default('channels', 1)
+        timesteps: int = self.default('timesteps', 1000)
+        beta_schedule: str = self.default('beta_schedule', 'cosine')
+        objective: str = self.default('objective', 'pred_noise')  # 'pred_noise' or 'pred_x_0'
+        p2_loss_weight_gamma: float = self.default('p2_loss_weight_gamma', 0.)  # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
+        p2_loss_weight_k: float = self.default('p2_loss_weight_k', 1.)
+        dynamic_threshold_percentile: float = self.default('dynamic_threshold_percentile', 0.995)
+
+        self.timesteps = timesteps
+        self.objective = objective
+        self.dynamic_threshold_percentile = dynamic_threshold_percentile
+        self.model = Unet(
+            dim,
+            dim_mults=dim_mults,
+            channels=channels
+        )
+
+        if beta_schedule == 'linear':
+            betas = linear_beta_schedule(timesteps)
+        elif beta_schedule == 'cosine':
+            betas = cosine_beta_schedule(timesteps)
+        else:
+            raise ValueError(f'unknown beta schedule {beta_schedule}')
+
+        alphas = 1. - betas
+        alphas_cumprod = torch.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.)
+
+        # Calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_recip_alphas_cumprod',
+                             torch.sqrt(1. / alphas_cumprod))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod',
+                             torch.sqrt(1. / alphas_cumprod - 1))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod',
+                             torch.sqrt(1. - alphas_cumprod))
+
+        # Calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        self.register_buffer('posterior_variance', posterior_variance)
+
+        self.register_buffer(
+            'posterior_log_variance_clipped',
+            torch.log(posterior_variance.clamp(min=1e-20))
+        )
+        self.register_buffer(
+            'posterior_mean_coef1',
+            betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)
+        )
+        self.register_buffer(
+            'posterior_mean_coef2',
+            (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod)
+        )
+
+        # p2 reweighting
+        p2_loss_weight = ((p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod))
+                          ** (-p2_loss_weight_gamma))
+        self.register_buffer('p2_loss_weight', p2_loss_weight)
+
+    def default(self, val, d):
+        return vars(self.config)[val] if val in self.config else d
+
+    def train_step(self, x_0: Tensor, cond: Optional[Tensor] = None, t:Optional[Tensor] = None) -> Tensor:
+        N, device = x_0.shape[0], x_0.device
+        
+        # If t is not none, use it, otherwise sample from uniform
+        if t is not None:
+            t = t.long().to(device)
+        else:
+            t = torch.randint(0, self.timesteps, (N,), device=device).long()  # (N)
+
+        model_out, noise = self(x_0, t, cond=cond)
+
+        if self.objective == 'pred_noise':
+            target = noise  # (N, C, H, W)
+        elif self.objective == 'pred_x_0':
+            target = x_0  # (N, C, H, W)
+        else:
+            raise ValueError(f'unknown objective {self.objective}')
+
+        loss = F.l1_loss(model_out, target, reduction='none')  # (N, C, H, W)
+        loss = reduce(loss, 'b ... -> b (...)', 'mean')  # (N, (C x H x W))
+
+        # p2 reweighting
+        loss = loss * get_index_from_list(self.p2_loss_weight, t, loss.shape)
+        return loss.mean()
+
+    def val_step(self, x_0: Tensor, cond: Optional[Tensor] = None, t_steps:Optional[int] = None) -> Tensor:
+        if not t_steps:
+            t_steps = self.timesteps
+        step_size = self.timesteps // t_steps
+        N, device = x_0.shape[0], x_0.device
+        losses = []
+        for t in range(0, self.timesteps, step_size):  
+            t = torch.ones((N,)) * t
+            t = t.long().to(device)
+            losses.append(self.train_step(x_0, cond, t))
+
+        return torch.stack(losses).mean()
+
+    def forward(self, x_0: Tensor, t: Tensor, cond: Optional[Tensor] = None) -> Tensor:
+        """
+        Noise x_0 for t timestep and get the model prediction.
+
+        :param x_0: Clean image, (N, C, H, W)
+        :param t: Timestep, (N,)
+        :param cond: element to condition the reconstruction on - eg image when x_0 is a segmentation (N, C', H, W)
+
+        :return pred: Model output, predicted noise or image, (N, C, H, W)
+        :return noise: Added noise, (N, C, H, W)
+        """
+        if self.config.normalize:
+            x_0 = normalize_to_neg_one_to_one(x_0)
+        if cond is not None and self.config.normalize:
+            cond = normalize_to_neg_one_to_one(cond)
+        x_t, noise = self.forward_diffusion_model(x_0, t)
+        return self.model(x_t, t, cond), noise
+
+    def forward_diffusion_model(
+        self,
+        x_0: Tensor,
+        t: Tensor,
+        noise: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Takes an image and a timestep as input and returns the noisy version
+        of it.
+
+        :param x_0: Image at timestep 0, (N, C, H, W)
+        :param t: Timestep, (N)
+        :param cond: element to condition the reconstruction on - eg image when x_0 is a segmentation (N, C', H, W)
+
+        :return x_t: Noisy image at timestep t, (N, C, H, W)
+        :return noise: Noise added to the image, (N, C, H, W)
+        """
+        noise = default(noise, lambda: torch.randn_like(x_0))
+
+        sqrt_alphas_cumprod_t = get_index_from_list(
+            self.sqrt_alphas_cumprod, t, x_0.shape)
+        sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
+            self.sqrt_one_minus_alphas_cumprod, t, x_0.shape)
+
+        # mean + variance
+        x_t = sqrt_alphas_cumprod_t * x_0 + sqrt_one_minus_alphas_cumprod_t * noise
+
+        return x_t, noise
+
+    @torch.no_grad()
+    def sample_timestep(self, x_t: Tensor, t: int, cond=Optional[Tensor]) -> Tensor:
+        """
+        Sample from the model.
+        :param x_t: Image noised t times, (N, C, H, W)
+        :param t: Timestep
+        :return: Sampled image, (N, C, H, W)
+        """
+        N = x_t.shape[0]
+        device = x_t.device
+        batched_t = torch.full((N,), t, device=device, dtype=torch.long)  # (N)
+        model_mean, model_log_variance, _ = self.p_mean_variance(x_t, batched_t, cond=cond)
+        noise = torch.randn_like(x_t) if t > 0 else 0.
+        pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
+        return pred_img
+
+    def p_mean_variance(self, x_t: Tensor, t: Tensor, clip_denoised: bool = True, cond:Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:
+        _, pred_x_0 = self.model_predictions(x_t, t, cond=cond)
+
+        if clip_denoised:
+            # pred_x_0.clamp_(-1., 1.)
+            # Dynamic thrsholding
+            s = torch.quantile(rearrange(pred_x_0, 'b ... -> b (...)').abs(),
+                               self.dynamic_threshold_percentile,
+                               dim=1)
+            s = torch.max(s, torch.tensor(1.0))[:, None, None, None]
+            pred_x_0 = torch.clip(pred_x_0, -s, s) / s
+
+        (model_mean,
+         posterior_log_variance) = self.q_posterior(pred_x_0, x_t, t)
+        return model_mean, posterior_log_variance, pred_x_0
+
+    def model_predictions(self, x_t: Tensor, t: Tensor, cond:Optional[Tensor] = None) \
+            -> Tuple[Tensor, Tensor]:
+        """
+        Return the predicted noise and x_0 for a given x_t and t.
+
+        :param x_t: Noised image at timestep t, (N, C, H, W)
+        :param t: Timestep, (N,)
+        :return pred_noise: Predicted noise, (N, C, H, W)
+        :return pred_x_0: Predicted x_0, (N, C, H, W)
+        """
+        model_output = self.model(x_t, t, cond)
+
+        if self.objective == 'pred_noise':
+            pred_noise = model_output
+            pred_x_0 = self.predict_x_0_from_noise(x_t, t, model_output)
+
+        elif self.objective == 'pred_x_start':
+            pred_noise = self.predict_noise_from_x_0(x_t, t, model_output)
+            pred_x_0 = model_output
+
+        return pred_noise, pred_x_0
+
+    def q_posterior(self, x_start: Tensor, x_t: Tensor, t: Tensor) \
+            -> Tuple[Tensor, Tensor]:
+        posterior_mean = (
+            get_index_from_list(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + get_index_from_list(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_log_variance_clipped = get_index_from_list(
+            self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_log_variance_clipped
+
+    def predict_x_0_from_noise(self, x_t: Tensor, t: Tensor, noise: Tensor) \
+            -> Tensor:
+        """
+        Get x_0 given x_t, t, and the known or predicted noise.
+
+        :param x_t: Noised image at timestep t, (N, C, H, W)
+        :param t: Timestep, (N,)
+        :param noise: Noise, (N, C, H, W)
+        :return: Predicted x_0, (N, C, H, W)
+        """
+        return (
+            get_index_from_list(
+                self.sqrt_recip_alphas_cumprod, t, x_t.shape)
+            * x_t
+            - get_index_from_list(
+                self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+            * noise
+        )
+
+    def predict_noise_from_x_0(self, x_t: Tensor, t: Tensor, x_0: Tensor) \
+            -> Tensor:
+        """
+        Get noise given the known or predicted x_0, x_t, and t
+
+        :param x_t: Noised image at timestep t, (N, C, H, W)
+        :param t: Timestep, (N,)
+        :param noise: Noise, (N, C, H, W)
+        :return: Predicted noise, (N, C, H, W)
+        """
+        return (
+            (get_index_from_list(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x_0)
+            / get_index_from_list(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+        )

models/global_local_cl.py

@@ -0,0 +1,111 @@
+from models.unet_model import Unet, default
+from torch import Tensor, nn
+import torch
+from typing import Optional, List   
+from einops.layers.torch import Rearrange
+
+
+class GlobalCL(Unet):
+    def __init__(self, 
+                 img_size,
+                 dim: int = 64,
+                 init_dim: Optional[int] = None,
+                 dim_mults: List[int] = [1, 2, 4, 8],
+                  **kwargs):
+        super().__init__(**kwargs)
+        init_dim = default(init_dim, dim)
+        # from the paper 
+        g_emb= 1024
+        g_out = 128
+        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
+        mid_dim = dims[-1]
+        mid_img_size = img_size
+        for _ in range(len(dims)-2):
+            mid_img_size = int((mid_img_size -1) / 2) + 1
+        self.g1 = nn.Sequential(
+            Rearrange('b c h w -> b (c h w)'),
+            nn.Linear(mid_dim * mid_img_size ** 2, g_emb, bias=False),
+            nn.ReLU(),
+            nn.Linear(g_emb, g_out, bias=False),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.init_conv(x)
+
+        t = None
+
+        for block1, block2, attn, downsample in self.downs:
+            x = block1(x, t)
+
+            x = block2(x, t)
+            x = attn(x)
+
+            x = downsample(x)
+
+        x = self.mid_block1(x, t)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, t)
+
+        x = self.g1(x)
+        return x
+
+
+class LocalCL(Unet):
+    def __init__(self, 
+                 img_size,
+                 dim: int = 64,
+                 init_dim: Optional[int] = None,
+                 dim_mults: List[int] = [1, 2, 4, 8],
+                  **kwargs):
+        super().__init__(**kwargs)
+        init_dim = default(init_dim, dim)
+        # from the paper 
+        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
+        #g_2 small network with two 1x1 convolutions
+        self.l = 2
+        mid_dim = dims[-self.l-1]
+        self.g2 = nn.Sequential(
+            nn.Conv2d(mid_dim, mid_dim, 1, bias=False),
+            nn.ReLU(),
+            nn.BatchNorm2d(mid_dim),
+            nn.Conv2d(mid_dim, mid_dim, 1, bias=False),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.init_conv(x)
+        r = x.clone()
+
+        t = None
+
+        h = []
+
+        for block1, block2, attn, downsample in self.downs:
+            x = block1(x, t)
+            h.append(x)
+
+            x = block2(x, t)
+            x = attn(x)
+            h.append(x)
+
+            x = downsample(x)
+
+        x = self.mid_block1(x, t)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, t)
+
+        for block1, block2, attn, upsample in self.ups[:self.l]:
+            x = torch.cat((x, h.pop()), dim=1)
+            x = block1(x, t)
+
+            x = torch.cat((x, h.pop()), dim=1)
+            x = block2(x, t)
+            x = attn(x)
+
+            x = upsample(x)
+        
+        x = self.g2(x)
+        return x
+
+
+
+

models/unet_model.py

@@ -0,0 +1,375 @@
+"""Adapted from https://github.com/lucidrains/denoising-diffusion-pytorch"""
+import math
+from collections import namedtuple
+from functools import partial
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import einsum, nn, Tensor
+
+from trainers.utils import default, exists
+
+# constants
+
+ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
+
+# helpers functions
+
+
+def l2norm(t: Tensor) -> Tensor:
+    """L2 normalize along last dimension"""
+    return F.normalize(t, dim=-1)
+
+
+# small helper modules
+
+
+class Residual(nn.Module):
+    """Residual of any Module -> x' = f(x) + x"""
+    def __init__(self, fn: nn.Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        return self.fn(x, *args, **kwargs) + x
+
+
+def Upsample(dim: int, dim_out: Optional[int] = None) -> nn.Sequential:
+    """UpsampleConv with factor 2"""
+    return nn.Sequential(
+        nn.Upsample(scale_factor=2, mode='nearest'),
+        nn.Conv2d(dim, default(dim_out, dim), 3, padding=1)
+    )
+
+
+def Downsample(dim: int, dim_out: Optional[int] = None) -> nn.Conv2d:
+    """Strided Conv2d for downsampling"""
+    return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+
+    def forward(self, x: Tensor) -> Tensor:
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
+        var = torch.var(x, dim=1, unbiased=False, keepdim=True)
+        mean = torch.mean(x, dim=1, keepdim=True)
+        return (x - mean) * (var + eps).rsqrt() * self.g
+
+
+class PreNorm(nn.Module):
+    """Apply LayerNorm before any Module"""
+    def __init__(self, dim: int, fn: nn.Module):
+        super().__init__()
+        self.fn = fn
+        self.norm = LayerNorm(dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.norm(x)
+        return self.fn(x)
+
+
+class SinusoidalPosEmb(nn.Module):
+    """Classical sinosoidal embedding"""
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t: Tensor) -> Tensor:
+        """
+        :param t: Batch of time steps (b,)
+        :return emb: Sinusoidal time embedding (b, dim)
+        """
+        device = t.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = t[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class LearnedSinusoidalPosEmb(nn.Module):
+    """ following @crowsonkb 's lead with learned sinusoidal pos emb """
+    """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
+    def __init__(self, dim: int):
+        super().__init__()
+        assert (dim % 2) == 0
+        half_dim = dim // 2
+        self.weights = nn.Parameter(torch.randn(half_dim))
+
+    def forward(self, t: Tensor) -> Tensor:
+        """
+        :param t: Batch of time steps (b,)
+        :return fouriered: Concatenation of t and time embedding (b, dim + 1)
+        """
+        t = rearrange(t, 'b -> b 1')
+        freqs = t * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
+        fouriered = torch.cat((t, fouriered), dim=-1)
+        return fouriered
+
+# building block modules
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int, dim_out: int, groups: int = 8):
+        super().__init__()
+        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)
+        self.norm = nn.GroupNorm(groups, dim_out)
+        self.act = nn.SiLU()
+
+    def forward(self, x: Tensor, scale_shift: Optional[Tensor] = None) -> Tensor:
+        x = self.proj(x)
+        x = self.norm(x)
+
+        if exists(scale_shift):
+            scale, shift = scale_shift
+            x = x * (scale + 1) + shift
+
+        x = self.act(x)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        *,
+        time_emb_dim: Optional[int] = None,
+        groups: int = 8
+    ):
+        super().__init__()
+
+        self.time_mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(time_emb_dim, dim_out * 2)
+        ) if exists(time_emb_dim) else None
+
+        self.block1 = Block(dim, dim_out, groups=groups)
+        self.block2 = Block(dim_out, dim_out, groups=groups)
+        if dim != dim_out:
+            self.res_conv = nn.Conv2d(dim, dim_out, 1)
+        else:
+            self.res_conv = nn.Identity()
+
+    def forward(self, x: Tensor, time_emb: Optional[Tensor] = None) -> Tensor:
+        """
+        :param x: Batch of input images (b, c, h, w)
+        :param time_emb: Batch of time embeddings (b, c)
+        """
+        scale_shift = None
+
+        if exists(self.time_mlp) and exists(time_emb):
+            time_emb = self.time_mlp(time_emb)
+            time_emb = rearrange(time_emb, 'b c -> b c 1 1')
+            scale_shift = time_emb.chunk(2, dim=1)
+
+        h = self.block1(x, scale_shift=scale_shift)
+        h = self.block2(h)
+        return h + self.res_conv(x)
+
+
+class LinearAttention(nn.Module):
+    """Attention with linear to_qtv"""
+    def __init__(self, dim: int, heads: int = 4, dim_head: int = 32):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Conv2d(hidden_dim, dim, 1),
+            LayerNorm(dim)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        :param x: Batch of input images (b, c, h, w)
+        """
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads), qkv)
+
+        q = q.softmax(dim=-2)
+        k = k.softmax(dim=-1)
+
+        q = q * self.scale
+        v = v / (h * w)
+
+        context = torch.einsum('b h d n, b h e n -> b h d e', k, v)
+
+        out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
+        out = rearrange(out, 'b h c (x y) -> b (h c) x y', h=self.heads, x=h, y=w)
+        return self.to_out(out)
+
+
+class Attention(nn.Module):
+    """Attention with convolutional to_qtv"""
+    def __init__(
+        self,
+        dim: int,
+        heads: int = 4,
+        dim_head: int = 32,
+        scale: int = 16
+    ):
+        super().__init__()
+        self.scale = scale
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads), qkv)
+
+        q, k = map(l2norm, (q, k))
+
+        sim = einsum('b h d i, b h d j -> b h i j', q, k) * self.scale
+        attn = sim.softmax(dim=-1)
+
+        out = einsum('b h i j, b h d j -> b h i d', attn, v)
+        out = rearrange(out, 'b h (x y) d -> b (h d) x y', x=h, y=w)
+        return self.to_out(out)
+
+# model
+
+
+class Unet(nn.Module):
+    def __init__(
+        self,
+        dim: int = 64,
+        init_dim: Optional[int] = None,
+        out_dim: Optional[int] = None,
+        dim_mults: List[int] = [1, 2, 4, 8],
+        channels: int = 1,
+        resnet_block_groups: int = 8,
+        learned_variance: bool = False,
+        learned_sinusoidal_cond: bool = False,
+        learned_sinusoidal_dim: int = 16,
+        **kwargs
+    ):
+        super().__init__()
+
+        # determine dimensions
+
+        self.channels = channels
+
+        init_dim = default(init_dim, dim)
+        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding=3)
+
+        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
+        in_out = list(zip(dims[:-1], dims[1:]))
+
+        block_class = partial(ResnetBlock, groups=resnet_block_groups)
+
+        # time embeddings
+
+        time_dim = dim * 4
+
+        self.learned_sinusoidal_cond = learned_sinusoidal_cond
+
+        if learned_sinusoidal_cond:
+            sinu_pos_emb = LearnedSinusoidalPosEmb(learned_sinusoidal_dim)
+            fourier_dim = learned_sinusoidal_dim + 1
+        else:
+            sinu_pos_emb = SinusoidalPosEmb(dim)
+            fourier_dim = dim
+
+        self.time_mlp = nn.Sequential(
+            sinu_pos_emb,
+            nn.Linear(fourier_dim, time_dim),
+            nn.GELU(),
+            nn.Linear(time_dim, time_dim)
+        )
+
+        # layers
+
+        self.downs = nn.ModuleList([])
+        self.ups = nn.ModuleList([])
+        num_resolutions = len(in_out)
+
+        for ind, (dim_in, dim_out) in enumerate(in_out):
+            is_last = ind >= (num_resolutions - 1)
+
+            self.downs.append(nn.ModuleList([
+                block_class(dim_in, dim_in, time_emb_dim=time_dim),
+                block_class(dim_in, dim_in, time_emb_dim=time_dim),
+                Residual(PreNorm(dim_in, LinearAttention(dim_in))),
+                Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(
+                    dim_in, dim_out, 3, padding=1)
+            ]))
+
+        mid_dim = dims[-1]
+        self.mid_block1 = block_class(mid_dim, mid_dim, time_emb_dim=time_dim)
+        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
+        self.mid_block2 = block_class(mid_dim, mid_dim, time_emb_dim=time_dim)
+
+        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
+            is_last = ind == (len(in_out) - 1)
+
+            self.ups.append(nn.ModuleList([
+                block_class(dim_out + dim_in, dim_out, time_emb_dim=time_dim),
+                block_class(dim_out + dim_in, dim_out, time_emb_dim=time_dim),
+                Residual(PreNorm(dim_out, LinearAttention(dim_out))),
+                Upsample(dim_out, dim_in) if not is_last else nn.Conv2d(
+                    dim_out, dim_in, 3, padding=1)
+            ]))
+
+        default_out_dim = channels * (1 if not learned_variance else 2)
+        self.out_dim = default(out_dim, default_out_dim)
+
+        self.final_res_block = block_class(dim * 2, dim, time_emb_dim=time_dim)
+        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+
+    def forward(self, x: Tensor, timestep: Optional[Tensor]=None, cond: Optional[Tensor]=None) -> Tensor:
+        x = self.init_conv(x)
+        r = x.clone()
+
+        t = self.time_mlp(timestep) if timestep is not None else None
+
+        h = []
+
+        for block1, block2, attn, downsample in self.downs:
+            x = block1(x, t)
+            h.append(x)
+
+            x = block2(x, t)
+            x = attn(x)
+            h.append(x)
+
+            x = downsample(x)
+
+        x = self.mid_block1(x, t)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, t)
+
+        for block1, block2, attn, upsample in self.ups:
+            x = torch.cat((x, h.pop()), dim=1)
+            x = block1(x, t)
+
+            x = torch.cat((x, h.pop()), dim=1)
+            x = block2(x, t)
+            x = attn(x)
+
+            x = upsample(x)
+
+        x = torch.cat((x, r), dim=1)
+
+        x = self.final_res_block(x, t)
+        return self.final_conv(x)
+
+
+if __name__ == '__main__':
+    model = Unet(channels=1)
+    x = torch.randn(1, 1, 128, 128)
+    y = model(x, timestep=torch.tensor([100]))
+    print(y.shape)

requirements.txt

@@ -0,0 +1,16 @@
+torch>=1.13.0
+torchvision
+tensorboard==2.13.0
+einops==0.6.1
+gradio==3.35.2
+matplotlib==3.7.1
+numpy==1.24.1
+opencv_python==4.8.0.74
+pandas==2.0.2
+Pillow==9.5.0
+scipy==1.11.1
+seaborn==0.12.2
+tqdm==4.65.0
+scikit_image==0.21.0
+protobuf==3.20
+fastapi==0.99.0

train.py

@@ -0,0 +1,56 @@
+from config import parser
+import argparse
+from pathlib import Path
+from trainers.train_CXR14 import main as train_CXR14
+from trainers.train_baseline import main as train_baseline
+from trainers.train_base_diffusion import main as train_JSRT
+from trainers.train_datasetDM import main as train_datasetDM
+from trainers.datasetDM_per_step import main as train_simple_datasetDM
+from trainers.train_global_cl import main as train_global_cl
+from trainers.train_local_cl import main as train_local_cl
+from trainers.finetune_glob_cl import main as train_global_finetune
+from trainers.finetune_glob_loc_cl import main as train_global_local_finetune
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser(parents=[parser], add_help=False)
+    config = parser.parse_args()
+
+    # catch exeptions
+    #if len(config.loss_weights) != 4:
+    #    raise ValueError('loss_weights must be a list of 4 values')
+    
+    config.normalize = True
+    config.log_dir = Path(config.log_dir).parent / config.experiment / str(config.n_labelled_images) /  Path(config.log_dir).name
+    config.channels = 1
+    config.out_channels = 1
+    if config.dataset == "CXR14":
+        config.data_dir = Path("<PATH_TO_DATA>/ChestXray-NIHCC/images")
+    elif config.dataset == "JSRT":
+        config.data_dir = Path("<PATH_TO_DATA>/JSRT")
+    else:
+        raise ValueError(f"Unknown dataset: {config.dataset}")
+    
+
+    if config.experiment == "img_only":
+        train_CXR14(config)
+    elif config.experiment == "baseline":
+        train_baseline(config)
+    elif config.experiment == "LEDM":
+        config.t_steps_to_save = [50, 150, 250]
+        train_datasetDM(config)
+    elif config.experiment == "LEDMe":
+        config.t_steps_to_save = [1, 10, 25, 50, 200, 400, 600, 800]
+        train_datasetDM(config)
+    elif config.experiment == "TEDM":
+        config.shared_weights_over_timesteps = True
+        config.t_steps_to_save = [1, 10, 25, 50, 200, 400, 600, 800]
+        train_datasetDM(config)
+    elif config.experiment == 'global_cl':
+        train_global_cl(config)
+    elif config.experiment == 'local_cl':
+        train_local_cl(config)
+    elif config.experiment == 'global_finetune':
+        train_global_finetune(config)
+    elif config.experiment == 'glob_loc_finetune':
+        train_global_local_finetune(config)

trainers/datasetDM_per_step.py

@@ -0,0 +1,115 @@
+from argparse import Namespace
+import os
+from pathlib import Path
+from dataloaders.JSRT import build_dataloaders
+import torch
+from tqdm.auto import tqdm
+from trainers.utils import seed_everything, TensorboardLogger
+from torch.cuda.amp import GradScaler
+from torch import Tensor, nn
+from typing import Dict, Optional
+from trainers.train_baseline import train
+from models.datasetDM_model import DatasetDM
+from einops import repeat
+from einops.layers.torch import Rearrange
+
+
+class ModDatasetDM(DatasetDM):
+    # the idea here is to pool info per timestep, 
+    # so that we can then use the aggregate for feature importance
+
+    def __init__(self, args: Namespace) -> None:
+        super().__init__(args)
+        self.mean = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.mean_squared = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.std = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.classifier = nn.Conv2d(len(self.steps) * 960, 1, 1)
+    
+    def forward(self, x: Tensor) -> Tensor:
+        features = self.extract_features(x).to(x.device)
+        out = (features - self.mean ) / self.std
+        out = self.classifier(features)
+        return out
+    
+class OneStepPredDatasetDM(DatasetDM):
+    # the idea here is to pool info per timestep, 
+    # so that we can then use the aggregate for feature importance
+
+    def __init__(self, args: Namespace) -> None:
+        super().__init__(args)
+        self.mean = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.mean_squared = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.std = torch.zeros(len(self.steps) * 960, args.img_size, args.img_size, requires_grad=False)
+        self.classifier = nn.Sequential(
+            Rearrange('b (step act) h w -> (b step) act h w', step=len(self.steps)),
+            nn.Conv2d(960, 128, 1),
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 32, 1),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.Conv2d(32, 1, args.out_channels)
+            )
+    
+
+    def forward(self, x: Tensor) -> Tensor:
+        features = self.extract_features(x).to(x.device)
+        out = (features - self.mean ) / self.std
+        out = self.classifier(features)
+        return out
+
+
+def main(config: Namespace) -> None:
+    # adjust logdir to include experiment name
+    os.makedirs(config.log_dir, exist_ok=True)
+    print('Experiment folder: %s' % (config.log_dir))
+
+    # save config namespace into logdir
+    with open(config.log_dir / 'config.txt', 'w') as f:
+        for k, v in vars(config).items():
+            if type(v) not in [str, int, float, bool]:
+                f.write(f'{k}: {str(v)}\n')
+            else:
+                f.write(f'{k}: {v}\n')
+
+    # Random seed
+    seed_everything(config.seed)
+
+    model = ModDatasetDM(config)
+    model = model.to(config.device)
+    model.train()
+
+    optimizer = torch.optim.Adam(model.classifier.parameters(), lr=config.lr, weight_decay=config.weight_decay)  # , betas=config.adam_betas)
+    step = 0
+
+    scaler = GradScaler()
+
+    dataloaders = build_dataloaders(
+        config.data_dir,
+        config.img_size,
+        config.batch_size,
+        config.num_workers,
+        config.n_labelled_images
+    )
+    train_dl = dataloaders['train']
+    val_dl = dataloaders['val']
+
+    # Logger
+    logger = TensorboardLogger(config.log_dir, enabled=not config.debug)
+
+
+    # do a loop to calculate mean and variance of the features
+    # then use those to normalize the features
+    model.to(config.device)
+    for x, _ in tqdm(train_dl, desc="Calculating mean and variance"):
+        x = x.to(config.device)
+        features = model.extract_features(x)
+        model.mean += features.sum(dim=0)
+        model.mean_squared += (features ** 2).sum(dim=0)
+    model.mean = model.mean / len(train_dl.dataset)
+    model.std = (model.mean_squared / len(train_dl.dataset) - model.mean ** 2).sqrt() + 1e-6
+
+    model.mean = model.mean.to(config.device)
+    model.std = model.std.to(config.device)
+        
+    train(config, model, optimizer, train_dl, val_dl, logger, scaler, step)

trainers/finetune_glob_cl.py

@@ -0,0 +1,172 @@
+import argparse
+import os
+from pathlib import Path
+import torch
+from torch import autocast, Tensor
+from torch.nn.functional import binary_cross_entropy_with_logits
+from torch.cuda.amp import GradScaler
+from tqdm import tqdm
+from config import parser
+from einops import rearrange, reduce, repeat
+from dataloaders.JSRT import build_dataloaders
+from models.unet_model import Unet
+from trainers.train_baseline import validate, save
+from trainers.utils import (TensorboardLogger, compare_configs, seed_everything, crop_batch)
+
+
+def train(config, model, optimizer, train_dl, val_dl, logger, scaler, step):
+    best_val_loss = float('inf')
+    train_losses = []
+    if config.dataset == "BRATS2D":
+        train_losses_per_class = []
+    elif config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+        train_losses_per_timestep = []
+        
+    pbar = tqdm(total=config.val_freq, desc='Training')
+    while True:
+        for x, y in train_dl:
+            if config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+                y = repeat(y, 'b c h w -> (b step) c h w', step=len(model.steps))
+            if config.augment_at_finetuning:
+                x, y = crop_batch([x, y], config.img_size, config.batch_size)
+                brightness = torch.rand((config.batch_size, 1, 1, 1), device=x.device)*.6 - .3        # random brightness adjustment between [-.3, .3]
+                contrast = torch.rand((config.batch_size, 1, 1, 1), device=x.device)*.6 + .7          # random contrast adjustment between [.7, 1.3]
+                x = (x + brightness) * contrast                                                # apply brightness and contrast
+
+            x = x.to(config.device)
+            y = y.to(config.device)
+
+            optimizer.zero_grad()
+            with autocast(device_type=config.device, enabled=config.mixed_precision):
+                pred = model(x)
+                # cross entropy loss
+                #loss = - ((y * torch.log(torch.sigmoid(pred)) + (1 - y) * torch.log(1 - torch.sigmoid(pred)))).mean()
+                if config.dataset == "BRATS2D":
+                    weights = repeat(torch.Tensor(config.loss_weights).to(config.device), 'c -> b c h w', b=y.shape[0], h=y.shape[2], w=y.shape[3])
+                else:
+                    weights = None
+                expanded_loss = reduce(binary_cross_entropy_with_logits(pred, y, weight=weights, reduction='none'), 'b c h w -> b c', 'mean') 
+                loss = expanded_loss.mean()
+            scaler.scale(loss).backward()
+            optimizer.step()
+
+            train_losses.append(loss.item())
+            if config.dataset == "BRATS2D":
+                loss_per_class = expanded_loss.mean(0)
+                train_losses_per_class.append(loss_per_class.detach().cpu())
+                pbar.set_description(f'Training loss: {loss.item():.4f} - {loss_per_class[0].item():.4f} - {loss_per_class[1].item():.4f} - {loss_per_class[2].item():.4f} - {loss_per_class[3].item():.4f}')
+            else:
+                pbar.set_description(f'Training loss: {loss.item():.4f}')
+
+            pbar.update(1)
+            step += 1
+            
+            if config.unfreeze_weights_at_step == step:
+                for name, param in model.named_parameters():
+                    if name.startswith('downs') or name.startswith('init_conv') or name.startswith('mid_'):
+                        param.requires_grad = True
+
+            if step % config.log_freq == 0 or config.debug:
+                avg_train_loss = sum(train_losses) / len(train_losses)
+                print(f'Step {step} - Train loss: {avg_train_loss:.4f}')
+                logger.log({'train/loss': avg_train_loss}, step=step)
+                if config.dataset == "BRATS2D":
+                    avg_train_loss_per_class = torch.stack(train_losses_per_class).mean(0)
+                    logger.log({'train_loss/0':avg_train_loss_per_class[0].item()}, step=step)
+                    logger.log({'train_loss/1':avg_train_loss_per_class[1].item()}, step=step)
+                    logger.log({'train_loss/2':avg_train_loss_per_class[2].item()}, step=step)
+                    logger.log({'train_loss/3':avg_train_loss_per_class[3].item()}, step=step)
+                if config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+                    avg_train_loss_per_timestep = torch.stack(train_losses_per_timestep).mean(0)
+                    for i, model_step in enumerate(model.steps):
+                        logger.log({'train_loss/step_' + str(model_step): avg_train_loss_per_timestep[i].item()}, step=step)
+
+            if step % config.val_freq == 0 or config.debug:
+                val_results = validate(config, model, val_dl)
+                logger.log(val_results, step=step)
+
+                if val_results['val/loss'] < best_val_loss and not config.debug:
+                    print(f'Step {step} - New best validation loss: '
+                          f'{val_results["val/loss"]:.4f}, saving model '
+                          f'in {config.log_dir}')
+                    best_val_loss = val_results['val/loss']
+                    save(
+                        model,
+                        optimizer,
+                        config,
+                        config.log_dir / 'best_model.pt',
+                        step
+                    )
+                elif val_results['val/loss'] > best_val_loss * 1.5 and config.early_stop:
+                    print(f'Step {step} - Validation loss increased by more than 50%')
+                    return model
+
+            if step >= config.max_steps or config.debug:
+                return model
+
+
+
+
+def load(config, path):
+    raise NotImplementedError
+
+def main(config):
+
+    os.makedirs(config.log_dir, exist_ok=True)
+    
+    # save config namespace into logdir
+    with open(config.log_dir / 'config.txt', 'w') as f:
+        for k, v in vars(config).items():
+            if type(v) not in [str, int, float, bool]:
+                f.write(f'{k}: {str(v)}\n')
+            else:
+                f.write(f'{k}: {v}\n')
+
+    # Random seed
+    seed_everything(config.seed)
+
+    # Init model and optimizer
+    if config.resume_path is not None:
+        print('Loading model from', config.resume_path)
+        model, optimizer, step = load(config, config.resume_path)
+    else:
+        model = Unet(
+            img_size=config.img_size,
+            dim=config.dim,
+            dim_mults=config.dim_mults,
+            channels=config.channels,
+            out_dim=config.out_channels)
+        state_dict = torch.load(config.global_model_path, map_location='cpu')['model_state_dict']
+        out = model.load_state_dict(state_dict=state_dict, strict=False)
+        print("Loaded state dict. \n\tMissing keys: {}\n\tUnexpected keys: {}".format(out.missing_keys, out.unexpected_keys))
+        print('Note that although the state dict of the decoder is loaded, its values are random.')
+        if config.unfreeze_weights_at_step !=0:
+            for name, param in model.named_parameters():
+                if name.startswith('downs') or name.startswith('init_conv') or name.startswith('mid_'):
+                    param.requires_grad = False
+
+        optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)  # , betas=config.adam_betas)
+        
+        step = 0
+    model.to(config.device)
+    model.train()
+
+    scaler = GradScaler()
+
+    # Load data
+    dataloaders = build_dataloaders(
+        config.data_dir,
+        config.img_size,
+        config.batch_size,
+        config.num_workers,
+        n_labelled_images=config.n_labelled_images,
+    )
+    train_dl = dataloaders['train']
+    val_dl = dataloaders['val']
+    print('Train dataset size:', len(train_dl.dataset))
+    print('Validation dataset size:', len(val_dl.dataset))
+
+    # Logger
+    logger = TensorboardLogger(config.log_dir, enabled=not config.debug)
+
+    train(config, model, optimizer, train_dl, val_dl, logger, scaler, step)

trainers/finetune_glob_loc_cl.py

@@ -0,0 +1,172 @@
+import argparse
+import os
+from pathlib import Path
+import torch
+from torch import autocast, Tensor
+from torch.nn.functional import binary_cross_entropy_with_logits
+from torch.cuda.amp import GradScaler
+from tqdm import tqdm
+from config import parser
+from einops import rearrange, reduce, repeat
+from dataloaders.JSRT import build_dataloaders
+from models.unet_model import Unet
+from trainers.train_baseline import validate, save
+from trainers.utils import (TensorboardLogger, compare_configs, seed_everything, crop_batch)
+
+
+def train(config, model, optimizer, train_dl, val_dl, logger, scaler, step):
+    best_val_loss = float('inf')
+    train_losses = []
+    if config.dataset == "BRATS2D":
+        train_losses_per_class = []
+    elif config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+        train_losses_per_timestep = []
+        
+    pbar = tqdm(total=config.val_freq, desc='Training')
+    while True:
+        for x, y in train_dl:
+            if config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+                y = repeat(y, 'b c h w -> (b step) c h w', step=len(model.steps))
+            if config.augment_at_finetuning:
+                x, y = crop_batch([x, y], config.img_size, config.batch_size)
+                brightness = torch.rand((config.batch_size, 1, 1, 1), device=x.device)*.6 - .3        # random brightness adjustment between [-.3, .3]
+                contrast = torch.rand((config.batch_size, 1, 1, 1), device=x.device)*.6 + .7          # random contrast adjustment between [.7, 1.3]
+                x = (x + brightness) * contrast                                                # apply brightness and contrast
+
+            x = x.to(config.device)
+            y = y.to(config.device)
+
+            optimizer.zero_grad()
+            with autocast(device_type=config.device, enabled=config.mixed_precision):
+                pred = model(x)
+                # cross entropy loss
+                #loss = - ((y * torch.log(torch.sigmoid(pred)) + (1 - y) * torch.log(1 - torch.sigmoid(pred)))).mean()
+                if config.dataset == "BRATS2D":
+                    weights = repeat(torch.Tensor(config.loss_weights).to(config.device), 'c -> b c h w', b=y.shape[0], h=y.shape[2], w=y.shape[3])
+                else:
+                    weights = None
+                expanded_loss = reduce(binary_cross_entropy_with_logits(pred, y, weight=weights, reduction='none'), 'b c h w -> b c', 'mean') 
+                loss = expanded_loss.mean()
+            scaler.scale(loss).backward()
+            optimizer.step()
+
+            train_losses.append(loss.item())
+            if config.dataset == "BRATS2D":
+                loss_per_class = expanded_loss.mean(0)
+                train_losses_per_class.append(loss_per_class.detach().cpu())
+                pbar.set_description(f'Training loss: {loss.item():.4f} - {loss_per_class[0].item():.4f} - {loss_per_class[1].item():.4f} - {loss_per_class[2].item():.4f} - {loss_per_class[3].item():.4f}')
+            else:
+                pbar.set_description(f'Training loss: {loss.item():.4f}')
+
+            pbar.update(1)
+            step += 1
+            
+            if config.unfreeze_weights_at_step == step:
+                for name, param in model.named_parameters():
+                    if name.startswith('downs') or name.startswith('init_conv') or name.startswith('mid_'):
+                        param.requires_grad = True
+
+            if step % config.log_freq == 0 or config.debug:
+                avg_train_loss = sum(train_losses) / len(train_losses)
+                print(f'Step {step} - Train loss: {avg_train_loss:.4f}')
+                logger.log({'train/loss': avg_train_loss}, step=step)
+                if config.dataset == "BRATS2D":
+                    avg_train_loss_per_class = torch.stack(train_losses_per_class).mean(0)
+                    logger.log({'train_loss/0':avg_train_loss_per_class[0].item()}, step=step)
+                    logger.log({'train_loss/1':avg_train_loss_per_class[1].item()}, step=step)
+                    logger.log({'train_loss/2':avg_train_loss_per_class[2].item()}, step=step)
+                    logger.log({'train_loss/3':avg_train_loss_per_class[3].item()}, step=step)
+                if config.shared_weights_over_timesteps and config.experiment == 'datasetDM':
+                    avg_train_loss_per_timestep = torch.stack(train_losses_per_timestep).mean(0)
+                    for i, model_step in enumerate(model.steps):
+                        logger.log({'train_loss/step_' + str(model_step): avg_train_loss_per_timestep[i].item()}, step=step)
+
+            if step % config.val_freq == 0 or config.debug:
+                val_results = validate(config, model, val_dl)
+                logger.log(val_results, step=step)
+
+                if val_results['val/loss'] < best_val_loss and not config.debug:
+                    print(f'Step {step} - New best validation loss: '
+                          f'{val_results["val/loss"]:.4f}, saving model '
+                          f'in {config.log_dir}')
+                    best_val_loss = val_results['val/loss']
+                    save(
+                        model,
+                        optimizer,
+                        config,
+                        config.log_dir / 'best_model.pt',
+                        step
+                    )
+                elif val_results['val/loss'] > best_val_loss * 1.5 and config.early_stop:
+                    print(f'Step {step} - Validation loss increased by more than 50%')
+                    return model
+
+            if step >= config.max_steps or config.debug:
+                return model
+
+
+
+
+def load(config, path):
+    raise NotImplementedError
+
+def main(config):
+
+    os.makedirs(config.log_dir, exist_ok=True)
+    
+    # save config namespace into logdir
+    with open(config.log_dir / 'config.txt', 'w') as f:
+        for k, v in vars(config).items():
+            if type(v) not in [str, int, float, bool]:
+                f.write(f'{k}: {str(v)}\n')
+            else:
+                f.write(f'{k}: {v}\n')
+
+    # Random seed
+    seed_everything(config.seed)
+
+    # Init model and optimizer
+    if config.resume_path is not None:
+        print('Loading model from', config.resume_path)
+        model, optimizer, step = load(config, config.resume_path)
+    else:
+        model = Unet(
+            img_size=config.img_size,
+            dim=config.dim,
+            dim_mults=config.dim_mults,
+            channels=config.channels,
+            out_dim=config.out_channels)
+        state_dict = torch.load(config.glob_loc_model_path, map_location='cpu')['model_state_dict']
+        out = model.load_state_dict(state_dict=state_dict, strict=False)
+        print("Loaded state dict. \n\tMissing keys: {}\n\tUnexpected keys: {}".format(out.missing_keys, out.unexpected_keys))
+        print('Note that although the state dict of the decoder is loaded, its values are random.')
+        if config.unfreeze_weights_at_step !=0:
+            for name, param in model.named_parameters():
+                if name.startswith('downs') or name.startswith('init_conv') or name.startswith('mid_'):
+                    param.requires_grad = False
+
+        optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)  # , betas=config.adam_betas)
+        
+        step = 0
+    model.to(config.device)
+    model.train()
+
+    scaler = GradScaler()
+
+    # Load data
+    dataloaders = build_dataloaders(
+        config.data_dir,
+        config.img_size,
+        config.batch_size,
+        config.num_workers,
+        n_labelled_images=config.n_labelled_images,
+    )
+    train_dl = dataloaders['train']
+    val_dl = dataloaders['val']
+    print('Train dataset size:', len(train_dl.dataset))
+    print('Validation dataset size:', len(val_dl.dataset))
+
+    # Logger
+    logger = TensorboardLogger(config.log_dir, enabled=not config.debug)
+
+    train(config, model, optimizer, train_dl, val_dl, logger, scaler, step)