PreviewLoader

class lib.training.data.loader.PreviewLoader(input_size: int, output_size: int, color_order: Literal['bgr', 'rgb'], input_folders: list[str], batch_size: int, sampler: None | type[RandomSampler | SequentialSampler] = None, num_samples: int = 0)

Bases: object

Generator for feeding faceswap models input data for generating preview images. Gets the next items from each of the configured loaders and collates them for feeding into a model

Parameters:

  • input_size (int) – The input size to the model

  • output_sizes – The output sizes to the model (list as some models have multi-scale outputs)

  • color_order (T.Literal['bgr', 'rgb']) – The color order of the model

  • input_folders (list[str]) – list of folders to read images from for each side being trained

  • batch_size (int) – The number of images being displayed in the preview

  • sampler (None | type[tch_data.RandomSampler | tch_data.SequentialSampler]) – The sampler to use for the data loaders. Default: None (RandomSampler)

  • num_samples (int) – Set to 0 for random previews from the image folder. Set to a positive integer for this number of images to use for a static timelapse. Default: 0

  • output_size (int)

Methods Summary

get_loader()

Obtain the dataloaders for each input/output for the model

Methods Documentation

get_loader() → DataLoader

Obtain the dataloaders for each input/output for the model

Return type:

The Training data loaders in side order

__next__() → tuple[Tensor, Tensor]

Obtain the next batch of data for each side for feeding the model

Returns:

  • inputs – The inputs to the model for each side of the model. The array is returned in (side, batch_size, *dims) where side 0 is “A” and side 1 is “B” etc.

  • targets – The full sized source image with mask in 4th channel for each side of the model in format (side, batch_size, *dims, 4) where `side 0 is “A” and side 1 is “B” etc.

Return type:

tuple[Tensor, Tensor]

get_loader() → DataLoader

Obtain the dataloaders for each input/output for the model

Return type:

The Training data loaders in side order