Noise paper

LION/experiments/ct_learned_denoising_experiments.py

@@ -0,0 +1,140 @@
+# =============================================================================
+# This file is part of LION library
+# License : GPL-3
+#
+# Author  : Ander Biguri, Max Kiss
+# =============================================================================
+
+
+import numpy as np
+import torch
+import pathlib
+import warnings
+from abc import ABC, abstractmethod, ABCMeta
+
+from LION.utils.parameter import LIONParameter
+import LION.CTtools.ct_geometry as ctgeo
+import LION.CTtools.ct_utils as ct
+from LION.data_loaders.deteCT import deteCT
+
+from LION.experiments.ct_experiments import Experiment
+
+
+class CTNoiseExperiment(Experiment):
+    def __init__(self, experiment_params=None, dataset="2DeteCT", datafolder=None):
+        dataset = "2DeteCT"
+        if dataset != "2DeteCT":
+            raise ValueError(
+                "Benchmarking experiments only supports 2DeteCT dataset, currently"
+            )
+        super().__init__(experiment_params, dataset, datafolder)
+
+class ExperimentalNoiseDenoising(CTNoiseExperiment):
+
+    def __init__(self, experiment_params=None, dataset="2DeteCT", datafolder=None):
+        super().__init__(experiment_params, dataset, datafolder)
+
+    @staticmethod
+    def default_parameters(dataset="2DeteCT"):
+        param = LIONParameter()
+
+        param.name = "Experimental Noise Denoising experiment with noisy sinograms from mode1 as input data and sinograms from mode2 as target data"
+
+        # Parameters for the geometry
+        param.geo = deteCT.get_default_geometry()
+
+        param.log_transform = False # if we do a sino2sino this should not be done
+
+        # leave this untouched
+        param.data_loader_params = Experiment.get_dataset_parameters(dataset)
+
+        # Change the data loader to be sino2sino
+        param.data_loader_params.task = "sino2sino"
+        param.data_loader_params.input_mode = "mode1"
+        param.data_loader_params.target_mode = "mode2"
+
+        return param
+
+class ArtificialNoiseDenoising(CTNoiseExperiment):
+
+    def __init__(self, experiment_params=None, dataset="2DeteCT", datafolder=None):
+        super().__init__(experiment_params, dataset, datafolder)
+
+    @staticmethod
+    def default_parameters(dataset="2DeteCT"):
+        param = LIONParameter()
+
+        param.name = "Artificial Noise Denoising experiment with noisy sinograms generated from mode2 as input data and sinograms from mode2 as target data"
+
+        # Parameters for the geometry
+        param.geo = deteCT.get_default_geometry()
+
+        param.log_transform = False # if we do a sino2sino this should not be done
+
+        # leave this untouched
+        param.data_loader_params = Experiment.get_dataset_parameters(dataset)
+
+        # Change the data loader to be sino2sino
+        param.data_loader_params.task = "sino2sino"
+        param.data_loader_params.input_mode = "mode2"
+        param.data_loader_params.target_mode = "mode2"
+
+        param.data_loader_params.add_noise = True
+        param.data_loader_params.noise_params = LIONParameter()
+        param.data_loader_params.noise_params.I0 = 200 #10000
+        param.data_loader_params.noise_params.cross_talk = 0.05 #from XCIST
+
+        return param
+
+class ExperimentalNoiseDenoisingRecon(CTNoiseExperiment):
+
+    def __init__(self, experiment_params=None, dataset="2DeteCT", datafolder=None):
+        super().__init__(experiment_params, dataset, datafolder)
+
+    @staticmethod
+    def default_parameters(dataset="2DeteCT"):
+        param = LIONParameter()
+
+        param.name = "Experimental Noise Denoising experiment with noisy sinograms from mode1 as input data and recons from mode2 as target data"
+
+        # Parameters for the geometry
+        param.geo = deteCT.get_default_geometry()
+
+        # leave this untouched
+        param.data_loader_params = Experiment.get_dataset_parameters(dataset)
+
+        # Change the data loader to be sino2sino
+        param.data_loader_params.task = "sino2recon"
+        param.data_loader_params.input_mode = "mode1"
+        param.data_loader_params.target_mode = "mode2"
+
+        return param
+
+class ArtificialNoiseDenoisingRecon(CTNoiseExperiment):
+
+    def __init__(self, experiment_params=None, dataset="2DeteCT", datafolder=None):
+        super().__init__(experiment_params, dataset, datafolder)
+
+    @staticmethod
+    def default_parameters(dataset="2DeteCT"):
+        param = LIONParameter()
+
+        param.name = "Artificial Noise Denoising experiment with noisy sinograms generated from mode2 as input data and recons from mode2 as target data"
+
+        # Parameters for the geometry
+        param.geo = deteCT.get_default_geometry()
+
+        # leave this untouched
+        param.data_loader_params = Experiment.get_dataset_parameters(dataset)
+
+        # Change the data loader to be sino2sino
+        param.data_loader_params.task = "sino2recon"
+        param.data_loader_params.input_mode = "mode2"
+        param.data_loader_params.target_mode = "mode2"
+
+        param.data_loader_params.add_noise = True
+        param.data_loader_params.noise_params = LIONParameter()
+        param.data_loader_params.noise_params.I0 = 200 #10000
+        param.data_loader_params.noise_params.cross_talk = 0.05 #from XCIST
+
+        return param

LION/models/CNNs/UNets/UNet_3.py

@@ -0,0 +1,120 @@
+# =============================================================================
+# This file is part of LION library
+# License : GPL-3
+#
+# From: https://github.com/usuyama/pytorch-unet
+# Authors: Max Kiss, Ander Biguri
+# =============================================================================
+import torch.nn as nn
+import torch
+import math
+import pdb
+
+from LION.models import LIONmodel
+from LION.utils.math import power_method
+from LION.utils.parameter import LIONParameter
+import LION.CTtools.ct_geometry as ct
+import LION.CTtools.ct_utils as ct_utils
+import LION.utils.utils as ai_utils
+from collections import OrderedDict
+import sys
+import os
+import torch.nn.functional as F
+
+
+class double_conv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(double_conv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, 3,padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, 3,padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+class inconv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(inconv, self).__init__()
+        self.conv = double_conv(in_ch, out_ch)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+class down(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(down, self).__init__()
+        self.mpconv = nn.Sequential(
+            nn.MaxPool2d(2),
+            double_conv(in_ch, out_ch)
+        )
+
+    def forward(self, x):
+        x = self.mpconv(x)
+        return x
+
+class up(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(up, self).__init__()
+        self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)
+        self.conv = double_conv(in_ch, out_ch)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        diffY = x2.size()[2] - x1.size()[2] #Won't work for 3d Images
+        diffX = x2.size()[3] - x1.size()[3]
+        x1 = F.pad(x1, (diffX // 2, diffX - diffX//2, diffY // 2, diffY - diffY//2))
+        x = torch.cat([x2, x1], dim=1)
+        x = self.conv(x)
+        return x
+
+class outconv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(outconv, self).__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, 1)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+class UNet(LIONmodel.LIONmodel):
+    def __init__(self, geometry_parameters: ct.Geometry, model_parameters: LIONParameter = None):
+        super().__init__(model_parameters, geometry_parameters)
+        self.hypers={}
+        self.inc = inconv(1, 64)
+        self.down1 = down(64, 128)
+        self.down2 = down(128, 256)
+        self.down3 = down(256, 512)
+        self.down4 = down(512, 512)
+        self.up1 = up(1024, 256)
+        self.up2 = up(512, 128)
+        self.up3 = up(256, 64)
+        self.up4 = up(128, 64)
+        self.outc = outconv(64, 1)
+
+    @staticmethod
+    def default_parameters():
+        UNet_params = LIONParameter()
+        return UNet_params
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        x = self.outc(x)
+        return x
+
+    def init_weights(self,m):
+        pass

LION/models/post_processing/FBPMSDNet.py

@@ -0,0 +1,108 @@
+# This file is part of LION library
+# License : BSD-3
+#
+# Author  : See ./MS-D. Daniel Pelt et al
+# Modifications: Ander Biguri
+# =============================================================================
+
+
+import torch
+import torch.nn as nn
+from LION.models import LIONmodel
+import LION.CTtools.ct_geometry as ct
+from LION.utils.parameter import LIONParameter
+import msd_pytorch as msd
+
+import tomosipo as ts
+from tomosipo.torch_support import to_autograd
+from ts_algorithms import fdk
+
+class FBPMS_D(LIONmodel.LIONmodel):
+    def __init__(self, geometry_parameters: ct.Geometry, model_parameters=None):
+
+        super().__init__(model_parameters, geometry_parameters)
+
+        if self.model_parameters.type == "regression":
+            model = msd.MSDRegressionModel(
+                self.model_parameters.c_in,
+                self.model_parameters.c_out,
+                self.model_parameters.depth,
+                self.model_parameters.width,
+                dilations=self.model_parameters.dilations,
+                loss="L2",
+            )
+
+        elif self.model_parameters.type == "segmentation":
+            if self.model_parameters.num_labels is None:
+                raise ValueError(
+                    "For a segmentation network, please set the model_parameters.num_labels vairable to the number of labels in the training set"
+                )
+            model = msd.MSDSegmentationModel(
+                self.model_parameters.c_in,
+                self.model_parameters.num_labels,
+                self.model_parameters.depth,
+                self.model_parameters.width,
+                dilations=self.model_parameters.dilations,
+            )
+        # We don't want MS-D to define our optimizer.
+        model.optimizer = None
+        self.net = model.net
+
+        self._make_operator()
+
+    @staticmethod
+    def default_parameters(mode="regression"):
+        param = LIONParameter()
+        if mode == "regression":
+            param.type = "regression"
+            param.c_in = 1
+            param.c_out = 1
+            param.depth = 100
+            param.width = 1
+            param.dilations = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        elif mode == "segmentation":
+            param.type = "segmentation"
+            param.num_labels = None
+            param.depth = 100
+            param.width = 1
+            param.dilations = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        return param
+
+    @staticmethod
+    def cite(cite_format="MLA"):
+        if cite_format == "MLA":
+            print("Pelt, Daniël M., and James A. Sethian.")
+            print(
+                '"A mixed-scale dense convolutional neural network for image analysis."'
+            )
+            print("\x1B[3mProceedings of the National Academy of Sciences  \x1B[0m")
+            print("115.2 (2018): 254-259.")
+        elif cite_format == "bib":
+            string = """
+            @article{pelt2018mixed,
+            title={A mixed-scale dense convolutional neural network for image analysis},
+            author={Pelt, Dani{\"e}l M and Sethian, James A},
+            journal={Proceedings of the National Academy of Sciences},
+            volume={115},
+            number={2},
+            pages={254--259},
+            year={2018},
+            publisher={National Acad Sciences}
+            }
+            """
+            print(string)
+        else:
+            raise AttributeError(
+                'cite_format not understood, only "MLA" and "bib" supported'
+            )
+
+    def forward(self, x):
+        B, C, W, H = x.shape
+
+        image = x.new_zeros(B, 1, *self.geo.image_shape[1:])
+        for i in range(B):
+            aux = fdk(self.op, x[i, 0])
+            aux = torch.clip(aux, min=0)
+            image[i] = aux
+
+        return self.net(image)