diff --git a/pyproject.toml b/pyproject.toml index cdc6cec..e2ae4b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -172,7 +172,6 @@ ignore = [ "BLE001", # blind `except Exception` "E501", # line too long (the formatter handles what it can) "E741", # ambiguous variable name (`l` for label is idiomatic here) - "F811", # redefinition (triggered by the __main__ demo blocks) "FURB171", # membership test against a single-item container "N801", # class name not CapWords (transform names mirror nnU-Net's) "N802", # function name not lowercase @@ -220,6 +219,12 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" # the drop-in contract with nnUNetTrainer. "smauglab/trainers/**" = ["B008"] +[tool.ruff.lint.isort] +# scripts/_common.py is imported by the standalone scripts as a bare `_common` +# (running `python scripts/foo.py` puts scripts/ on sys.path). Without this, isort +# files it under third-party and sorts it in among monai/torch/wandb. +known-first-party = ["_common"] + [tool.ruff.lint.mccabe] max-complexity = 20 diff --git a/smauglab/utils/utils.py b/scripts/_common.py similarity index 59% rename from smauglab/utils/utils.py rename to scripts/_common.py index 319c109..19fb070 100644 --- a/smauglab/utils/utils.py +++ b/scripts/_common.py @@ -1,4 +1,21 @@ -import argparse +"""Helpers shared by the standalone scripts in this directory. + +These used to live in `smauglab/utils/utils.py` and so shipped in the wheel, but +nothing under `smauglab/` ever imported them once the `__main__` demo blocks were +removed -- they are MONAI-training and data-loading support for `train_monai.py` and +`generate_augmentations.py`, not part of the augmentation library. `config2parser` +and `sig_fn` came along too and had no callers at all; they are gone. + +`smauglab/utils/image.py` deliberately did NOT move: five modules in the sibling +segtransferaug repository import `smauglab.utils.image.Image`, so it is a real part +of the public API despite having no in-package consumer. + +Scripts here are run as `python scripts/.py`, which puts this directory on +sys.path, so `from _common import ...` resolves. +""" + +from __future__ import annotations + import json import os @@ -6,8 +23,9 @@ from progress.bar import Bar -def fetch_image_config(config_data, split="TRAINING"): - """ +def fetch_image_config(config_data: dict, split: str = "TRAINING") -> tuple[list[dict], list]: + """Resolve a data config's image/label pairs for one split. + :param config_data: Config dict where every label used for TRAINING, VALIDATION and/or TESTING has its path specified :param split: Split of the data needed in the config file ('TRAINING', 'VALIDATION', 'TESTING'). :return: out_list: list of dictionary with image and label paths (like monai load_decathlon_datalist) @@ -28,7 +46,7 @@ def fetch_image_config(config_data, split="TRAINING"): err = [] out_list = [] - for di in dict_list: + for i, di in enumerate(dict_list): input_img_path = os.path.join(config_data["DATASETS_PATH"], di["IMAGE"]) input_seg_path = os.path.join(config_data["DATASETS_PATH"], di["LABEL"]) if not os.path.exists(input_img_path): @@ -36,27 +54,18 @@ def fetch_image_config(config_data, split="TRAINING"): else: out_list.append({"image": os.path.abspath(input_img_path), "segmentation": os.path.abspath(input_seg_path)}) - # Plot progress - bar.suffix = f"{dict_list.index(di) + 1}/{len(dict_list)}" + # Plot progress. Indexing by enumerate, not dict_list.index(di): the latter is + # a linear scan per item (quadratic overall) and reports the wrong number + # whenever two entries are equal. + bar.suffix = f"{i + 1}/{len(dict_list)}" bar.next() bar.finish() return out_list, err -def config2parser(config_path): - """ - Create a parser object from a json file - """ - # Read json file and create a dictionary - with open(config_path) as file: - config_dict = json.load(file) - - return argparse.Namespace(**config_dict) +def parser2config(args, path_out: str) -> None: + """Extract the parameters from an input parser to create a config json file. - -def parser2config(args, path_out): - """ - Extract the parameters from an input parser to create a config json file :param args: parser arguments :param path_out: path out of the config file """ @@ -78,31 +87,25 @@ def parser2config(args, path_out): outfile.write(json_object) -def tuple_type_int(strings): - """ - Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument - """ +def tuple_type_int(strings: str) -> tuple[int, ...]: + """Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument""" strings = strings.replace("(", "").replace(")", "") - mapped_int = map(int, strings.split(",")) - return tuple(mapped_int) + return tuple(map(int, strings.split(","))) -def tuple_type_float(strings): - """ - Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument - """ +def tuple_type_float(strings: str) -> tuple[float, ...]: + """Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument""" strings = strings.replace("(", "").replace(")", "") - mapped_float = map(float, strings.split(",")) - return tuple(mapped_float) + return tuple(map(float, strings.split(","))) -def tuple2string(t): +def tuple2string(t) -> str: return str(t).replace(" ", "").replace("(", "").replace(")", "").replace(",", "-") -def adjust_learning_rate(optimizer, lr, gamma): - """ - Sets the learning rate to the initial LR decayed by schedule +def adjust_learning_rate(optimizer, lr: float, gamma: float) -> float: + """Set the learning rate to the initial LR decayed by schedule. + Copied from https://github.com/spinalcordtoolbox/disc-labeling-hourglass """ lr *= gamma @@ -111,8 +114,9 @@ def adjust_learning_rate(optimizer, lr, gamma): return lr -def compute_dsc(gt_mask, pred_mask, sigmoid=False): - """ +def compute_dsc(gt_mask, pred_mask, sigmoid: bool = False): + """Dice similarity coefficient. + :param gt_mask: Ground truth mask used as the reference :param pred_mask: Prediction mask :param sigmoid: Apply sigmoid on prediction if True (default=False) @@ -120,26 +124,22 @@ def compute_dsc(gt_mask, pred_mask, sigmoid=False): :return: dsc=2*intersection/(number of non zero pixels) """ if sigmoid: - pred_mask = sig_fn(pred_mask) + pred_mask = 1 / (1 + np.exp(-pred_mask)) numerator = 2 * (gt_mask * pred_mask).sum() denominator = gt_mask.sum() + pred_mask.sum() if denominator == 0: # Both ground truth and prediction are empty return 0 - else: - return numerator / denominator - + return numerator / denominator -def sig_fn(z): - return 1 / (1 + np.exp(-z)) - -def get_validation_image(in_img, target_img, pred_img, sigmoid=False): +def get_validation_image(in_img, target_img, pred_img, sigmoid: bool = False): + """Stack input / target / prediction mid-slices into one image for logging.""" in_img = in_img.data.cpu().numpy() target_img = target_img.data.cpu().numpy() pred_img = pred_img.data.cpu().numpy() if sigmoid: - pred_img = sig_fn(pred_img) + pred_img = 1 / (1 + np.exp(-pred_img)) in_all = [] target_all = [] pred_all = [] @@ -156,9 +156,9 @@ def get_validation_image(in_img, target_img, pred_img, sigmoid=False): y_pred = y_pred[shape[0] // 2, :, :] # Normalize intensity - x = normalize(x) * 255 - y = normalize(y) * 255 - y_pred = normalize(y_pred) * 255 + x = normalize_percentile(x) * 255 + y = normalize_percentile(y) * 255 + y_pred = normalize_percentile(y_pred) * 255 # Regroup batch in_all.append(x) @@ -176,11 +176,13 @@ def get_validation_image(in_img, target_img, pred_img, sigmoid=False): return img_result, target_line_arr, pred_line_arr -def normalize(arr): - """ - Normalize image using percentiles +def normalize_percentile(arr: np.ndarray) -> np.ndarray: + """Rescale using the 10th/90th percentiles. + + Renamed from `normalize`: three functions in this repository shared that name and + two of them computed something else (min-max, in the GPU demo blocks). The name now + says which one this is. See `normalize_minmax` in demo_augmentations.py. """ - # Use 10th percentile p10 = np.percentile(arr, 10) p90 = np.percentile(arr, 90) return (arr - p10) / (p90 - p10 + 0.00001) diff --git a/scripts/demo_augmentations.py b/scripts/demo_augmentations.py new file mode 100644 index 0000000..ecb542f --- /dev/null +++ b/scripts/demo_augmentations.py @@ -0,0 +1,270 @@ +"""Render a before/after montage of a SmaugLab pipeline on a real volume. + +This replaces three `if __name__ == "__main__":` blocks that used to live inside the +shipped package -- one each in `transforms/gpu/transforms.py`, +`transforms/gpu/transforms_list.py` and `transforms/cpu/transforms.py`. The two GPU +blocks were ~205 lines of near-identical copy (they differed only in a hardcoded home +directory and `cuda(device=7)` vs `cuda()`), and all three hardcoded absolute paths +into one person's machine, so nobody else could run them. They also imported `cv2`, +which is not a SmaugLab dependency; this script writes PNGs through torchvision, +which is. + + # GPU pipeline, one subject + python scripts/demo_augmentations.py --image sub-01_T1w.nii.gz --seg sub-01_dseg.nii.gz + + # GPU pipeline, two subjects batched together (what the old GPU demos did) + python scripts/demo_augmentations.py \ + --image sub-01_T1w.nii.gz --seg sub-01_dseg.nii.gz \ + --image sub-02_T2w.nii.gz --seg sub-02_dseg.nii.gz \ + --config smauglab/configs/transform_params_gpu.json --device cuda:0 + + # CPU pipeline, a grid of repeated draws (what the old CPU demo did) + python scripts/demo_augmentations.py --backend cpu --repeats 24 \ + --image sub-01_T1w.nii.gz --seg sub-01_dseg.nii.gz +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import torch + +from smauglab.utils.image import Image, resample_nib + + +def default_config_path() -> Path: + """The packaged GPU config, resolved the way smauglab.add_trainer resolves its own.""" + import importlib.resources + + from smauglab import configs + + return Path(str(importlib.resources.files(configs))) / "transform_params_gpu.json" + + +#: Segmentation values pulled into their own channels. The old demos hardcoded two +#: different sets, one per dataset; they are only used to build a multi-channel mask +#: so the pipeline's mask handling gets exercised, so any distinct labels will do. +DEFAULT_SEG_LABELS = (12, 13, 14, 15, 16) + + +# --- small array helpers ---------------------------------------------------------- +# +# `normalize` here is min-max, and is the one the GPU demos carried (identically, in +# two files). It is deliberately NOT the percentile-based `normalize` that used to sit +# in smauglab/utils/utils.py -- three functions shared that name and computed two +# different things. + + +def normalize_minmax(arr: np.ndarray) -> np.ndarray: + """Rescale to [0, 1].""" + min_val = np.min(arr) + max_val = np.max(arr) + return (arr - min_val) / (max_val - min_val + 1e-8) + + +def pad_to(arr: np.ndarray, shape: tuple[int, ...]) -> np.ndarray: + """Zero-pad `arr` up to `shape`, centred. Axes already at or over size are left alone.""" + pad_width = [] + for i in range(len(shape)): + total = max(0, shape[i] - arr.shape[i]) + pad_width.append((total // 2, total - total // 2)) + return np.pad(arr, pad_width, mode="constant", constant_values=0) + + +def mid_slices(volume: np.ndarray, pad_shape: tuple[int, int]) -> np.ndarray: + """The three orthogonal centre slices of a [D, H, W] volume, side by side.""" + return np.concatenate( + [ + normalize_minmax(pad_to(volume[volume.shape[0] // 2, :, :], pad_shape)), + normalize_minmax(pad_to(volume[:, :, volume.shape[2] // 2], pad_shape)), + normalize_minmax(pad_to(volume[:, volume.shape[1] // 2, :], pad_shape)), + ], + axis=1, + ) + + +def write_png(image: np.ndarray, path: Path) -> None: + """Write a 2-D float array in [0, 1] as an 8-bit greyscale PNG.""" + from torchvision.utils import save_image + + path.parent.mkdir(parents=True, exist_ok=True) + save_image(torch.from_numpy(np.ascontiguousarray(image)).float().clamp(0, 1).unsqueeze(0), str(path)) + print(f"wrote {path}") + + +# --- loading ---------------------------------------------------------------------- + + +def load_volume(path: str, interpolation: str) -> torch.Tensor: + """Load a NIfTI, reorient to RSP and resample to 1 mm isotropic.""" + image = Image(path).change_orientation("RSP") + image = resample_nib(image, new_size=[1, 1, 1], new_size_type="mm", interpolation=interpolation) + return torch.from_numpy(image.data.copy()) + + +def centre_crop(tensor: torch.Tensor, shape: list[int]) -> torch.Tensor: + """Centre-crop a [D, H, W] tensor down to `shape`.""" + gap = (torch.tensor(tensor.shape) - torch.tensor(shape)) // 2 + return tensor[gap[0] : gap[0] + shape[0], gap[1] : gap[1] + shape[1], gap[2] : gap[2] + shape[2]] + + +def load_subject(image_path: str, seg_path: str) -> tuple[torch.Tensor, torch.Tensor]: + """One subject as (image [D,H,W] float, label map [D,H,W]).""" + return load_volume(image_path, "linear").to(torch.float32), load_volume(seg_path, "nn") + + +def stack_subjects(subjects: list[tuple[torch.Tensor, torch.Tensor]], labels: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor]: + """Batch subjects into (image [B,2,D,H,W], mask [B,C,D,H,W]). + + Subjects rarely share a shape, so everything is centre-cropped to the smallest + common one first. Channel 1 of the image is the binarised segmentation: the + pipeline must leave it alone, which is what the montage's last row shows. + """ + common = [min(subject[0].shape[dim] for subject in subjects) for dim in range(3)] + + images, masks = [], [] + for image, seg_all in subjects: + image_c = centre_crop(image, common) + seg_c = centre_crop(seg_all, common) + + mask = torch.zeros((1, len(labels), *seg_c.shape)) + for i, value in enumerate(labels): + mask[0, i] = seg_c == value + + images.append(torch.cat([image_c.unsqueeze(0), seg_c.bool().int().unsqueeze(0)], dim=0).unsqueeze(0)) + masks.append(mask) + + return torch.cat(images, dim=0), torch.cat(masks, dim=0) + + +# --- the two demos ---------------------------------------------------------------- + + +def demo_gpu(args: argparse.Namespace) -> int: + from smauglab.transforms.gpu.transforms import AugTransformsGPU + + subjects = [load_subject(image, seg) for image, seg in zip(args.image, args.seg)] + image_tensor, mask_tensor = stack_subjects(subjects, args.labels) + + augmentor = AugTransformsGPU(args.config).to(args.device) + image_tensor = image_tensor.to(args.device) + mask_tensor = mask_tensor.to(args.device) + + augmented_image, augmented_mask = augmentor(image_tensor.clone(), mask_tensor.clone()) + + # The old demos asserted these inline and are the reason the script is worth + # keeping: a pipeline that changes shape or emits NaN is broken in a way the unit + # tests' 24-voxel volumes do not always surface. + if augmented_image.shape != image_tensor.shape: + raise ValueError(f"augmented image shape {tuple(augmented_image.shape)} != input {tuple(image_tensor.shape)}") + if augmented_mask.shape != mask_tensor.shape: + raise ValueError(f"augmented mask shape {tuple(augmented_mask.shape)} != input {tuple(mask_tensor.shape)}") + if torch.isnan(augmented_image).any(): + raise ValueError("NaNs in the augmented image") + if torch.isnan(augmented_mask).any(): + raise ValueError("NaNs in the augmented mask") + + image_np = image_tensor.cpu().detach().numpy() + augmented_image_np = augmented_image.cpu().detach().numpy() + # Collapse the one-hot mask channels so the montage shows one picture per subject. + mask_np = mask_tensor.cpu().detach().numpy().sum(axis=1) + augmented_mask_np = augmented_mask.cpu().detach().numpy().sum(axis=1) + + pad_shape = 2 * (max(image_np.shape[2:]),) + out_dir = Path(args.out_dir) + + for b in range(image_np.shape[0]): + montage = np.concatenate( + [ + mid_slices(image_np[b, 0], pad_shape), + mid_slices(mask_np[b], pad_shape), + mid_slices(augmented_image_np[b, 0], pad_shape), + mid_slices(augmented_mask_np[b], pad_shape), + # Channel 1 is the untouched segmentation channel; it should look + # exactly like the input mask row above. + mid_slices(augmented_image_np[b, 1], pad_shape), + ], + axis=0, + ) + write_png(montage, out_dir / f"combined_{b}.png") + + print(augmentor) + return 0 + + +def demo_cpu(args: argparse.Namespace) -> int: + from smauglab.transforms.cpu.transforms import AugTransforms + + image, seg_all = load_subject(args.image[0], args.seg[0]) + image_tensor = image.unsqueeze(0) + + mask = torch.zeros((len(args.labels), *seg_all.shape)) + for i, value in enumerate(args.labels): + mask[i] = seg_all == value + + augmentor = AugTransforms( + json_path=args.config, + do_dummy_2d_data_aug=False, + patch_size=tuple(args.patch_size), + rotation_for_DA=(-10, 10), + ) + + draws = [augmentor(image=image_tensor.detach().clone(), segmentation=mask.detach().clone()) for _ in range(args.repeats)] + + out_dir = Path(args.out_dir) + slice_index = image_tensor.shape[-3] // 2 + for key in ("image", "segmentation"): + tiles = [normalize_minmax(draw[key].detach().numpy().sum(axis=0)[slice_index]) for draw in draws] + rows = [np.concatenate(tiles[i : i + args.columns], axis=1) for i in range(0, len(tiles), args.columns)] + # A short final row would not concatenate against the full-width ones. + rows = [row for row in rows if row.shape[1] == rows[0].shape[1]] + write_png(np.concatenate(rows, axis=0), out_dir / f"transforms_{key}.png") + + print(augmentor) + return 0 + + +# --- wiring ----------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--image", action="append", required=True, help="NIfTI image; repeat for a multi-subject batch") + parser.add_argument("--seg", action="append", required=True, help="matching NIfTI segmentation; repeat alongside --image") + parser.add_argument("--config", default=None, help="config JSON (default: the packaged GPU config)") + parser.add_argument("--backend", choices=["gpu", "cpu"], default="gpu") + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--out-dir", default="img", help="where the PNGs go (default: img/)") + parser.add_argument( + "--labels", + type=int, + nargs="+", + default=list(DEFAULT_SEG_LABELS), + help=f"segmentation values to split into mask channels (default: {' '.join(map(str, DEFAULT_SEG_LABELS))})", + ) + parser.add_argument("--repeats", type=int, default=24, help="cpu backend: how many draws to render") + parser.add_argument("--columns", type=int, default=6, help="cpu backend: tiles per row") + parser.add_argument("--patch-size", type=int, nargs=3, default=[128, 128, 128], help="cpu backend: patch size") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + + if len(args.image) != len(args.seg): + print(f"got {len(args.image)} --image but {len(args.seg)} --seg; they pair up one to one") + return 2 + if args.config is None: + args.config = str(default_config_path()) + args.labels = tuple(args.labels) + + return demo_gpu(args) if args.backend == "gpu" else demo_cpu(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/generate_augmentations.py b/scripts/generate_augmentations.py index 89dd017..d6a01b6 100644 --- a/scripts/generate_augmentations.py +++ b/scripts/generate_augmentations.py @@ -10,9 +10,9 @@ import torch from tqdm.contrib.concurrent import process_map +from _common import fetch_image_config from smauglab.transforms.cpu.transforms import AugTransforms from smauglab.utils.image import Image, resample_nib, zeros_like -from smauglab.utils.utils import fetch_image_config warnings.filterwarnings("ignore") diff --git a/scripts/train_monai.py b/scripts/train_monai.py index 4e84e5d..b37f6ea 100644 --- a/scripts/train_monai.py +++ b/scripts/train_monai.py @@ -28,13 +28,10 @@ from torch import optim from tqdm import tqdm -from smauglab import configs - -# Import SmaugLab GPU transforms 🐞 -from smauglab.transforms.gpu.transforms import AugTransformsGPU - # Import SmaugLab custom transforms -from smauglab.utils.utils import ( +# Script-local helpers. These used to be smauglab.utils.utils, which shipped in the +# wheel despite nothing in the library importing it. +from _common import ( adjust_learning_rate, compute_dsc, fetch_image_config, @@ -44,6 +41,10 @@ tuple_type_float, tuple_type_int, ) +from smauglab import configs + +# Import SmaugLab GPU transforms 🐞 +from smauglab.transforms.gpu.transforms import AugTransformsGPU def get_parser(): diff --git a/smauglab/transforms/cpu/transforms.py b/smauglab/transforms/cpu/transforms.py index bbfb049..97cdaa3 100644 --- a/smauglab/transforms/cpu/transforms.py +++ b/smauglab/transforms/cpu/transforms.py @@ -336,77 +336,3 @@ def _build_transforms(self): ) return transforms - - -if __name__ == "__main__": - # Example usage - import importlib - - import cv2 - - from smauglab import configs - from smauglab.transforms.gpu.transforms import AugTransformsGPU - from smauglab.utils.image import Image, resample_nib - from smauglab.utils.utils import normalize - - configs_path = importlib.resources.files(configs) - json_path = str(configs_path / "transform_params_hybrid_TAGE.json") - - # Load images and masks tensors - img_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz" - img = Image(img_path).change_orientation("RSP") - img = resample_nib(img, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") - img_tensor = torch.from_numpy(img.data.copy()).to(torch.float32).unsqueeze(0) - - seg_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz" - seg = Image(seg_path).change_orientation("RSP") - seg = resample_nib(seg, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") - seg_tensor_all = torch.from_numpy(seg.data.copy()) - - # Add segmentation values to different channels - seg_tensor = torch.zeros((5, *seg_tensor_all.shape)) - for i, value in enumerate([12, 13, 14, 15, 16]): - seg_tensor[i] = seg_tensor_all == value - - # Example usage - aug_transforms = AugTransforms( - json_path=json_path, do_dummy_2d_data_aug=False, patch_size=(128, 128, 128), rotation_for_DA=(-10, 10), mirror_axes=None - ) - - augmentor_gpu = AugTransformsGPU(json_path) - - # Apply transforms - tensor_dict = {} - gpu = False - for i in range(24): - tensor_dict[f"transfo_{i + 1!s}"] = aug_transforms(image=img_tensor.detach().clone(), segmentation=seg_tensor.detach().clone()) - - if gpu: - augmented_img, augmented_seg = augmentor_gpu( - tensor_dict[f"transfo_{i + 1!s}"]["image"].cuda().unsqueeze(0).clone(), - tensor_dict[f"transfo_{i + 1!s}"]["segmentation"].cuda().unsqueeze(0).clone(), - ) - tensor_dict[f"transfo_{i + 1!s}"]["image"] = augmented_img.cpu().squeeze(0) - tensor_dict[f"transfo_{i + 1!s}"]["segmentation"] = augmented_seg.cpu().squeeze(0) - - nb_img = len(tensor_dict.keys()) - nb_col = 6 - for key in ["image", "segmentation"]: - output = [] - line: list[np.ndarray] = [] - aug: list[list[str]] = [[]] - for _idx, (augment, _dic) in enumerate(tensor_dict.items()): - if len(line) < nb_col: - img = 255 * normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0, 64]) - line.append(img) - aug[-1].append(augment) - else: - output.append(np.concatenate(line, axis=1)) - img = 255 * normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0, 64]) - line = [img] - aug.append([augment]) - output.append(np.concatenate(line, axis=1)) - - out_img = np.concatenate(output, axis=0) - cv2.imwrite(f"img/transforms_default+plus_{key}.png", out_img) - print(aug_transforms) diff --git a/smauglab/transforms/gpu/transforms.py b/smauglab/transforms/gpu/transforms.py index 42f158d..05d1ca6 100644 --- a/smauglab/transforms/gpu/transforms.py +++ b/smauglab/transforms/gpu/transforms.py @@ -505,210 +505,3 @@ def pad_numpy_array(arr, shape): ] padded_arr = np.pad(arr, pad_width, mode="constant", constant_values=0) return padded_arr - - -if __name__ == "__main__": - # Example usage - import importlib - - from smauglab import configs - from smauglab.utils.image import Image, resample_nib - - configs_path = importlib.resources.files(configs) - json_path = str(configs_path / "transform_params_gpu.json") - augmentor = AugTransformsGPU(json_path) - - # Load images and masks tensors - img_path = "/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz" - img = Image(img_path).change_orientation("RSP") - img = resample_nib(img, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") - img_tensor = torch.from_numpy(img.data.copy()).to(torch.float32) - - seg_path = "/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz" - seg = Image(seg_path).change_orientation("RSP") - seg = resample_nib(seg, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") - seg_tensor_all = torch.from_numpy(seg.data.copy()) - - img2_path = "/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/sub-002/anat/sub-002_acq-lowresSag_T2w.nii.gz" - img2 = Image(img2_path).change_orientation("RSP") - img2 = resample_nib(img2, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") - img2_tensor = torch.from_numpy(img2.data.copy()).to(torch.float32) - - seg2_path = "/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/derivatives/labels/sub-002/anat/sub-002_acq-lowresSag_T2w_label-spine_dseg.nii.gz" - seg2 = Image(seg2_path).change_orientation("RSP") - seg2 = resample_nib(seg2, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") - seg2_tensor_all = torch.from_numpy(seg2.data.copy()) - - # Combine two images to same size - new_shape = [] - for dim in range(3): - size1 = img_tensor.shape[dim] - size2 = img2_tensor.shape[dim] - min_size = min(size1, size2) - new_shape.append(min_size) - - new_img_tensor = torch.zeros(new_shape) - new_img2_tensor = torch.zeros(new_shape) - new_seg_tensor_all = torch.zeros(new_shape) - new_seg2_tensor_all = torch.zeros(new_shape) - - gap = (torch.tensor(img_tensor.shape) - torch.tensor(new_shape)) // 2 - gap2 = (torch.tensor(img2_tensor.shape) - torch.tensor(new_shape)) // 2 - new_img_tensor = img_tensor[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] - new_img2_tensor = img2_tensor[gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2]] - new_seg_tensor_all = seg_tensor_all[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] - new_seg2_tensor_all = seg2_tensor_all[ - gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2] - ] - - # Add segmentation values to different channels - seg_tensor = torch.zeros((1, 5, *new_seg_tensor_all.shape)) - for i, value in enumerate([12, 13, 14, 15, 16]): - seg_tensor[0, i] = new_seg_tensor_all == value - - seg2_tensor = torch.zeros((1, 5, *new_seg2_tensor_all.shape)) - for i, value in enumerate([50, 45, 44, 43, 42]): - seg2_tensor[0, i] = new_seg2_tensor_all == value - - # Format tensors to match expected input shape (B, C, D, H, W) - img_tensor = torch.cat([new_img_tensor.unsqueeze(0), new_seg_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( - 0 - ) # Add batch dimension and second channel - img2_tensor = torch.cat([new_img2_tensor.unsqueeze(0), new_seg2_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( - 0 - ) # Add batch dimension and second channel - - # Add batch - img_tensor = torch.cat([img_tensor, img2_tensor], dim=0) - seg_tensor = torch.cat([seg_tensor, seg2_tensor], dim=0) - - # Move to GPU - img_tensor = img_tensor.cuda(device=7) - seg_tensor = seg_tensor.cuda(device=7) - augmentor = augmentor.cuda(device=7) - - # Apply augmentations - augmented_img, augmented_seg = augmentor(img_tensor.clone(), seg_tensor.clone()) - - if augmented_img.shape != img_tensor.shape: - raise ValueError("Augmented image shape does not match input shape.") - if augmented_seg.shape != seg_tensor.shape: - raise ValueError("Augmented segmentation shape does not match input shape.") - # Check if nans are present - if torch.isnan(augmented_img).any(): - raise ValueError("NaNs found in augmented image.") - if torch.isnan(augmented_seg).any(): - raise ValueError("NaNs found in augmented segmentation.") - - import os - import warnings - - import cv2 - import numpy as np - - warnings.simplefilter("always") - - # Convert tensors to numpy arrays - img_tensor_np = img_tensor.cpu().detach().numpy() - seg_tensor_np = seg_tensor.cpu().detach().numpy() - augmented_img_np = augmented_img.cpu().detach().numpy() - augmented_seg_np = augmented_seg.cpu().detach().numpy() - - # Concatenate segmentation channels for visualization - seg_tensor_np = np.sum(seg_tensor_np, axis=1) - augmented_seg_np = np.sum(augmented_seg_np, axis=1) - - pad_shape = 2 * (np.max(img_tensor_np.shape[2:]),) - - # Combine tensors into single output for visualization - os.makedirs("img", exist_ok=True) - img_line = np.concatenate( - [ - normalize(pad_numpy_array(img_tensor_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_img_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - seg_line = np.concatenate( - [ - normalize(pad_numpy_array(seg_tensor_np[0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_seg_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_seg_np[0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - not_augmented_channel_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[0, 1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - combined_img = np.concatenate([img_line, seg_line, augmented_img_line, augmented_seg_line, not_augmented_channel_line], axis=0) - cv2.imwrite("img/combined.png", combined_img * 255) - - img_line2 = np.concatenate( - [ - normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_img_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - seg_line2 = np.concatenate( - [ - normalize(pad_numpy_array(seg_tensor_np[1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_seg_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_seg_np[1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - not_augmented_channel_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[1, 1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - combined_img2 = np.concatenate([img_line2, seg_line2, augmented_img_line2, augmented_seg_line2, not_augmented_channel_line2], axis=0) - cv2.imwrite("img/combined2.png", combined_img2 * 255) - - # cv2.imwrite('img/orig_img.png', normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape))*255) - # cv2.imwrite('img/aug_img.png', normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape))*255) - - print(augmentor) diff --git a/smauglab/transforms/gpu/transforms_list.py b/smauglab/transforms/gpu/transforms_list.py index 32bd632..35f08f8 100644 --- a/smauglab/transforms/gpu/transforms_list.py +++ b/smauglab/transforms/gpu/transforms_list.py @@ -779,208 +779,3 @@ def pad_numpy_array(arr, shape): ] padded_arr = np.pad(arr, pad_width, mode="constant", constant_values=0) return padded_arr - - -if __name__ == "__main__": - # Example usage - import importlib - - from smauglab import configs - from smauglab.transforms.gpu.transforms import AugTransformsGPU - from smauglab.utils.image import Image, resample_nib - - configs_path = importlib.resources.files(configs) - json_path = str(configs_path / "transform_params_gpu.json") - augmentor = AugTransformsGPU(json_path) - - # Load images and masks tensors - img_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz" - img = Image(img_path).change_orientation("RSP") - img = resample_nib(img, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") - img_tensor = torch.from_numpy(img.data.copy()).to(torch.float32) - - seg_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz" - seg = Image(seg_path).change_orientation("RSP") - seg = resample_nib(seg, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") - seg_tensor_all = torch.from_numpy(seg.data.copy()) - - img2_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/spider-challenge-2023/sub-002/anat/sub-002_acq-lowresSag_T2w.nii.gz" - img2 = Image(img2_path).change_orientation("RSP") - img2 = resample_nib(img2, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") - img2_tensor = torch.from_numpy(img2.data.copy()).to(torch.float32) - - seg2_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/spider-challenge-2023/derivatives/labels/sub-002/anat/sub-002_acq-lowresSag_T2w_label-spine_dseg.nii.gz" - seg2 = Image(seg2_path).change_orientation("RSP") - seg2 = resample_nib(seg2, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") - seg2_tensor_all = torch.from_numpy(seg2.data.copy()) - - # Combine two images to same size - new_shape = [] - for dim in range(3): - size1 = img_tensor.shape[dim] - size2 = img2_tensor.shape[dim] - min_size = min(size1, size2) - new_shape.append(min_size) - - new_img_tensor = torch.zeros(new_shape) - new_img2_tensor = torch.zeros(new_shape) - new_seg_tensor_all = torch.zeros(new_shape) - new_seg2_tensor_all = torch.zeros(new_shape) - - gap = (torch.tensor(img_tensor.shape) - torch.tensor(new_shape)) // 2 - gap2 = (torch.tensor(img2_tensor.shape) - torch.tensor(new_shape)) // 2 - new_img_tensor = img_tensor[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] - new_img2_tensor = img2_tensor[gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2]] - new_seg_tensor_all = seg_tensor_all[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] - new_seg2_tensor_all = seg2_tensor_all[ - gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2] - ] - - # Add segmentation values to different channels - seg_tensor = torch.zeros((1, 5, *new_seg_tensor_all.shape)) - for i, value in enumerate([12, 13, 14, 15, 16]): - seg_tensor[0, i] = new_seg_tensor_all == value - - seg2_tensor = torch.zeros((1, 5, *new_seg2_tensor_all.shape)) - for i, value in enumerate([50, 45, 44, 43, 42]): - seg2_tensor[0, i] = new_seg2_tensor_all == value - - # Format tensors to match expected input shape (B, C, D, H, W) - img_tensor = torch.cat([new_img_tensor.unsqueeze(0), new_seg_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( - 0 - ) # Add batch dimension and second channel - img2_tensor = torch.cat([new_img2_tensor.unsqueeze(0), new_seg2_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( - 0 - ) # Add batch dimension and second channel - - # Add batch - img_tensor = torch.cat([img_tensor, img2_tensor], dim=0) - seg_tensor = torch.cat([seg_tensor, seg2_tensor], dim=0) - - # Move to GPU - img_tensor = img_tensor.cuda() - seg_tensor = seg_tensor.cuda() - augmentor = augmentor.cuda() - - # Apply augmentations - augmented_img, augmented_seg = augmentor(img_tensor.clone(), seg_tensor.clone()) - - if augmented_img.shape != img_tensor.shape: - raise ValueError("Augmented image shape does not match input shape.") - if augmented_seg.shape != seg_tensor.shape: - raise ValueError("Augmented segmentation shape does not match input shape.") - # Check if nans are present - if torch.isnan(augmented_img).any(): - raise ValueError("NaNs found in augmented image.") - if torch.isnan(augmented_seg).any(): - raise ValueError("NaNs found in augmented segmentation.") - - import os - import warnings - - import cv2 - import numpy as np - - warnings.simplefilter("always") - - # Convert tensors to numpy arrays - img_tensor_np = img_tensor.cpu().detach().numpy() - seg_tensor_np = seg_tensor.cpu().detach().numpy() - augmented_img_np = augmented_img.cpu().detach().numpy() - augmented_seg_np = augmented_seg.cpu().detach().numpy() - - # Concatenate segmentation channels for visualization - seg_tensor_np = np.sum(seg_tensor_np, axis=1) - augmented_seg_np = np.sum(augmented_seg_np, axis=1) - - pad_shape = 2 * (np.max(img_tensor_np.shape[2:]),) - - # Combine tensors into single output for visualization - os.makedirs("img", exist_ok=True) - img_line = np.concatenate( - [ - normalize(pad_numpy_array(img_tensor_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_img_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - seg_line = np.concatenate( - [ - normalize(pad_numpy_array(seg_tensor_np[0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_seg_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_seg_np[0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - not_augmented_channel_line = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[0, 1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[0, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - combined_img = np.concatenate([img_line, seg_line, augmented_img_line, augmented_seg_line, not_augmented_channel_line], axis=0) - cv2.imwrite("img/combined.png", combined_img * 255) - - img_line2 = np.concatenate( - [ - normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(img_tensor_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_img_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - seg_line2 = np.concatenate( - [ - normalize(pad_numpy_array(seg_tensor_np[1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(seg_tensor_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - augmented_seg_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_seg_np[1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_seg_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - not_augmented_channel_line2 = np.concatenate( - [ - normalize(pad_numpy_array(augmented_img_np[1, 1, img_tensor_np.shape[2] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), - normalize(pad_numpy_array(augmented_img_np[1, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), - ], - axis=1, - ) - combined_img2 = np.concatenate([img_line2, seg_line2, augmented_img_line2, augmented_seg_line2, not_augmented_channel_line2], axis=0) - cv2.imwrite("img/combined2.png", combined_img2 * 255) - - print(augmentor)