pipeline package

Module contents

class pipeline.Pipeline(model_type: str, modality: int, num_of_labels: int, model_path: str = '', debug: bool = False)

Bases: object

Class for managing machine learning pipeline for medical image semantic segmentation. It assists with loading models and data for training, and it automatically records metrics and save check points.

debug_mode

Whether the pipeline is in debug mode or not

Type:

bool

model

The neural network used for training or inference.

model_type

Neural network architecture of model. Currently supports UNETR and SWINUNETR

Type:

str

train_transforms

Transformations applied on the training dataset

val_transforms

Transformations applied on the validation dataset

modality

Input dimension of the loaded dataset

Type:

int

num_of_labels

Number of output classes of the dataset

Type:

int

dataset_name

Name of the dataset

Type:

str

num_train_images

Number of training images in the dataset

Type:

int

num_val_images

Number of validation images in the dataset

Type:

int

train_batch_size

Batchsize for training

Type:

int

create_model(model_type: str, modality: int, num_of_labels: int, model_path: str = '') → None

Creates a new model for the pipeline

Parameters:

• model_type (str) – Type of model the pipeline uses, takes value “UNETR” or “SWINUNETR” for their respective model types.

• modality (int) – Modality of the dataset. Determines the input dimension of the model.

• num_of_labels (int) – Number of labels to the dataset.

• model_path (str) – File path to the saved model of the same type as model_type.

inference(data_folder, output_folder, transforms) → None

Runs the prediction on the files located under self.data_folder, will save the files as Nifti (.nii.gz) format under output_folder. If output_folder is not specified, then it will be saved to the folder where the data was originally taken from.

Parameters:

• data_folder (str) – The folder where the data is located as string. All files in this folder should be medical images.

• output_folder – The folder path to save the nifti images as a string. If None, then it will save to the folder where the data files are located. (self.data_folder)

• transforms – Transformations to apply onto images before inferece. Should be similar to transformation done on validation dataset

load_data(dataset_path: str, train_transforms, val_transforms, cache_num_train: int = 24, train_batch_size: int = 1, cache_num_val: int = 6, val_batch_size: int = 1, workers: int = 4) → None

Load the training and validation data from the json file for the dataset.

Parameters:

• dataset_path (str) – File path to the json file of the dataset.

• train_transforms – Transformation done on the dataset during training.

• val_transforms – Transformation done on the dataset during validation.

• cache_num_train (int) – Number of cached data for training dataset.

• train_batch_size (int) – Batch size for training.

• cache_num_val (int) – Number of cached data for validation dataset.

• val_batch_size (int) – Batch size for validation.

• workers (int) – Number of workers working in parallel.

load_model(model_path: str) → None

Load the saved model.

Parameters:

model_path (str) – File path to the saved model of the same type as model_type.

train(max_epoch: int, epoch_val: int, learning_rate: float = 0.0001, weight_decay: float = 1e-05) → None

Initiate training for the loaded model on the loaded dataset.

Parameters:

• max_epoch – Total number of epoch to train.

• epoch_val – Number of epochs between every validation and saving the model

• learning_rate – learning rate of the training process with AdamW optimizer

• weight_decay – Weight decay for the AdamW optimizer