initial commit

app.py

@@ -0,0 +1,58 @@
+import streamlit as st
+import numpy as np
+from apps.data_preprocessing import preprocess_input
+import os
+import tempfile
+import nibabel as nib
+
+st.title("NIfTI File Uploader and Video Generator")
+st.write("This tool accepts NIfTI `.nii` or `.nii.gz` files smaller than a specific size.")
+uploaded_file = st.file_uploader("Upload a NIfTI file (.nii or .nii.gz)", type=["nii", "nii.gz"])
+
+if uploaded_file is not None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        try:
+            # Save uploaded file to a temporary location
+            nifti_path = os.path.join(temp_dir, "uploaded_file.nii.gz")
+            with open(nifti_path, "wb") as f:
+                f.write(uploaded_file.read())
+
+            # Preprocess the uploaded NIfTI file
+            output_tensor, dataset = preprocess_input(nifti_path)
+            imgs = dataset[0]["CT"]["data"]
+            pred = output_tensor.argmax(0)
+            numpy_array = pred.numpy().astype(np.int16)
+            affine = np.eye(4)
+
+            # Save finalized mask
+            mask_path = os.path.join(temp_dir, "finalized_mask.nii")
+            nifti_image_mask = nib.Nifti1Image(numpy_array, affine)
+            nib.save(nifti_image_mask, mask_path)
+
+            # Save finalized image
+            nifti_array = imgs.squeeze(0).numpy()
+            image_path = os.path.join(temp_dir, "finalized_image.nii.gz")
+            nifti_image_img = nib.Nifti1Image(nifti_array, affine)
+            nib.save(nifti_image_img, image_path)
+
+            # Provide download buttons for the finalized files
+            with open(mask_path, "rb") as f:
+                st.download_button(
+                    label="Download Finalized Mask",
+                    data=f,
+                    file_name="finalized_mask.nii",
+                    mime="application/octet-stream"
+                )
+
+            with open(image_path, "rb") as f:
+                st.download_button(
+                    label="Download Finalized Image",
+                    data=f,
+                    file_name="finalized_image.nii.gz",
+                    mime="application/octet-stream"
+                )
+
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+else:
+    st.info("Please upload a NIfTI file to generate the video.")

apps/best_model_36.ckpt

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

apps/data_preprocessing.py

@@ -0,0 +1,30 @@
+import torchio as tio
+import torch
+from apps.model import model
+def preprocess_input(uploaded_file):
+    subject = tio.Subject({"CT": tio.ScalarImage(uploaded_file)})
+    normalize_orientation = tio.ToCanonical()
+    preprocess_spatial = tio.Compose([
+        normalize_orientation,
+        tio.RescaleIntensity((0, 1)),
+        tio.Resize((300, 300, 400))
+    ])
+    transform = preprocess_spatial
+    dataset = tio.SubjectsDataset([subject], transform=transform)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt_path = 'apps/best_model_36.ckpt'
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    model.load_state_dict(checkpoint['state_dict'])
+    model.to(device)
+    model.eval()
+    grid_sampler = tio.inference.GridSampler(dataset[0], 96, (8, 8, 8))
+    aggregator = tio.inference.GridAggregator(grid_sampler)
+    patch_loader = tio.data.SubjectsLoader(grid_sampler, batch_size=4)
+    with torch.no_grad():
+        for patches_batch in patch_loader:
+            input_tensor = patches_batch['CT']["data"].to(device)  # Get batch of patches
+            locations = patches_batch[tio.LOCATION]  # Get locations of patches
+            pred = model(input_tensor)  # Compute prediction
+            aggregator.add_batch(pred, locations)
+    output_tensor = aggregator.get_output_tensor()
+    return output_tensor, dataset

apps/model.py

@@ -0,0 +1,11 @@
+import pytorch_lightning as pl
+from apps.project_model2 import UNet
+
+class Segmenter(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = UNet()
+    def forward(self, data):
+        pred = self.model(data)
+        return pred
+model=Segmenter()

apps/project_model2.py

@@ -0,0 +1,91 @@
+import torch
+
+class DoubleConv(torch.nn.Module):
+    """
+    Helper Class which implements the intermediate Convolutions
+    """
+    def __init__(self, in_channels, out_channels):
+        
+        super().__init__()
+        self.step = torch.nn.Sequential(torch.nn.Conv3d(in_channels, out_channels, 3, padding=1),
+                                        torch.nn.ReLU(),
+                                        torch.nn.Conv3d(out_channels, out_channels, 3, padding=1),
+                                        torch.nn.ReLU())
+        
+    def forward(self, X):
+        return self.step(X)
+
+    
+    
+
+class UNet(torch.nn.Module):
+    """
+    This class implements a UNet for the Segmentation
+    We use 3 down- and 3 UpConvolutions and two Convolutions in each step
+    """
+
+    def __init__(self):
+        """Sets up the U-Net Structure
+        """
+        super().__init__()
+        
+        
+        ############# DOWN SAMPLING #####################
+        self.layer1 = DoubleConv(1, 32)
+        self.layer2 = DoubleConv(32, 64)
+        self.layer3 = DoubleConv(64, 128)
+        self.layer4 = DoubleConv(128, 256)
+
+        #########################################
+
+        ############## UP SAMPLING #######################
+        self.layer5 = DoubleConv(256 + 128, 128)
+        self.layer6 = DoubleConv(128+64, 64)
+        self.layer7 = DoubleConv(64+32, 32)
+        self.layer8 = torch.nn.Conv3d(32, 6, 1)  # Output: 5 values -> background, upper jaw, lower jaw,upper teeth, lower teeth, artery
+        #########################################
+
+        self.maxpool = torch.nn.MaxPool3d(2)
+
+    def forward(self, x):
+        
+        ####### DownConv 1#########
+        x1 = self.layer1(x)
+        x1m = self.maxpool(x1)
+        ###########################
+        
+        ####### DownConv 2#########        
+        x2 = self.layer2(x1m)
+        x2m = self.maxpool(x2)
+        ###########################
+
+        ####### DownConv 3#########        
+        x3 = self.layer3(x2m)
+        x3m = self.maxpool(x3)
+        ###########################
+        
+        ##### Intermediate Layer ## 
+        x4 = self.layer4(x3m)
+        ###########################
+
+        ####### UpCONV 1#########        
+        x5 = torch.nn.Upsample(scale_factor=2, mode="trilinear")(x4)  # Upsample with a factor of 2
+        x5 = torch.cat([x5, x3], dim=1)  # Skip-Connection
+        x5 = self.layer5(x5)
+        ###########################
+
+        ####### UpCONV 2#########        
+        x6 = torch.nn.Upsample(scale_factor=2, mode="trilinear")(x5)        
+        x6 = torch.cat([x6, x2], dim=1)  # Skip-Connection AKA downsampling   
+        x6 = self.layer6(x6)
+        ###########################
+        
+        ####### UpCONV 3#########        
+        x7 = torch.nn.Upsample(scale_factor=2, mode="trilinear")(x6)
+        x7 = torch.cat([x7, x1], dim=1)       
+        x7 = self.layer7(x7)
+        ###########################
+        
+        ####### Predicted segmentation#########        
+        ret = self.layer8(x7)
+        return ret

requirements.txt

@@ -0,0 +1,5 @@
+streamlit==1.41.1
+nibabel==5.3.2
+torch==2.5.1
+torchio==0.20.3
+pytorch-lightning==2.5.0.post0