Merge pull request #221 from Refound-445/master

Support For 3D Custom Dataset

Author: WuJunde

guided_diffusion/custom_dataset_loader.py

@@ -14,12 +14,14 @@
 from skimage.transform import rotate
 from glob import glob
 from sklearn.model_selection import train_test_split
+import nibabel
+
 
 class CustomDataset(Dataset):
-    def __init__(self, args, data_path , transform = None, mode = 'Training',plane = False):
+    def __init__(self, args, data_path, transform=None, mode="Training", plane=False):
 
-        print("loading data from the directory :",data_path)
-        path=data_path
+        print("loading data from the directory :", data_path)
+        path = data_path
         images = sorted(glob(os.path.join(path, "images/*.png")))
         masks = sorted(glob(os.path.join(path, "masks/*.png")))
 
@@ -37,12 +39,12 @@ def __getitem__(self, index):
         """Get the images"""
         name = self.name_list[index]
         img_path = os.path.join(name)
-        
+
         mask_name = self.label_list[index]
         msk_path = os.path.join(mask_name)
 
-        img = Image.open(img_path).convert('RGB')
-        mask = Image.open(msk_path).convert('L')
+        img = Image.open(img_path).convert("RGB")
+        mask = Image.open(msk_path).convert("L")
 
         # if self.mode == 'Training':
         #     label = 0 if self.label_list[index] == 'benign' else 1
@@ -60,3 +62,60 @@ def __getitem__(self, index):
         #     return (img, mask, name)
         # else:
         #     return (img, mask, name)
+
+
+class CustomDataset3D(torch.utils.data.Dataset):
+    def __init__(self, data_path, transform):
+        super().__init__()
+        
+        print("loading data from the directory :", data_path)
+        path = data_path
+        images = sorted(glob(os.path.join(path, "images/*.nii.gz")))
+        masks = sorted(glob(os.path.join(path, "masks/*.nii.gz")))
+
+        assert len(images) == len(masks), "Number of images and masks must be the same"
+        
+        self.valid_cases = [(img_path, seg_path) for img_path, seg_path in zip(images, masks)]
+
+        self.all_slices = []
+        for case_idx, (img_path, seg_path) in enumerate(self.valid_cases):
+            seg_vol = nibabel.load(seg_path)
+            img = nibabel.load(img_path)
+            assert (
+                img.shape == seg_vol.shape
+            ), f"Image and segmentation shape mismatch: {img.shape} vs {seg_vol.shape}, Flies: {img_path}, {seg_path}"
+            num_slices = img.shape[-1]
+            self.all_slices.extend(
+                [(case_idx, slice_idx) for slice_idx in range(num_slices)]
+            )
+            
+        self.data_path = path
+        
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.all_slices)
+
+    def __getitem__(self, x):
+        case_idx, slice_idx = self.all_slices[x]
+        img_path, seg_path = self.valid_cases[case_idx]
+
+        nib_img = nibabel.load(img_path)
+        nib_seg = nibabel.load(seg_path)
+
+        image = torch.tensor(nib_img.get_fdata(),dtype=torch.float32)[:, :, slice_idx].unsqueeze(0).unsqueeze(0)
+        label = torch.tensor(nib_seg.get_fdata(),dtype=torch.float32)[:, :, slice_idx].unsqueeze(0).unsqueeze(0)
+        label = torch.where(
+            label > 0, 1, 0
+        ).float()  # merge all tumor classes into one
+
+        if self.transform:
+            state = torch.get_rng_state()
+            image = self.transform(image)
+            torch.set_rng_state(state)
+            label = self.transform(label)
+        return (
+            image,
+            label,
+            img_path.split(".nii")[0] + "_slice" + str(slice_idx) + ".nii",
+        )  # virtual path

scripts/segmentation_train.py

@@ -7,14 +7,15 @@
 from guided_diffusion.resample import create_named_schedule_sampler
 from guided_diffusion.bratsloader import BRATSDataset, BRATSDataset3D
 from guided_diffusion.isicloader import ISICDataset
-from guided_diffusion.custom_dataset_loader import CustomDataset
+from guided_diffusion.custom_dataset_loader import CustomDataset,CustomDataset3D
 from guided_diffusion.script_util import (
     model_and_diffusion_defaults,
     create_model_and_diffusion,
     args_to_dict,
     add_dict_to_argparser,
 )
 import torch as th
+from pathlib import Path
 from guided_diffusion.train_util import TrainLoop
 from visdom import Visdom
 viz = Visdom(port=8850)
@@ -40,7 +41,13 @@ def main():
 
         ds = BRATSDataset3D(args.data_dir, transform_train, test_flag=False)
         args.in_ch = 5
-    else :
+    elif any(Path(args.data_dir).glob("*\*.nii.gz")):
+        tran_list = [transforms.Resize((args.image_size,args.image_size)),]
+        transform_train = transforms.Compose(tran_list)
+        print("Your current directory : ",args.data_dir)
+        ds = CustomDataset3D(args, args.data_dir, transform_train)
+        args.in_ch = 4
+    else:
         tran_list = [transforms.Resize((args.image_size,args.image_size)), transforms.ToTensor(),]
         transform_train = transforms.Compose(tran_list)
         print("Your current directory : ",args.data_dir)