Initial working pipeline

lungtumormask/dataprocessing.py

@@ -3,6 +3,7 @@ from lungmask import mask
 from monai import transforms
 from monai.transforms.intensity.array import ThresholdIntensity
 from monai.transforms.spatial.array import Resize, Spacing
+from monai.transforms.utility.dictionary import ToTensord
 import torch
 from tqdm import tqdm
 import numpy as np
@@ -113,8 +114,9 @@ def process_lung_scan(scan_dict, extremes):
     transformer_1 = Compose(
         [
             DivisiblePadd(keys=["image"], k=16, mode='constant'),
-            SqueezeDimd(keys=["image"], dim = 0),
-            ToNumpyd(keys=["image"]),
+            ToTensord(keys=['image'])
+            #SqueezeDimd(keys=["image"], dim = 0),
+            #ToNumpyd(keys=["image"]),
         ]
     )
 
@@ -151,8 +153,8 @@ def preprocess(image_path):
     
     preprocess_dump['affine'] = left_lung_processed[1]
 
-    preprocess_dump['right_lung'] = right_lung_processed[0]
-    preprocess_dump['left_lung'] = left_lung_processed[0]
+    preprocess_dump['right_lung'] = right_lung_processed[0].unsqueeze(0)
+    preprocess_dump['left_lung'] = left_lung_processed[0].unsqueeze(0)
 
     return preprocess_dump
 
@@ -223,8 +225,19 @@ def stitch(org_shape, cropped, roi):
     return holder
 
 def post_process(left_mask, right_mask, preprocess_dump):
-    left = remove_pad(left_mask, preprocess_dump['left_lung'])
-    right = remove_pad(right_mask, preprocess_dump['right_lung'])
+
+    left_mask[left_mask >= 0.5] = 1
+    left_mask[left_mask < 0.5] = 0
+
+    left_mask = left_mask.astype(int)
+
+    right_mask[right_mask >= 0.5] = 1
+    right_mask[right_mask < 0.5] = 0
+
+    right_mask = right_mask.astype(int)
+
+    left = remove_pad(left_mask, preprocess_dump['left_lung'].squeeze(0).squeeze(0).numpy())
+    right = remove_pad(right_mask, preprocess_dump['right_lung'].squeeze(0).squeeze(0).numpy())
 
     left = voxel_space(left, preprocess_dump['left_extremes'])
     right = voxel_space(right, preprocess_dump['right_extremes'])

lungtumormask/mask.py

@@ -0,0 +1,34 @@
+from numpy import load
+from lungtumormask.dataprocessing import preprocess, post_process
+from lungtumormask.network import UNet_double
+import torch as T
+import nibabel
+
+def load_model():
+    model = UNet_double(3, 1, 1, tuple([64, 128, 256, 512, 1024]), tuple([2 for i in range(4)]), num_res_units = 0)
+    state_dict = T.hub.load_state_dict_from_url("https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0/dc_student.pth", progress=True, map_location=T.device('cuda:0'))
+    #model.load_state_dict(T.load("D:\\OneDrive\\Skole\\Universitet\\10. Semester\\Masteroppgave\\bruk_for_full_model.pth", map_location="cuda:0"))
+    model.load_state_dict(state_dict) 
+    model.eval()
+    return model
+
+def mask(image_path, save_path):
+    print("Loading model...")
+    model = load_model()
+
+    print("Preprocessing image...")
+    preprocess_dump = preprocess(image_path)
+
+    print("Looking for tumors...")
+    left = model(preprocess_dump['left_lung']).squeeze(0).squeeze(0).detach().numpy()
+    right = model(preprocess_dump['right_lung']).squeeze(0).squeeze(0).detach().numpy()
+
+    print("Post-processing image...")
+    infered = post_process(left, right, preprocess_dump)
+
+    print(f"Storing segmentation at {save_path}")
+    nimage = nibabel.Nifti1Image(infered, preprocess_dump['org_affine'])
+    nibabel.save(nimage, save_path)
+
+if __name__ == "__main__":
+    mask("D:\\Datasets\\MSD\\Images\\lung_003.nii.gz", "D:\\Datasets\\MSD\\Images\\out3.nii.gz")

lungtumormask/network.py

@@ -0,0 +1,151 @@
+from typing import Sequence, Union
+
+import torch
+import torch.nn as nn
+
+from monai.networks.blocks.convolutions import Convolution, ResidualUnit
+from monai.networks.layers.factories import Act, Norm
+from monai.networks.layers.simplelayers import SkipConnection
+from monai.utils import alias, export
+
+class UNet_double(nn.Module):
+    def __init__(
+        self,
+        dimensions: int,
+        in_channels: int,
+        out_channels: int,
+        channels: Sequence[int],
+        strides: Sequence[int],
+        kernel_size: Union[Sequence[int], int] = 3,
+        up_kernel_size: Union[Sequence[int], int] = 3,
+        num_res_units: int = 0,
+        act=Act.PRELU,
+        norm=Norm.INSTANCE,
+        dropout=0.0,) -> None:
+
+        super().__init__()
+
+        self.dimensions = dimensions
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.channels = channels
+        self.strides = strides
+        self.kernel_size = kernel_size
+        self.up_kernel_size = up_kernel_size
+        self.num_res_units = num_res_units
+        self.act = act
+        self.norm = norm
+        self.dropout = dropout
+
+        def _create_block(
+            inc: int, outc: int, channels: Sequence[int], strides: Sequence[int], is_top: bool) -> nn.Sequential:
+            
+            c = channels[0]
+            s = strides[0]
+
+            subblock: nn.Module
+
+            if len(channels) > 2:
+                subblock1, subblock2 = _create_block(c, c, channels[1:], strides[1:], False)  # continue recursion down
+                upc = c * 2
+            else:
+                # the next layer is the bottom so stop recursion, create the bottom layer as the sublock for this layer
+                subblock = self._get_bottom_layer(c, channels[1])
+                upc = c + channels[1]
+
+                down = self._get_down_layer(inc, c, s, is_top)  # create layer in downsampling path
+                up1 = self._get_up_layer(upc, outc, s, is_top)  # create layer in upsampling path
+                up2 = self._get_up_layer(upc, outc, s, is_top)
+
+                return nn.Sequential(down, SkipConnection(subblock), up1), nn.Sequential(down, SkipConnection(subblock), up2)
+
+            down = self._get_down_layer(inc, c, s, is_top)  # create layer in downsampling path
+            up1 = self._get_up_layer(upc, outc, s, is_top)  # create layer in upsampling path
+            up2 = self._get_up_layer(upc, outc, s, is_top)
+
+            return nn.Sequential(down, SkipConnection(subblock1), up1), nn.Sequential(down, SkipConnection(subblock2), up2)
+
+        self.model1, self.model2 = _create_block(in_channels, out_channels, self.channels, self.strides, True)
+        self.activation = nn.Sigmoid()
+
+    def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_top: bool) -> nn.Module:
+        """
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            strides: convolution stride.
+            is_top: True if this is the top block.
+        """
+        if self.num_res_units > 0:
+            return ResidualUnit(
+                self.dimensions,
+                in_channels,
+                out_channels,
+                strides=strides,
+                kernel_size=self.kernel_size,
+                subunits=self.num_res_units,
+                act=self.act,
+                norm=self.norm,
+                dropout=self.dropout,
+            )
+        return Convolution(
+            self.dimensions,
+            in_channels,
+            out_channels,
+            strides=strides,
+            kernel_size=self.kernel_size,
+            act=self.act,
+            norm=self.norm,
+            dropout=self.dropout,
+        )
+
+    def _get_bottom_layer(self, in_channels: int, out_channels: int) -> nn.Module:
+        """
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+        """
+        return self._get_down_layer(in_channels, out_channels, 1, False)
+
+    def _get_up_layer(self, in_channels: int, out_channels: int, strides: int, is_top: bool) -> nn.Module:
+        """
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            strides: convolution stride.
+            is_top: True if this is the top block.
+        """
+        conv: Union[Convolution, nn.Sequential]
+
+        conv = Convolution(
+            self.dimensions,
+            in_channels,
+            out_channels,
+            strides=strides,
+            kernel_size=self.up_kernel_size,
+            act=self.act,
+            norm=self.norm,
+            dropout=self.dropout,
+            conv_only=is_top and self.num_res_units == 0,
+            is_transposed=True,
+        )
+
+        if self.num_res_units > 0:
+            ru = ResidualUnit(
+                self.dimensions,
+                out_channels,
+                out_channels,
+                strides=1,
+                kernel_size=self.kernel_size,
+                subunits=1,
+                act=self.act,
+                norm=self.norm,
+                dropout=self.dropout,
+                last_conv_only=is_top,
+            )
+            conv = nn.Sequential(conv, ru)
+
+        return conv
+
+    def forward(self, x, box = None) -> torch.Tensor:
+        return self.activation(self.model1(x))

setup.py

@@ -0,0 +1,11 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="LungTumorMask",
+    packages=find_packages(),
+    version='1.0.1',
+    author="Svein Ole M Sevle, Vemund Fredriksen",
+    url="https://github.com/VemundFredriksen/LungTumorMask",
+    license="MIT",
+    python_requires='>=3.6'
+)