From b6baaac551343a87f4c80f163d7fa72a604d193f Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 12:46:05 +0000 Subject: [PATCH] refactor: move the demo blocks and script helpers out of the package Three `if __name__ == "__main__":` blocks -- 486 lines across transforms/gpu/transforms.py, transforms/gpu/transforms_list.py and transforms/cpu/transforms.py -- become one parameterised scripts/demo_augmentations.py. None of them could run for anyone but their author. All three hardcoded absolute paths into one person's machine, and the two GPU blocks differed only in that path and in cuda(device=7) vs cuda(); they were otherwise a 205-line copy of each other. They also imported cv2, which is not a SmaugLab dependency, so they would have failed at import for anyone who did try. The script writes PNGs through torchvision, which is a dependency. They shipped in the wheel, which is the real cost: 486 lines of dead demonstration code in the installed package. The wheel is now 26 modules. Dropping them removes the only reason for the F811 ruff carve-out, so that goes too. smauglab/utils/utils.py moves to scripts/_common.py. Nothing under smauglab/ imported it once the demo blocks were gone -- it is MONAI training-loop helpers, argparse tuple parsers and a Dice function, used only by scripts/train_monai.py and scripts/generate_augmentations.py. config2parser and sig_fn had no callers at all. Its `normalize` becomes normalize_percentile, because three functions in this repository shared that name and two of them computed different things; the min-max one now lives in the demo script as normalize_minmax. smauglab/utils/image.py deliberately stays. Five modules in the sibling segtransferaug repository import smauglab.utils.image.Image, so it is public API in practice despite having no in-package consumer. Verified: the 24-config seeded digests are unchanged, the wheel still builds and carries 26 modules and 30 configs with no utils/utils.py, and both scripts still import and parse their arguments. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 7 +- smauglab/utils/utils.py => scripts/_common.py | 108 +++---- scripts/demo_augmentations.py | 270 ++++++++++++++++++ scripts/generate_augmentations.py | 2 +- scripts/train_monai.py | 13 +- smauglab/transforms/cpu/transforms.py | 74 ----- smauglab/transforms/gpu/transforms.py | 207 -------------- smauglab/transforms/gpu/transforms_list.py | 205 ------------- 8 files changed, 339 insertions(+), 547 deletions(-) rename smauglab/utils/utils.py => scripts/_common.py (59%) create mode 100644 scripts/demo_augmentations.py 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)