Upload 6 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/paper_images.jpg filter=lfs diff=lfs merge=lfs -text
+images/paper_images.pdf filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,55 @@
----
-license: mit
----
+# Deep Colorization modules for Medical Image Analysis
+
+
+This project aims at bridging the gap between medical image analysis by introducing a light colorization module. Three different modules are proposed and implemented (DECONV, PixelShuffle and ColorU)
+
+
+![Modules](/images/MODULES.png)
+
+The modules are trained jointly with a backbone pre-trained on ImageNet. A  multi-stage transfer learning pipeline is summarized here. 
+First, the colorization module is trained from scratch together with the classifier, while the pre-trained CNN backbone is kept frozen, to learn the mapping which maximizes classification accuracy. 
+Then, the entire network is fine-tuned to learn useful features for the target task, while simultaneously adjusting the colorization mapping. The figure below shows the output of each colorization module when only the colorization module is trained, and after the entire network is fine-tuned.
+
+
+![Colorization](/images/paper_images.jpg)
+
+## Dependencies
+
++ Linux 
++ Python 3.7
++ PyTorch 1.4.0
+
+## Download 
+
+Trained models with DenseNet121 and ResNet18 backbones are available [here](https://drive.google.com/drive/folders/1uwLd-rzkt7Fcph6RqR1Eq41aTh85XGb-?usp=sharing)
+A detailed list of models is available [here](README_FILES.md)
+
+All models were trained on [CheXpert](https://stanfordmlgroup.github.io/competitions/chexpert/)  to predict the presence/absence of 5 labels:
++ Atelectasis
++ Cardiomegaly
++ Consolidation
++ Edema
++ Pleural Effusion
+
+
+## Image normalization
+If you wish to use the above models, please bear in mind that images were normalized with statistics calculated on the CheXPert dataset:
+
++ mean: [0.5028, 0.5028, 0.5028]
++ std: [0.2902, 0.2902, 0.2902]
+
+
+# Citation
+
+If you use the models in your research, please cite our paper:
+
+```
+@article{morra2020bridging,
+  title="Bridging the gap between Natural and Medical Images through Deep Colorization",
+  author="Morra, Lia and Piano, Luca and Lamberti, Fabrizio and Tommasi, Tatiana",
+  year="2020"
+  }
+```
+
+
+

colorization_modules.py

@@ -0,0 +1,469 @@
+import logging as log
+import os
+from pathlib import Path
+
+import torch
+import torchvision.transforms as transforms
+import torchvision.models as models
+from torch import nn
+from torch.nn import functional as F
+
+from enum import Enum
+
+
+class AdvEnum(Enum):
+    @classmethod
+    def list(cls):
+        return list(map(lambda c: c.value, cls))
+
+    @classmethod
+    def list_name_value(cls):
+        return list(map(lambda c: (c.name, c.value), cls))
+
+
+class DecoNetMode(AdvEnum):
+    FREEZE_DECO = 0
+    FREEZE_PTMODEL = 1
+    FREEZE_PTMODEL_NO_FC = 2
+    UNFREEZE_ALL = 3
+    FREEZE_ALL = 4
+    FREEZE_ALL_NO_FC = 5
+
+
+class DecoType(AdvEnum):
+    NO = 0
+    DECONV = 1
+    RESIZE_CONV = 2
+    ColorUDECO = 16
+    PIXEL_SHUFFLE = 20
+
+
+def get_deco_model(use_deco, out_deco) -> nn.Module:
+    if use_deco in [DecoType.DECONV, DecoType.DECONV_NORM]:
+        return StandardDECO(out_deco, deconv=True)
+    elif use_deco in [DecoType.RESIZE_CONV]:
+        return StandardDECO(out_deco, deconv=False)
+    elif use_deco is DecoType.PIXEL_SHUFFLE:
+        return PixelShuffle(out_deco, lrelu=False)
+    elif use_deco is DecoType.ColorUDECO:
+        return ColorUDECO(out_deco)
+    else:
+        raise ValueError("Module not found")
+
+
+class PreTrainedModel(AdvEnum):
+    DENSENET_121 = 0
+    RESNET_18 = 1
+    RESNET_34 = 2
+    RESNET_50 = 3
+    VGG11 = 4
+    VGG11_BN = 5
+
+
+def get_pt_model(model, output, pretrained=True):
+    input = 224
+    if not isinstance(model, PreTrainedModel):
+        model = PreTrainedModel(model)
+    pt_model = None
+    if model == PreTrainedModel.DENSENET_121:
+        pt_model = models.densenet121(pretrained=pretrained)
+        num_ftrs = pt_model.classifier.in_features
+        pt_model.classifier = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "classifier"
+    elif model == PreTrainedModel.RESNET_18:
+        pt_model = models.resnet18(pretrained=pretrained)
+        num_ftrs = pt_model.fc.in_features
+        pt_model.fc = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "fc"
+    elif model == PreTrainedModel.RESNET_34:
+        pt_model = models.resnet34(pretrained=pretrained)
+        num_ftrs = pt_model.fc.in_features
+        pt_model.fc = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "fc"
+    elif model == PreTrainedModel.RESNET_50:
+        pt_model = models.resnet50(pretrained=pretrained)
+        num_ftrs = pt_model.fc.in_features
+        pt_model.fc = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "fc"
+    elif model == PreTrainedModel.VGG11:
+        pt_model = models.vgg11(pretrained=pretrained)
+        num_ftrs = pt_model.classifier[6].in_features
+        pt_model.classifier[6] = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "classifier.6"
+    elif model == PreTrainedModel.VGG11_BN:
+        pt_model = models.vgg11_bn(pretrained=pretrained)
+        num_ftrs = pt_model.classifier[6].in_features
+        pt_model.classifier[6] = nn.Linear(num_ftrs, output)
+        pt_model.last_layer_name = "classifier.6"
+    else:
+        raise ValueError("Model not found")
+
+    return pt_model, input
+
+
+class DecoNet(nn.Module):
+    """
+    Colorization module(optional)+Model
+    """
+
+    def __init__(self, output=14,
+                 deco_type=DecoType.ColorUDECO,
+                 pt_model=PreTrainedModel.RESNET_18,
+                 pre_trained=True,
+                 training_mode=DecoNetMode.FREEZE_PTMODEL_NO_FC,
+                 use_aap=False):
+        super().__init__()
+        # Pre-trained Model
+        self.deco_type = deco_type
+        self.training_mode = training_mode
+        self.use_aap = use_aap
+        pt_model, self.out_deco = get_pt_model(pt_model, output, pre_trained)
+        self.last_layer_name = pt_model.last_layer_name
+        # DECO if needed
+        if self.deco_type is not DecoType.NO:
+            self.deco = get_deco_model(self.deco_type, self.out_deco)
+        else:
+            self.deco = None
+        self.pt_model = pt_model
+        self.set_mode(training_mode)
+
+    def set_mode(self, mode, print=True):
+        if not isinstance(mode, DecoNetMode):
+            mode = DecoNetMode(mode)
+        if mode == DecoNetMode.UNFREEZE_ALL:
+            for param in self.parameters():
+                param.requires_grad = True
+        elif mode == DecoNetMode.FREEZE_DECO:
+            self.set_mode(DecoNetMode.UNFREEZE_ALL, False)
+            for param in self.deco.parameters():
+                param.requires_grad = False
+        elif mode == DecoNetMode.FREEZE_PTMODEL:
+            self.set_mode(DecoNetMode.UNFREEZE_ALL, False)
+            for param in self.pt_model.parameters():
+                param.requires_grad = False
+        elif mode == DecoNetMode.FREEZE_PTMODEL_NO_FC:
+            self.set_mode(DecoNetMode.UNFREEZE_ALL, False)
+            for name, param in self.pt_model.named_parameters():
+                if self.last_layer_name not in name:
+                    param.requires_grad = False
+        elif mode == DecoNetMode.FREEZE_ALL:
+            for param in self.parameters():
+                param.requires_grad = False
+        elif mode == DecoNetMode.FREEZE_ALL_NO_FC:
+            self.set_mode(DecoNetMode.FREEZE_ALL, False)
+            # Unfreeze last layer
+            for name, param in self.pt_model.named_parameters():
+                if self.last_layer_name in name:
+                    param.requires_grad = True
+
+        if print:
+            log.info("#############################################")
+            log.info("PARAMETERS STATUS:")
+            for name, param in self.named_parameters():
+                log.info("{} : {}".format(name, param.requires_grad))
+            log.info("#############################################")
+
+    def get_layer_weight(self, sel_name: str = ""):
+        if sel_name == "":
+            sel_name = self.last_layer_name
+        res = []
+        for name, param in self.pt_model.named_parameters():
+            if sel_name in name:
+                res.append(param)
+
+        return res
+
+    def forward(self, xb):
+        """
+        @:param xb : tensor
+          Batch of input images
+
+        @:return tensor
+          A batch of output images
+        """
+        if self.deco is not None:
+            xb = self.deco(xb)
+            if self.use_aap:
+                xb = F.adaptive_avg_pool2d(xb, (self.out_deco, self.out_deco))
+        return self.pt_model(xb)
+
+    def clean_last_layer(self):
+        pt_model_type = self.pt_model
+
+        if pt_model_type == PreTrainedModel.VGG11_BN or pt_model_type == PreTrainedModel.VGG11:
+            self.pt_model.classifier[6].reset_parameters()
+        else:
+            last_layer_name = list(self.pt_model._modules)[-1]
+            self.pt_model._modules[last_layer_name].reset_parameters()
+
+        log.info("Last layer cleaned!")
+
+    def last_layer_size(self):
+        pt_model_type = self.pt_model
+        if pt_model_type == PreTrainedModel.VGG11_BN or pt_model_type == PreTrainedModel.VGG11:
+            return self.pt_model.classifier[6].weight.shape[-1]
+        else:
+            last_layer_name = list(self.pt_model._modules)[-1]
+            return self.pt_model._modules[last_layer_name].shape[-1]
+
+    def load_deco_state_dict(self, state_dict):
+        if self.deco is None:
+            self.deco = get_deco_model(self.deco_type, self.out_deco)
+        if hasattr(self.deco, "load_state_dict"):
+            self.deco.load_state_dict(state_dict)
+        else:
+            return False
+        self.set_mode(self.training_mode)
+        return True
+
+
+def default_deco__weight_init(m):
+    if isinstance(m, nn.Conv2d):
+        # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+        # m.weight.data.normal_(0, math.sqrt(2. / n))
+        torch.nn.init.xavier_uniform_(m.weight)
+    elif isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1)
+        m.bias.data.zero_()
+
+
+def bn_weight_init(m):
+    if isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1)
+        m.bias.data.zero_()
+
+
+class BaseDECO(nn.Module):
+    def __init__(self, out=224, init=None):
+        super().__init__()
+        self.out_s = out
+        self.init = init
+
+    def set_output_size(self, out_s):
+        self.out_s = out_s
+
+    def init_weights(self):
+        if self.init is None:
+            pass
+        elif self.init == 0:
+            self.apply(default_deco__weight_init)
+        elif self.init == 1:
+            self.apply(bn_weight_init)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, ni, nf=None, kernel=3, stride=1, padding=1):
+        super().__init__()
+        if nf is None:
+            nf = ni
+        self.conv1 = conv_layer(ni, nf, kernel=kernel, stride=stride, padding=padding)
+        self.conv2 = conv_layer(nf, nf, kernel=kernel, stride=stride, padding=padding)
+
+    def forward(self, x):
+        return x + self.conv2(self.conv1(x))
+
+
+def conv_layer(in_layer, out_layer, kernel=3, stride=1, padding=1, instanceNorm=False):
+    return nn.Sequential(
+        nn.Conv2d(in_layer, out_layer, kernel_size=kernel, stride=stride, padding=padding),
+        nn.BatchNorm2d(out_layer) if not instanceNorm else nn.InstanceNorm2d(out_layer),
+        nn.LeakyReLU(inplace=True)
+    )
+
+
+def _make_res_layers(nl, ni, kernel=3, stride=1, padding=1):
+    layers = []
+    for i in range(nl):
+        layers.append(ResBlock(ni, kernel=kernel, stride=stride, padding=padding))
+
+    return nn.Sequential(*layers)
+
+
+class StandardDECO(BaseDECO):
+    """
+    Standard DECO Module
+    """
+
+    def __init__(self, out=224, init=0, deconv=False):
+        super().__init__(out, init)
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=2)
+        self.bn1 = nn.BatchNorm2d(64)
+        # ReLU
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.resblocks = _make_res_layers(8, 64)
+        self.conv_last = nn.Conv2d(64, 3, kernel_size=1)
+        self.deconv = deconv
+        if deconv:
+            # TODO: Check if use "groups = 1"
+            self.deconv = nn.ConvTranspose2d(in_channels=3, out_channels=3, kernel_size=8, padding=2, stride=4,
+                                             groups=3, bias=False)
+        else:
+            self.pad = nn.ReflectionPad2d(1)
+            self.conv_up = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=0, stride=1)
+
+        self.init_weights()
+
+    def forward(self, xb):
+        """
+        @:param xb : Tensor
+          Batch of input images
+
+        @:return tensor
+          A batch of output images
+        """
+        _xb = self.maxpool(F.leaky_relu(self.bn1(self.conv1(xb))))
+        _xb = self.resblocks(_xb)
+        _xb = self.conv_last(_xb)
+        if self.deconv:
+            _xb = self.deconv(_xb, output_size=xb.shape)
+        else:
+            _xb = self.conv_up(self.pad(F.interpolate(_xb, scale_factor=4, mode='nearest')))
+        return _xb
+
+
+def icnr(x, scale=4, init=nn.init.kaiming_normal_):
+    """ ICNR init of `x`, with `scale` and `init` function.
+
+        Checkerboard artifact free sub-pixel convolution: https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf
+    """
+    ni, nf, h, w = x.shape
+    ni2 = int(ni / (scale ** 2))
+    k = init(torch.zeros([ni2, nf, h, w])).transpose(0, 1)
+    k = k.contiguous().view(ni2, nf, -1)
+    k = k.repeat(1, 1, scale ** 2)
+    k = k.contiguous().view([nf, ni, h, w]).transpose(0, 1)
+    x.data.copy_(k)
+
+
+class PixelShuffle_ICNR(nn.Module):
+    """ Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`, `icnr` init,
+        and `weight_norm`.
+
+        "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts":
+        https://arxiv.org/abs/1806.02658
+    """
+
+    def __init__(self, ni: int, nf: int = None, scale: int = 4, icnr_init=True, blur_k=2, blur_s=1,
+                 blur_pad=(1, 0, 1, 0), lrelu=True):
+        super().__init__()
+        nf = ni if nf is None else nf
+        self.conv = conv_layer(ni, nf * (scale ** 2), kernel=1, padding=0, stride=1) if lrelu else nn.Sequential(
+            nn.Conv2d(64, 3 * (scale ** 2), 1, 1, 0), nn.BatchNorm2d(3 * (scale ** 2)))
+        if icnr_init:
+            icnr(self.conv[0].weight, scale=scale)
+        self.act = nn.LeakyReLU(inplace=False) if lrelu else nn.Hardtanh(-10000, 10000)
+        self.shuf = nn.PixelShuffle(scale)
+        # Blurring over (h*w) kernel
+        self.pad = nn.ReplicationPad2d(blur_pad)
+        self.blur = nn.AvgPool2d(blur_k, stride=blur_s)
+
+    def forward(self, x):
+        x = self.shuf(self.act(self.conv(x)))
+        return self.blur(self.pad(x))
+
+
+class PixelShuffle(BaseDECO):
+    """
+    PixelShuffle Module
+    """
+
+    def __init__(self, out=224, init=1, scale=4, lrelu=False):
+        super().__init__(out, init)
+        # Which value should I use for stride and padding?
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=2)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.act1 = nn.LeakyReLU()
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.resblocks = _make_res_layers(8, 64)
+        self.pixel_shuffle = PixelShuffle_ICNR(ni=64, nf=3, scale=scale, lrelu=lrelu)
+        self.init_weights()
+
+    def forward(self, xb):
+        """
+        @:param xb : Tensor
+          Batch of input images
+
+        @:return tensor
+          A batch of output images
+        """
+        _xb = self.maxpool(self.act1(self.bn1(self.conv1(xb))))
+        _xb = self.resblocks(_xb)
+
+        return self.pixel_shuffle(_xb)
+
+
+class ColorUDECO(BaseDECO):
+    """
+    ColorUDECO Module
+    """
+
+    def __init__(self, out=224, init=0, in_ch=1, out_ch=3):
+        super().__init__(out, init)
+        self.dw1 = ColorDown(in_ch, 16)
+        self.dw2 = ColorDown(16, 32)
+        self.dw3 = ColorDown(32, 64)
+        self.up1 = ColorUp(64, 32)
+        self.up2 = ColorUp(64, 16)
+        self.out = ColorOut(32, 16, out_ch)
+
+    def forward(self, x1):
+        """
+        @:param x1 : Tensor
+          Batch of input images
+
+        @:return tensor
+          A batch of output images
+        """
+        x1 = self.dw1(x1)
+        x2 = self.dw2(x1)
+        x3 = self.dw3(x2)
+        x3 = self.up1(x3)
+        x2 = self.up2(torch.cat([x2, x3], dim=1))
+        return self.out(torch.cat([x1, x2], dim=1))
+
+
+class ColorDown(nn.Module):
+    def __init__(self, in_ch, out_ch, htan=False):
+        super(ColorDown, self).__init__()
+        self.d = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.Conv2d(out_ch, out_ch, kernel_size=4, stride=2, padding=1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.BatchNorm2d(out_ch)
+        )
+
+    def forward(self, x):
+        return self.d(x)
+
+
+class ColorUp(nn.Module):
+    def __init__(self, in_ch, out_ch, htan=False):
+        super(ColorUp, self).__init__()
+        self.u = nn.Sequential(
+            nn.ConvTranspose2d(in_ch, out_ch, 4, 2, 1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.BatchNorm2d(out_ch)
+        )
+
+    def forward(self, x):
+        return self.u(x)
+
+
+class ColorOut(nn.Module):
+    def __init__(self, in_ch, out_ch, out_last, htan=False):
+        super(ColorOut, self).__init__()
+        self.u = nn.Sequential(
+            nn.ConvTranspose2d(in_ch, out_ch, 4, 2, 1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU() if not htan else nn.Hardtanh(),
+            nn.Conv2d(out_ch, out_last, kernel_size=1, stride=1, padding=0),
+        )
+
+    def forward(self, x):
+        return self.u(x)

images/empty

@@ -0,0 +1 @@
+

images/paper_images.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d46d8be627e5d96c592be306e9c74b91a2412cbbc05bead81e3a12082915b89
+size 1271680