              if input_shape[0] / faces.shape[1] != 1.0:
                  feeds[side] = self._resize_sample(side, faces, input_shape[0])
                  feeds[side] = feeds[side].reshape((-1, ) + input_shape)
              else:
                  feeds[side] = faces
  
          preds = self._get_predictions(feeds["a"], feeds["b"])
  
          for side, samples in self.images.items():
              other_side = "a" if side == "b" else "b"
 
</s> remove             if input_shape[0] / faces.shape[1] != 1.0:
                feeds[side] = self._resize_sample(side, faces, input_shape[0])
 </s> add             if input_shape[0] / self.images[side][0].shape[1] != 1.0:
                feeds[side] = self._resize_sample(side, self.images[side][1], input_shape[0]) </s> remove         feeds = {}
        figures = {}
        headers = {}
        for idx, side in enumerate(("a", "b")):
            samples = self.images[side]
            faces = samples[1]
 </s> add         feeds: Dict[Literal["a", "b"], np.ndarray] = {}
        for idx, side in enumerate(get_args(Literal["a", "b"])): </s> remove 
        for side, samples in self.images.items():
            other_side = "a" if side == "b" else "b"
            predictions = [preds[f"{side}_{side}"],
                           preds[f"{other_side}_{side}"]]
            display = self._to_full_frame(side, samples, predictions)
            headers[side] = self._get_headers(side, display[0].shape[1])
            figures[side] = np.stack([display[0], display[1], display[2], ], axis=1)
            if self.images[side][0].shape[0] % 2 == 1:
                figures[side] = np.concatenate([figures[side],
                                                np.expand_dims(figures[side][0], 0)])

        width = 4
        side_cols = width // 2
        if side_cols != 1:
            headers = self._duplicate_headers(headers, side_cols)

        header = np.concatenate([headers["a"], headers["b"]], axis=1)
        figure = np.concatenate([figures["a"], figures["b"]], axis=0)
        height = int(figure.shape[0] / width)
        figure = figure.reshape((width, height) + figure.shape[1:])
        figure = _stack_images(figure)
        figure = np.concatenate((header, figure), axis=0)

        logger.debug("Compiled sample")
        return np.clip(figure * 255, 0, 255).astype('uint8')
 </s> add         return self._compile_preview(preds) </s> remove     def _to_full_frame(self, side, samples, predictions):
 </s> add     def _compile_preview(self, predictions: Dict[str, np.ndarray]) -> np.ndarray:
        """ Compile predictions and images into the final preview image.

        Parameters
        ----------
        predictions: dict
            The predictions from the model

        Returns
        -------
        :class:`numpy.ndarry`
            A compiled preview image ready for display or saving
        """
        figures: Dict[Literal["a", "b"], np.ndarray] = {}
        headers: Dict[Literal["a", "b"], np.ndarray] = {}

        for side, samples in self.images.items():
            other_side = "a" if side == "b" else "b"
            preds = [predictions[f"{side}_{side}"],
                     predictions[f"{other_side}_{side}"]]
            display = self._to_full_frame(side, samples, preds)
            headers[side] = self._get_headers(side, display[0].shape[1])
            figures[side] = np.stack([display[0], display[1], display[2], ], axis=1)
            if self.images[side][1].shape[0] % 2 == 1:
                figures[side] = np.concatenate([figures[side],
                                                np.expand_dims(figures[side][0], 0)])

        width = 4
        if width // 2 != 1:
            headers = self._duplicate_headers(headers, width // 2)

        header = np.concatenate([headers["a"], headers["b"]], axis=1)
        figure = np.concatenate([figures["a"], figures["b"]], axis=0)
        height = int(figure.shape[0] / width)
        figure = figure.reshape((width, height) + figure.shape[1:])
        figure = _stack_images(figure)
        figure = np.concatenate((header, figure), axis=0)

        logger.debug("Compiled sample")
        return np.clip(figure * 255, 0, 255).astype('uint8')

    def _to_full_frame(self,
                       side: Literal["a", "b"],
                       samples: List[np.ndarray],
                       predictions: List[np.ndarray]) -> List[np.ndarray]: </s> remove         full, faces = samples[:2]
 </s> add         faces, full = samples[:2]
                feeds[side] = self.images[side][0]
              if input_shape[0] / faces.shape[1] != 1.0:
                  feeds[side] = self._resize_sample(side, faces, input_shape[0])
                  feeds[side] = feeds[side].reshape((-1, ) + input_shape)
              else:
                feeds[side] = self.images[side][0]
  
          preds = self._get_predictions(feeds["a"], feeds["b"])
  
          for side, samples in self.images.items():
              other_side = "a" if side == "b" else "b"
keep keep keep keep 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 
plugins/train/trainer/_base.py