ModelABC

class ModelABC[source]

Abstract base class for models used in tiatoolbox.

Initialize Abstract class ModelABC.

Methods

forward

Torch method, this contains logic for using layers defined in init.

infer_batch

Run inference on an input batch.

load_weights_from_file

Helper function to load a torch model.

postproc

Define the post-processing of this class of model.

preproc

Define the pre-processing of this class of model.

to

Transfers model to cpu/gpu.

Attributes

postproc_func

Return the current post-processing function of this instance.

preproc_func

Return the current pre-processing function of this instance.

training

abstract forward(*args, **kwargs)[source]

Torch method, this contains logic for using layers defined in init.

Parameters:
Return type:

None

abstract static infer_batch(model, batch_data, *, on_gpu)[source]

Run inference on an input batch.

Contains logic for forward operation as well as I/O aggregation.

Parameters:

  • model (nn.Module) – PyTorch defined model.

  • batch_data (np.ndarray) – A batch of data generated by torch.utils.data.DataLoader.

  • on_gpu (bool) – Whether to run inference on a GPU.

Return type:

None

load_weights_from_file(weights)[source]

Helper function to load a torch model.

Parameters:

  • self (ModelABC) – A torch model as ModelABC.

  • weights (str or Path) – Path to pretrained weights.

Returns:

Torch model with pretrained weights loaded on CPU.

Return type:

torch.nn.Module

static postproc(image)[source]

Define the post-processing of this class of model.

Parameters:

image (np.ndarray)

Return type:

np.ndarray

property postproc_func: Callable

Return the current post-processing function of this instance.

static preproc(image)[source]

Define the pre-processing of this class of model.

Parameters:

image (np.ndarray)

Return type:

np.ndarray

property preproc_func: Callable

Return the current pre-processing function of this instance.

to(device='cpu')[source]

Transfers model to cpu/gpu.

Parameters:

  • model (torch.nn.Module) – PyTorch defined model.

  • device (str) – Transfers model to the specified device. Default is “cpu”.

  • self (ModelABC)

Returns:

The model after being moved to cpu/gpu.

Return type:

torch.nn.Module