# Copyright (c) MONAI Consortium

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import logging

import os

from collections.abc import Sequence

import numpy as np

from monai.config import PathLike

from monai.transforms import Compose, EnsureChannelFirstd, LoadImaged, Orientationd, Spacingd, SqueezeDimd, Transform

from monai.utils import GridSampleMode

def create_dataset(

datalist: list[dict],

output_dir: str,

dimension: int,

pixdim: Sequence[float] | float,

image_key: str = "image",

label_key: str = "label",

base_dir: PathLike | None = None,

limit: int = 0,

relative_path: bool = False,

transforms: Transform | None = None,

) -> list[dict]:

"""

Utility to pre-process and create dataset list for Deepgrow training over on existing one.

The input data list is normally a list of images and labels (3D volume) that needs pre-processing

for Deepgrow training pipeline.

Args:

datalist: A list of data dictionary. Each entry should at least contain 'image_key': <image filename>.

For example, typical input data can be a list of dictionaries::

[{'image': <image filename>, 'label': <label filename>}]

output_dir: target directory to store the training data for Deepgrow Training

pixdim: output voxel spacing.

dimension: dimension for Deepgrow training. It can be 2 or 3.

image_key: image key in input datalist. Defaults to 'image'.

label_key: label key in input datalist. Defaults to 'label'.

base_dir: base directory in case related path is used for the keys in datalist. Defaults to None.

limit: limit number of inputs for pre-processing. Defaults to 0 (no limit).

relative_path: output keys values should be based on relative path. Defaults to False.

transforms: explicit transforms to execute operations on input data.

Raises:

ValueError: When ``dimension`` is not one of [2, 3]

ValueError: When ``datalist`` is Empty

Returns:

A new datalist that contains path to the images/labels after pre-processing.

Example::

datalist = create_dataset(

datalist=[{'image': 'img1.nii', 'label': 'label1.nii'}],

base_dir=None,

output_dir=output_2d,

dimension=2,

image_key='image',

label_key='label',

pixdim=(1.0, 1.0),

limit=0,

relative_path=True

)

print(datalist[0]["image"], datalist[0]["label"])

"""

if dimension not in [2, 3]:

raise ValueError("Dimension can be only 2 or 3 as Deepgrow supports only 2D/3D Training")

if not len(datalist):

raise ValueError("Input datalist is empty")

transforms = _default_transforms(image_key, label_key, pixdim) if transforms is None else transforms

new_datalist = []

for idx, item in enumerate(datalist):

if limit and idx >= limit:

break

image = item[image_key]

label = item.get(label_key, None)

if base_dir:

image = os.path.join(base_dir, image)

label = os.path.join(base_dir, label) if label else None

image = os.path.abspath(image)

label = os.path.abspath(label) if label else None

logging.info(f"Image: {image}; Label: {label if label else None}")

data = transforms({image_key: image, label_key: label})

vol_image = data[image_key]

vol_label = data.get(label_key)

logging.info(f"Image (transform): {vol_image.shape}; Label: {None if vol_label is None else vol_label.shape}")

vol_image = np.moveaxis(vol_image, -1, 0)

if vol_label is not None:

vol_label = np.moveaxis(vol_label, -1, 0)

logging.info(f"Image (final): {vol_image.shape}; Label: {None if vol_label is None else vol_label.shape}")

if dimension == 2:

data = _save_data_2d(

vol_idx=idx,

vol_image=vol_image,

vol_label=vol_label,

dataset_dir=output_dir,

relative_path=relative_path,

)

else:

data = _save_data_3d(

vol_idx=idx,

vol_image=vol_image,

vol_label=vol_label,

dataset_dir=output_dir,

relative_path=relative_path,

)

new_datalist.extend(data)

return new_datalist

def _default_transforms(image_key, label_key, pixdim):

keys = [image_key] if label_key is None else [image_key, label_key]

mode = [GridSampleMode.BILINEAR, GridSampleMode.NEAREST] if len(keys) == 2 else [GridSampleMode.BILINEAR]

return Compose(

[

LoadImaged(keys=keys),

EnsureChannelFirstd(keys=keys),

Orientationd(keys=keys, axcodes="RAS"),

Spacingd(keys=keys, pixdim=pixdim, mode=mode),

SqueezeDimd(keys=keys),

]

)

def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):

data_list: list[dict[str, str | int]] = []

image_count = 0

label_count = 0

unique_labels_count = 0

for sid in range(vol_image.shape[0]):

image = vol_image[sid, ...]

label = vol_label[sid, ...] if vol_label is not None else None

if vol_label is not None and np.sum(label) == 0:

continue

image_file_prefix = f"vol_idx_{vol_idx:0>4d}_slice_{sid:0>3d}"

image_file = os.path.join(dataset_dir, "images", image_file_prefix)

image_file += ".npy"

os.makedirs(os.path.join(dataset_dir, "images"), exist_ok=True)

np.save(image_file, image)

image_count += 1

# Test Data

if vol_label is None:

data_list.append(

{"image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file}

)

continue

# For all Labels

unique_labels = np.unique(label.flatten())

unique_labels = unique_labels[unique_labels != 0]

unique_labels_count = max(unique_labels_count, len(unique_labels))

for idx in unique_labels:

label_file_prefix = f"{image_file_prefix}_region_{int(idx):0>2d}"

label_file = os.path.join(dataset_dir, "labels", label_file_prefix)

label_file += ".npy"

os.makedirs(os.path.join(dataset_dir, "labels"), exist_ok=True)

curr_label = (label == idx).astype(np.float32)

np.save(label_file, curr_label)

label_count += 1

data_list.append(

{

"image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file,

"label": label_file.replace(dataset_dir + os.pathsep, "") if relative_path else label_file,

"region": int(idx),

}

)

if unique_labels_count >= 20:

logging.warning(f"Unique labels {unique_labels_count} exceeds 20. Please check if this is correct.")

logging.info(

f"{vol_idx} => Image Shape: {vol_image.shape} => {image_count};"

f" Label Shape: {vol_label.shape if vol_label is not None else None} => {label_count};"

f" Unique Labels: {unique_labels_count}"

)

return data_list

def _save_data_3d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):

data_list: list[dict[str, str | int]] = []

image_count = 0

label_count = 0

unique_labels_count = 0

image_file_prefix = f"vol_idx_{vol_idx:0>4d}"

image_file = os.path.join(dataset_dir, "images", image_file_prefix)

image_file += ".npy"

os.makedirs(os.path.join(dataset_dir, "images"), exist_ok=True)

np.save(image_file, vol_image)

image_count += 1

# Test Data

if vol_label is None:

data_list.append({"image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file})

else:

# For all Labels

unique_labels = np.unique(vol_label.flatten())

unique_labels = unique_labels[unique_labels != 0]

unique_labels_count = max(unique_labels_count, len(unique_labels))

for idx in unique_labels:

label_file_prefix = f"{image_file_prefix}_region_{int(idx):0>2d}"

label_file = os.path.join(dataset_dir, "labels", label_file_prefix)

label_file += ".npy"

curr_label = (vol_label == idx).astype(np.float32)

os.makedirs(os.path.join(dataset_dir, "labels"), exist_ok=True)

np.save(label_file, curr_label)

label_count += 1

data_list.append(

{

"image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file,

"label": label_file.replace(dataset_dir + os.pathsep, "") if relative_path else label_file,

"region": int(idx),

}

)

if unique_labels_count >= 20:

logging.warning(f"Unique labels {unique_labels_count} exceeds 20. Please check if this is correct.")

logging.info(

f"{vol_idx} => Image Shape: {vol_image.shape} => {image_count};"

f" Label Shape: {vol_label.shape if vol_label is not None else None} => {label_count};"

f" Unique Labels: {unique_labels_count}"

)

return data_list