Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| from typing import Literal | |
| import numpy as np | |
| import torch | |
| from monai.data import PersistentDataset, load_decathlon_datalist | |
| from monai.transforms import ( | |
| Compose, | |
| ConcatItemsd, | |
| DeleteItemsd, | |
| EnsureTyped, | |
| LoadImaged, | |
| RandCropByPosNegLabeld, | |
| RandRotate90d, | |
| ToTensord, | |
| Transform, | |
| Transposed, | |
| ) | |
| from torch.utils.data.dataloader import default_collate | |
| from .custom_transforms import ( | |
| ClipMaskIntensityPercentilesd, | |
| ElementwiseProductd, | |
| NormalizeIntensity_customd, | |
| NormalizePSAd, | |
| ) | |
| def list_data_collate(batch: list): | |
| """ | |
| Combine instances from a list of dicts into a single dict, by stacking them along first dim | |
| [{'image' : 3xHxW}, {'image' : 3xHxW}, {'image' : 3xHxW}...] - > {'image' : Nx3xHxW} | |
| followed by the default collate which will form a batch BxNx3xHxW | |
| """ | |
| for i, item in enumerate(batch): | |
| data = item[0] | |
| data["image"] = torch.stack([ix["image"] for ix in item], dim=0) | |
| if all("final_heatmap" in ix for ix in item): | |
| data["final_heatmap"] = torch.stack([ix["final_heatmap"] for ix in item], dim=0) | |
| data["smooth_mask"] = torch.stack([ix["smooth_mask"] for ix in item], dim=0) | |
| batch[i] = data | |
| return default_collate(batch) | |
| def data_transform(args: argparse.Namespace, split) -> Transform: | |
| if split == "train": | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"), | |
| RandCropByPosNegLabeld( | |
| keys=["image", "final_heatmap", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| RandRotate90d( | |
| keys=["image", "final_heatmap", "smooth_mask"], | |
| prob=0.6, | |
| spatial_axes=(0, 1), | |
| max_k=3, | |
| ), | |
| NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std), | |
| EnsureTyped(keys=["label", "psa"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]), | |
| ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]), | |
| ] | |
| ) | |
| else: | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"), | |
| RandCropByPosNegLabeld( | |
| keys=["image", "final_heatmap", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std), | |
| EnsureTyped(keys=["label", "psa"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]), | |
| ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]), | |
| ] | |
| ) | |
| return transform | |
| """ | |
| def data_transform(args: argparse.Namespace) -> Transform: | |
| if args.use_heatmap: | |
| if args.use_psa: | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"), | |
| RandCropByPosNegLabeld( | |
| keys=["image", "final_heatmap", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std), | |
| EnsureTyped(keys=["label", "psa"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]), | |
| ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]), | |
| ] | |
| ) | |
| else: | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"), | |
| #RandRotate90d(keys=["image", "final_heatmap", "smooth_mask"], prob=0.5, spatial_axes=(0, 1)), | |
| RandCropByPosNegLabeld( | |
| keys=["image", "final_heatmap", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| EnsureTyped(keys=["label"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]), | |
| ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask"]), | |
| ] | |
| ) | |
| else: | |
| if args.use_psa: | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc","smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| RandCropByPosNegLabeld( | |
| keys=["image", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std), | |
| EnsureTyped(keys=["label", "psa"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=[ "dwi", "adc", "mask"]), | |
| ToTensord(keys=["image", "label", "smooth_mask", "psa"]), | |
| ] | |
| ) | |
| else: | |
| transform = Compose( | |
| [ | |
| LoadImaged( | |
| keys=["image", "mask", "dwi", "adc","smooth_mask"], | |
| reader="ITKReader", | |
| ensure_channel_first=True, | |
| dtype=np.float32, | |
| ), | |
| ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"), | |
| ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["image"], mask_key="mask"), | |
| NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"), | |
| ConcatItemsd( | |
| keys=["image", "dwi", "adc"], name="image", dim=0 | |
| ), # stacks to (3, H, W) | |
| RandCropByPosNegLabeld( | |
| keys=["image", "smooth_mask"], | |
| label_key="smooth_mask", | |
| spatial_size=(args.tile_size, args.tile_size, args.depth), | |
| pos=1, | |
| neg=0, | |
| num_samples=args.tile_count, | |
| ), | |
| EnsureTyped(keys=["label"], dtype=torch.float32), | |
| Transposed(keys=["image"], indices=(0, 3, 1, 2)), | |
| DeleteItemsd(keys=[ "dwi", "adc", "mask"]), | |
| ToTensord(keys=["image", "label", "smooth_mask"]), | |
| ] | |
| ) | |
| return transform | |
| """ | |
| def get_dataloader( | |
| args: argparse.Namespace, split: Literal["train", "test"] | |
| ) -> torch.utils.data.DataLoader: | |
| data_list = load_decathlon_datalist( | |
| data_list_file_path=args.dataset_json, | |
| data_list_key=split, | |
| base_dir=args.data_root, | |
| ) | |
| data_list_updated = [{**i, "psa": i.get("psa", [0, 0])} for i in data_list] | |
| cache_dir_ = os.path.join(args.logdir, "cache") | |
| os.makedirs(os.path.join(cache_dir_, split), exist_ok=True) | |
| transform = data_transform(args, split) | |
| dataset = PersistentDataset( | |
| data=data_list_updated, transform=transform, cache_dir=os.path.join(cache_dir_, split) | |
| ) | |
| loader = torch.utils.data.DataLoader( | |
| dataset, | |
| batch_size=args.batch_size, | |
| shuffle=(split == "train"), | |
| num_workers=args.workers if split == "train" else 2, | |
| pin_memory=True, | |
| multiprocessing_context="fork" if args.workers > 0 else None, | |
| sampler=None, | |
| collate_fn=list_data_collate, | |
| ) | |
| return loader | |