Upload 2 files

dataload.py

@@ -0,0 +1,235 @@
+import json
+from typing import Optional, Sequence
+
+import numpy as np
+import torch
+import torch.distributed as ptdist
+from monai.data import (
+    CacheDataset,
+    PersistentDataset,
+    partition_dataset,
+)
+from monai.data.utils import pad_list_data_collate
+from monai.transforms import (
+    Compose,
+    CropForegroundd,
+    EnsureChannelFirstd,
+    LoadImaged,
+    Orientationd,
+    RandSpatialCropSamplesd,
+    ScaleIntensityRanged,
+    Spacingd,
+    SpatialPadd,
+    ToTensord,
+    Transform,
+)
+
+
+class PermuteImage(Transform):
+    """Permute the dimensions of the image"""
+
+    def __call__(self, data):
+        data["image"] = data["image"].permute(
+            3, 0, 1, 2
+        )  # Adjust permutation order as needed
+        return data
+
+
+class CTDataset:
+    def __init__(
+        self,
+        json_path: str,
+        img_size: int,
+        depth: int,
+        mask_patch_size: int,
+        patch_size: int,
+        downsample_ratio: Sequence[float],
+        cache_dir: str,
+        batch_size: int = 1,
+        val_batch_size: int = 1,
+        num_workers: int = 4,
+        cache_num: int = 0,
+        cache_rate: float = 0.0,
+        dist: bool = False,
+    ):
+        super().__init__()
+        self.json_path = json_path
+        self.img_size = img_size
+        self.depth = depth
+        self.mask_patch_size = mask_patch_size
+        self.patch_size = patch_size
+        self.cache_dir = cache_dir
+        self.downsample_ratio = downsample_ratio
+        self.batch_size = batch_size
+        self.val_batch_size = val_batch_size
+        self.num_workers = num_workers
+        self.cache_num = cache_num
+        self.cache_rate = cache_rate
+        self.dist = dist
+
+        data_list = json.load(open(json_path, "r"))
+
+        if "train" in data_list.keys():
+            self.train_list = data_list["train"]
+        if "validation" in data_list.keys():
+            self.val_list = data_list["validation"]
+
+    def val_transforms(
+        self,
+    ):
+        return self.train_transforms()
+
+    def train_transforms(
+        self,
+    ):
+        transforms = Compose(
+            [
+                LoadImaged(keys=["image"]),
+                EnsureChannelFirstd(keys=["image"]),
+                Orientationd(keys=["image"], axcodes="RAS"),
+                Spacingd(
+                    keys=["image"],
+                    pixdim=self.downsample_ratio,
+                    mode=("bilinear"),
+                ),
+                ScaleIntensityRanged(
+                    keys=["image"],
+                    a_min=-175,
+                    a_max=250,
+                    b_min=0.0,
+                    b_max=1.0,
+                    clip=True,
+                ),
+                CropForegroundd(keys=["image"], source_key="image"),
+                RandSpatialCropSamplesd(
+                    keys=["image"],
+                    roi_size=(self.img_size, self.img_size, self.depth),
+                    random_size=False,
+                    num_samples=1,
+                ),
+                SpatialPadd(
+                    keys=["image"],
+                    spatial_size=(self.img_size, self.img_size, self.depth),
+                ),
+                # RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
+                # RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
+                ToTensord(keys=["image"]),
+                PermuteImage(),
+            ]
+        )
+
+        return transforms
+
+    def setup(self, stage: Optional[str] = None):
+        # Assign Train split(s) for use in Dataloaders
+        if stage in [None, "train"]:
+            if self.dist:
+                train_partition = partition_dataset(
+                    data=self.train_list,
+                    num_partitions=ptdist.get_world_size(),
+                    shuffle=True,
+                    even_divisible=True,
+                    drop_last=False,
+                )[ptdist.get_rank()]
+                valid_partition = partition_dataset(
+                    data=self.val_list,
+                    num_partitions=ptdist.get_world_size(),
+                    shuffle=False,
+                    even_divisible=True,
+                    drop_last=False,
+                )[ptdist.get_rank()]
+                # self.cache_num //= ptdist.get_world_size()
+            else:
+                train_partition = self.train_list
+                valid_partition = self.val_list
+
+            if any([self.cache_num, self.cache_rate]) > 0:
+                train_ds = CacheDataset(
+                    train_partition,
+                    cache_num=self.cache_num,
+                    cache_rate=self.cache_rate,
+                    num_workers=self.num_workers,
+                    transform=self.train_transforms(),
+                )
+                valid_ds = CacheDataset(
+                    valid_partition,
+                    cache_num=self.cache_num // 4,
+                    cache_rate=self.cache_rate,
+                    num_workers=self.num_workers,
+                    transform=self.val_transforms(),
+                )
+            else:
+                train_ds = PersistentDataset(
+                    train_partition,
+                    transform=self.train_transforms(),
+                    cache_dir=self.cache_dir,
+                )
+                valid_ds = PersistentDataset(
+                    valid_partition,
+                    transform=self.val_transforms(),
+                    cache_dir=self.cache_dir,
+                )
+
+            return {"train": train_ds, "validation": valid_ds}
+
+        if stage in [None, "test"]:
+            if any([self.cache_num, self.cache_rate]) > 0:
+                test_ds = CacheDataset(
+                    self.val_list,
+                    cache_num=self.cache_num // 4,
+                    cache_rate=self.cache_rate,
+                    num_workers=self.num_workers,
+                    transform=self.val_transforms(),
+                )
+            else:
+                test_ds = PersistentDataset(
+                    self.val_list,
+                    transform=self.val_transforms(),
+                    cache_dir=self.cache_dir,
+                )
+
+            return {"test": test_ds}
+
+        return {"train": None, "validation": None}
+
+    def train_dataloader(self, train_ds):
+        # def collate_fn(examples):
+        #     pixel_values = torch.stack([example["image"] for example in examples])
+        #     mask = torch.stack([example["mask"] for example in examples])
+        #     return {"pixel_values": pixel_values, "bool_masked_pos": mask}
+
+        return torch.utils.data.DataLoader(
+            train_ds,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            shuffle=True,
+            collate_fn=pad_list_data_collate,
+            # collate_fn=collate_fn
+            # drop_last=False,
+            # prefetch_factor=4,
+        )
+
+    def val_dataloader(self, valid_ds):
+        return torch.utils.data.DataLoader(
+            valid_ds,
+            batch_size=self.val_batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            shuffle=False,
+            # drop_last=False,
+            collate_fn=pad_list_data_collate,
+            # prefetch_factor=4,
+        )
+
+    def test_dataloader(self, test_ds):
+        return torch.utils.data.DataLoader(
+            test_ds,
+            batch_size=self.val_batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            shuffle=False,
+            # drop_last=False,
+            collate_fn=pad_list_data_collate,
+            # prefetch_factor=4,
+        )

tumor_mask.py

@@ -0,0 +1,84 @@
+import numpy as np
+from scipy import ndimage
+
+
+def extract_tumor_and_peritumoral(mask_volume, peritumoral_margin=2, patch_size=(16, 16, 16)):
+    """
+    Extract tumor and peritumoral regions from a 3D annotation mask.
+    Flattens dilated mask into sequence of patches and creates position mask.
+
+    Parameters:
+    mask_volume: 3D numpy array (z, y, x) with tumor annotations (1 for tumor, 0 for background)
+    peritumoral_margin: Integer specifying the margin size (in voxels) for peritumoral region
+    patch_size: Tuple (z,y,x) specifying size of patches to use
+
+    Returns:
+    tumor_coords: List of coordinates (z, y, x) for tumor region
+    peritumoral_coords: List of coordinates (z, y, x) for peritumoral region
+    patch_mask: Binary mask indicating if patches contain tumor (1) or not (0)
+    """
+
+    # Get tumor coordinates
+    tumor_coords = np.where(mask_volume == 1)
+    tumor_coords = list(zip(tumor_coords[0], tumor_coords[1], tumor_coords[2]))
+
+    # Create dilated mask for peritumoral region
+    dilated_mask = ndimage.binary_dilation(
+        mask_volume,
+        structure=np.ones((peritumoral_margin * 2 + 1, peritumoral_margin * 2 + 1, peritumoral_margin * 2 + 1)),
+    )
+
+    # Create patch position mask
+    z_steps = mask_volume.shape[0] // patch_size[0]
+    y_steps = mask_volume.shape[1] // patch_size[1]
+    x_steps = mask_volume.shape[2] // patch_size[2]
+
+    patch_mask = np.zeros((z_steps, y_steps, x_steps))
+
+    for z in range(z_steps):
+        for y in range(y_steps):
+            for x in range(x_steps):
+                patch = dilated_mask[
+                    z * patch_size[0] : (z + 1) * patch_size[0],
+                    y * patch_size[1] : (y + 1) * patch_size[1],
+                    x * patch_size[2] : (x + 1) * patch_size[2],
+                ]
+                if np.any(patch):
+                    patch_mask[z, y, x] = 1
+
+    return tumor_coords, patch_mask.flatten()
+
+
+# Example usage
+def main():
+    # Create sample data for testing
+    volume_shape = (96, 96, 96)
+    mask_volume = np.zeros(volume_shape)
+
+    # Create a synthetic tumor mask in the middle
+    mask_volume[40:60, 40:60, 40:60] = 1
+
+    # Test parameters
+    patch_size = (16, 16, 16)
+    peritumoral_margin = 5
+
+    # Call function and get results
+    tumor_coords, patch_mask = extract_tumor_and_peritumoral(
+        mask_volume, peritumoral_margin=peritumoral_margin, patch_size=patch_size
+    )
+
+    # Print test results
+    print(f"Volume shape: {volume_shape}")
+    print(f"Tumor volume: {len(tumor_coords)}")
+    print(f"Number of total patches: {patch_mask.shape}")
+    print(f"Number of patches containing tumor/peritumoral region: {np.sum(patch_mask)}")
+
+    # Validate results
+    assert len(tumor_coords) > 0, "No tumor coordinates found"
+    assert len(patch_mask) == np.prod(np.array(volume_shape) // np.array(patch_size)), "Incorrect patch mask size"
+
+    return tumor_coords, patch_mask
+
+
+if __name__ == "__main__":
+    tumor_coords, patch_mask = main()