Upload folder using huggingface_hub

base_model.py

@@ -0,0 +1,69 @@
+import os
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+from torchsummary import summary
+
+class BaseModel(nn.Module, ABC):
+  def __init__(self):
+    super().__init__()
+    self.best_loss = 1000000
+
+  @abstractmethod
+  def forward(self, x):
+    pass
+
+  @abstractmethod
+  def test(self):
+    pass
+
+  @property
+  def device(self):
+    return next(self.parameters()).device
+
+  def restore_checkpoint(self, ckpt_file, optimizer=None, affect_weights=True):
+    """
+    Restores checkpoint from a pth file and restores optimizer state.
+    Args:
+      ckpt_file (str): A PyTorch pth file containing model weights.
+      optimizer (Optimizer): A vanilla optimizer to have its state restored from.
+    Returns:
+      int: Global step variable where the model was last checkpointed.
+    """
+    if not ckpt_file:
+      raise ValueError("No checkpoint file to be restored.")
+
+    try:
+      ckpt_dict = torch.load(ckpt_file)
+
+    except RuntimeError:
+      ckpt_dict = torch.load(ckpt_file, map_location=lambda storage, loc: storage)
+    # Restore model weights if needed
+    if affect_weights:
+      self.load_state_dict(ckpt_dict['model_state_dict'])
+    # Restore optimizer status if existing. Evaluation doesn't need this
+    if optimizer:
+      optimizer.load_state_dict(ckpt_dict['optimizer_state_dict'])
+    # Return global step
+    return ckpt_dict, optimizer
+
+  def count_params(self):
+    """
+    Computes the number of parameters in this model.
+    Args: None
+    Returns:
+        int: Total number of weight parameters for this model.
+        int: Total number of trainable parameters for this model.
+    """
+    num_total_params = sum(p.numel() for p in self.parameters())
+    num_trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+    return num_total_params, num_trainable_params
+
+  def inference(self, input_tensor):
+    self.eval()
+    with torch.no_grad():
+      output = self.forward(input_tensor)
+      if isinstance(output, tuple):
+        output = output[0]
+      return output.cpu().detach()

inference.py

@@ -0,0 +1,75 @@
+import torch 
+import tempfile
+import numpy as np
+import nibabel as nib
+from fastapi import FastAPI, UploadFile, File
+
+def preprocess_input(flair_file: UploadFile, t1_file: UploadFile, t1ce_file: UploadFile, t2_file: UploadFile):
+    with tempfile.NamedTemporaryFile(suffix=".nii") as temp_flair, \
+            tempfile.NamedTemporaryFile(suffix=".nii") as temp_t1, \
+            tempfile.NamedTemporaryFile(suffix=".nii") as temp_t1ce, \
+            tempfile.NamedTemporaryFile(suffix=".nii") as temp_t2:
+        
+        # Save uploaded files to temporary NIfTI files
+        flair_content = flair_file.file.read()
+        t1_content = t1_file.file.read()
+        t1ce_content = t1ce_file.file.read()
+        t2_content = t2_file.file.read()
+        
+        temp_flair.write(flair_content)
+        temp_t1.write(t1_content)
+        temp_t1ce.write(t1ce_content)
+        temp_t2.write(t2_content)
+
+        # Load and preprocess the NIfTI files
+        flair = nib.load(temp_flair.name).get_fdata()
+        t1 = nib.load(temp_t1.name).get_fdata()
+        t1ce = nib.load(temp_t1ce.name).get_fdata()
+        t2 = nib.load(temp_t2.name).get_fdata()
+
+        flair_tensor = torch.tensor(flair[56:184, 56:184, 13:141], dtype=torch.float32).unsqueeze(0)
+        t1_tensor = torch.tensor(t1[56:184, 56:184, 13:141], dtype=torch.float32).unsqueeze(0)
+        t1ce_tensor = torch.tensor(t1ce[56:184, 56:184, 13:141], dtype=torch.float32).unsqueeze(0)
+        t2_tensor = torch.tensor(t2[56:184, 56:184, 13:141], dtype=torch.float32).unsqueeze(0)
+
+    return flair_tensor, t1_tensor, t1ce_tensor, t2_tensor
+
+
+def segment_wt(wt_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor):
+    wt_input = torch.cat((flair_tensor, t1_tensor, t1ce_tensor, t2_tensor), dim=0).unsqueeze(0)
+    with torch.no_grad():
+        wt_output = wt_model(wt_input)
+    return wt_output
+
+
+def segment_tc(tc_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output):
+    tc_input = torch.cat((flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output.squeeze(0), 
+                        wt_output.squeeze(0), wt_output.squeeze(0), wt_output.squeeze(0)), dim=0).unsqueeze(0)
+    with torch.no_grad():
+        tc_output = tc_model(tc_input)
+    return tc_output
+
+def segment_et(et_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output, tc_output):
+    et_input = torch.cat((flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output.squeeze(0), 
+                        wt_output.squeeze(0), tc_output.squeeze(0), tc_output.squeeze(0)), dim=0).unsqueeze(0)
+    with torch.no_grad():
+        et_output = et_model(et_input)
+    return et_output
+
+def segment_all(wt_model, tc_model, et_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor):
+    wt_output = segment_wt(wt_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor)
+    tc_output = segment_tc(tc_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output)
+    et_output = segment_et(et_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor, wt_output, tc_output)
+    
+    wt_label = torch.sigmoid(wt_output.squeeze(0, 1)) > 0.5
+    tc_label = torch.sigmoid(tc_output.squeeze(0, 1)) > 0.5
+    et_label = torch.sigmoid(et_output.squeeze(0, 1)) > 0.5
+
+    output = np.zeros((128, 128, 128))
+    output[(tc_label == 1) & (wt_label == 0)] = 2
+    output[(et_label == 1) & (tc_label == 0)] = 3
+    output[(et_label == 0) & (tc_label == 0)] = 1
+
+    output_in_original_size = np.zeros((240, 240, 155))
+    output_in_original_size[56:184, 56:184, 13:141] = output
+    return output_in_original_size

model_loader.py

@@ -0,0 +1,21 @@
+import torch
+from vnet_light_arch import VNetLight
+
+def load_light_model(checkpoint_path, channels):
+  print(f"Loading The Saved Model at {checkpoint_path}")
+  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+  model = VNetLight(in_channels=channels, classes=1)
+  model.to(device)
+  model.restore_checkpoint(checkpoint_path)
+  for param in model.parameters():
+    param.requires_grad = False
+  return model
+
+def load_vChain_model(wt_chkpt, tc_chkpt, et_chkpt):
+  wt_model = load_light_model(wt_chkpt, 4)
+  tc_model = load_light_model(tc_chkpt, 8)
+  et_model = load_light_model(et_chkpt, 8)
+  return wt_model, tc_model, et_model
+
+# see this for modifying loss function:
+# https://docs.monai.io/en/latest/losses.html

vnet_blocks.py

@@ -0,0 +1,110 @@
+import os
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+from torchsummary import summary
+
+from base_model import BaseModel
+
+def passthrough(x, **kwargs):
+  return x
+
+def ELUCons(elu, nchan):
+  if elu:
+    return nn.ELU(inplace=True)
+  else:
+    return nn.PReLU(nchan)
+
+class LUConv(nn.Module):
+  def __init__(self, nchan, elu):
+    super(LUConv, self).__init__()
+    self.relu1 = ELUCons(elu, nchan)
+    self.conv1 = nn.Conv3d(nchan, nchan, kernel_size=5, padding=2)
+    self.bn1 = torch.nn.BatchNorm3d(nchan)
+
+  def forward(self, x):
+    out = self.relu1(self.bn1(self.conv1(x)))
+    return out
+
+def _make_nConv(nchan, depth, elu):
+    layers = []
+    for _ in range(depth):
+        layers.append(LUConv(nchan, elu))
+    return nn.Sequential(*layers)
+
+class InputTransition(nn.Module):
+    def __init__(self, in_channels, elu):
+        super(InputTransition, self).__init__()
+        self.num_features = 16
+        self.in_channels = in_channels
+        self.conv1 = nn.Conv3d(self.in_channels, self.num_features, kernel_size=5, padding=2)
+        self.bn1 = torch.nn.BatchNorm3d(self.num_features)
+        self.relu1 = ELUCons(elu, self.num_features)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        repeat_rate = int(self.num_features / self.in_channels)
+        out = self.bn1(out)
+        x16 = x.repeat(1, repeat_rate, 1, 1, 1)
+        return self.relu1(torch.add(out, x16))
+
+class DownTransition(nn.Module):
+  def __init__(self, inChans, nConvs, elu, dropout=False):
+    super(DownTransition, self).__init__()
+    outChans = 2 * inChans
+    self.down_conv = nn.Conv3d(inChans, outChans, kernel_size=2, stride=2)
+    self.bn1 = torch.nn.BatchNorm3d(outChans)
+
+    self.do1 = passthrough
+    self.relu1 = ELUCons(elu, outChans)
+    self.relu2 = ELUCons(elu, outChans)
+    if dropout:
+        self.do1 = nn.Dropout3d()
+    self.ops = _make_nConv(outChans, nConvs, elu)
+
+  def forward(self, x):
+    down = self.relu1(self.bn1(self.down_conv(x)))
+    out = self.do1(down)
+    out = self.ops(out)
+    out = self.relu2(torch.add(out, down))
+    return out
+
+class UpTransition(nn.Module):
+  def __init__(self, inChans, outChans, nConvs, elu, dropout=False):
+    super(UpTransition, self).__init__()
+    self.up_conv = nn.ConvTranspose3d(inChans, outChans // 2, kernel_size=2, stride=2)
+
+    self.bn1 = torch.nn.BatchNorm3d(outChans // 2)
+    self.do1 = passthrough
+    self.do2 = nn.Dropout3d()
+    self.relu1 = ELUCons(elu, outChans // 2)
+    self.relu2 = ELUCons(elu, outChans)
+    if dropout:
+        self.do1 = nn.Dropout3d()
+    self.ops = _make_nConv(outChans, nConvs, elu)
+
+  def forward(self, x, skipx):
+    out = self.do1(x)
+    skipxdo = self.do2(skipx)
+    out = self.relu1(self.bn1(self.up_conv(out)))
+    xcat = torch.cat((out, skipxdo), 1)
+    out = self.ops(xcat)
+    out = self.relu2(torch.add(out, xcat))
+    return out
+
+class OutputTransition(nn.Module):
+  def __init__(self, in_channels, classes, elu):
+    super(OutputTransition, self).__init__()
+    self.classes = classes
+    self.conv1 = nn.Conv3d(in_channels, classes, kernel_size=5, padding=2)
+    self.bn1 = torch.nn.BatchNorm3d(classes)
+
+    self.conv2 = nn.Conv3d(classes, classes, kernel_size=1)
+    self.relu1 = ELUCons(elu, classes)
+
+  def forward(self, x):
+    # convolve 32 down to channels as the desired classes
+    out = self.relu1(self.bn1(self.conv1(x)))
+    out = self.conv2(out)
+    return out

vnet_light_arch.py

@@ -0,0 +1,38 @@
+from vnet_blocks import *
+
+class VNetLight(BaseModel):
+    """
+    A lighter version of Vnet that skips down_tr256 and up_tr256 in oreder to reduce time and space complexity
+    """
+    def __init__(self, elu=True, in_channels=1, classes=4):
+        super(VNetLight, self).__init__()
+        self.classes = classes
+        self.in_channels = in_channels
+
+        self.in_tr = InputTransition(in_channels, elu)
+        self.down_tr32 = DownTransition(16, 1, elu)
+        self.down_tr64 = DownTransition(32, 2, elu)
+        self.down_tr128 = DownTransition(64, 3, elu, dropout=True)
+        self.up_tr128 = UpTransition(128, 128, 2, elu, dropout=True)
+        self.up_tr64 = UpTransition(128, 64, 1, elu)
+        self.up_tr32 = UpTransition(64, 32, 1, elu)
+        self.out_tr = OutputTransition(32, classes, elu)
+
+    def forward(self, x):
+        out16 = self.in_tr(x)
+        out32 = self.down_tr32(out16)
+        out64 = self.down_tr64(out32)
+        out128 = self.down_tr128(out64)
+        out = self.up_tr128(out128, out64)
+        out = self.up_tr64(out, out32)
+        out = self.up_tr32(out, out16)
+        out = self.out_tr(out)
+        return out
+
+    def test(self,device='cpu'):
+        input_tensor = torch.rand(1, self.in_channels, 32, 32, 32)
+        ideal_out = torch.rand(1, self.classes, 32, 32, 32)
+        out = self.forward(input_tensor)
+        assert ideal_out.shape == out.shape
+        summary(self.to(torch.device(device)), (self.in_channels, 32, 32, 32),device=device)
+        print("Vnet Light test is complete")