          256px image (:math:`256 * 0.625 = 160`)
      config: dict
          The configuration `dict` generated from :file:`config.train.ini` containing the trainer
          plugin configuration options.
  
      Attributes
      ----------
      initialized: bool
          Flag to indicate whether :class:`ImageAugmentation` has been initialized with the training
          image size in order to cache certain augmentation operations (see :func:`initialize`)
      is_display: bool
          Flag to indicate whether these augmentations are for time-lapses/preview images (``True``)
          or standard training data (``False``)
      """
      def __init__(self, batchsize, is_display, input_size, output_shapes, coverage_ratio, config):
          logger.debug("Initializing %s: (batchsize: %s, is_display: %s, input_size: %s, "
                       "output_shapes: %s, coverage_ratio: %s, config: %s)",
                       self.__class__.__name__, batchsize, is_display, input_size, output_shapes,
 
</s> remove     def __init__(self, batchsize, is_display, input_size, output_shapes, coverage_ratio, config):
        logger.debug("Initializing %s: (batchsize: %s, is_display: %s, input_size: %s, "
                     "output_shapes: %s, coverage_ratio: %s, config: %s)",
                     self.__class__.__name__, batchsize, is_display, input_size, output_shapes,
                     coverage_ratio, config)

        self.initialized = False
        self.is_display = is_display

        # Set on first image load from initialize
        self._training_size = 0
        self._constants = None

 </s> add     def __init__(self,
                 batchsize: int,
                 processing_size: int,
                 config: "ConfigType") -> None:
        logger.debug("Initializing %s: (batchsize: %s, processing_size: %s, "
                     "config: %s)",
                     self.__class__.__name__, batchsize, processing_size, config)

        self._processing_size = processing_size </s> remove     is_display: bool
        Whether the images being fed through will be used for Preview or Time-lapse. Disables
        the "warp" augmentation for these images.
    input_size: int
        The expected input size for the model. It is assumed that the input to the model is always
        a square image. This is the size, in pixels, of the `width` and the `height` of the input
        to the model.
    output_shapes: list
        A list of tuples defining the output shapes from the model, in the order that the outputs
        are returned. The tuples should be in (`height`, `width`, `channels`) format.
    coverage_ratio: float
        The ratio of the training image to be trained on. Dictates how much of the image will be
        cropped out. E.G: a coverage ratio of 0.625 will result in cropping a 160px box from a
        256px image (:math:`256 * 0.625 = 160`)
 </s> add     processing_size: int
        The largest input or output size of the model. This is the size that images are processed
        at. </s> remove     def __init__(self, model_input_size, model_output_shapes, coverage_ratio, color_order,
                 augment_color, no_flip, no_warp, warp_to_landmarks, config):
        logger.debug("Initializing %s: (model_input_size: %s, model_output_shapes: %s, "
                     "coverage_ratio: %s, color_order: %s, augment_color: %s, no_flip: %s, "
                     "no_warp: %s, warp_to_landmarks: %s, config: %s)",
                     self.__class__.__name__, model_input_size, model_output_shapes,
                     coverage_ratio, color_order, augment_color, no_flip, no_warp,
                     warp_to_landmarks, config)
        self._config = config
        self._model_input_size = model_input_size
        self._model_output_shapes = model_output_shapes
        self._coverage_ratio = coverage_ratio
        self._color_order = color_order.lower()
        self._augment_color = augment_color
        self._no_flip = no_flip
        self._warp_to_landmarks = warp_to_landmarks
        self._no_warp = no_warp

        # Batchsize and processing class are set when this class is called by a feeder
        # from lib.training_data
        self._batchsize = 0
        self._face_cache = None
        self._nearest_landmarks = dict()
        self._processing = None
        logger.debug("Initialized %s", self.__class__.__name__)
 </s> add     def __init__(self,
                 config: ConfigType,
                 model: "ModelBase",
                 side: Literal["a", "b"],
                 images: List[str],
                 batch_size: int) -> None:
        super().__init__(config, model, side, images, batch_size)
        self._augment_color = not model.command_line_arguments.no_augment_color
        self._no_flip = model.command_line_arguments.no_flip
        self._no_warp = model.command_line_arguments.no_warp
        self._warp_to_landmarks = (not self._no_warp
                                   and model.command_line_arguments.warp_to_landmarks)
        self._model_input_size = model.input_shape[1] </s> remove     def __init__(self, images, model, batch_size, config):
 </s> add     def __init__(self,
                 images: Dict[Literal["a", "b"], List[str]],
                 model: 'ModelBase',
                 batch_size: int,
                 config: ConfigType) -> None: </s> remove     model_input_size: int
        The expected input size for the model. It is assumed that the input to the model is always
        a square image. This is the size, in pixels, of the `width` and the `height` of the input
        to the model.
    model_output_shapes: list
        A list of tuples defining the output shapes from the model, in the order that the outputs
        are returned. The tuples should be in (`height`, `width`, `channels`) format.
    coverage_ratio: float
        The ratio of the training image to be trained on. Dictates how much of the image will be
        cropped out. E.G: a coverage ratio of 0.625 will result in cropping a 160px box from a
        256px image (:math:`256 * 0.625 = 160`).
    color_order: ["rgb", "bgr"]
        The color order that the model expects as input
    augment_color: bool
        ``True`` if color is to be augmented, otherwise ``False``
    no_flip: bool
        ``True`` if the image shouldn't be randomly flipped as part of augmentation, otherwise
        ``False``
    no_warp: bool
        ``True`` if the image shouldn't be warped as part of augmentation, otherwise ``False``
    warp_to_landmarks: bool
        ``True`` if the random warp method should warp to similar landmarks from the other side,
        ``False`` if the standard random warp method should be used.
    face_cache: dict
        A thread safe dictionary containing a cache of information relating to all faces being
        trained on
 </s> add     model: :class:`~plugins.train.model.ModelBase`
        The model that this data generator is feeding
          256px image (:math:`256 * 0.625 = 160`)
      config: dict
          The configuration `dict` generated from :file:`config.train.ini` containing the trainer
          plugin configuration options.









      """
      def __init__(self, batchsize, is_display, input_size, output_shapes, coverage_ratio, config):
          logger.debug("Initializing %s: (batchsize: %s, is_display: %s, input_size: %s, "
                       "output_shapes: %s, coverage_ratio: %s, config: %s)",
                       self.__class__.__name__, batchsize, is_display, input_size, output_shapes,
keep keep keep keep replace replace replace replace replace replace replace replace replace keep keep keep keep keep
 Data Augmentation update (#1263)

- lib.detected_face
    - Subclass Masks for Landmark based masks
    - Add training mask propery + methods to DetectedFace
  - lib.training_training
    - subclass TrainingDataGenerator for training and preview data
    - Split cache into own module
    - Reduce thread count to 1 to prevent image corruption + data re-use
    - Process on largest model input/output size rather than stored image size
    - Size and crop masks during caching stage
    - Implement ring buffer for data flow
    - Fix preview reload bug
  - augmentation
    - typing
    - switch color aug order
    - better initialization
    - Fix warp + landmark warp to correctly apply at different image scales
    - Slightly improved warp caching
    - Don't store whether image is_preview. Handle all data as training images implicitly
  - plugins.trainer: Typing and fixes to work with trainingdata refactor 
lib/training/augmentation.py