diff --git a/plugins/train/trainer/_base.py b/plugins/train/trainer/_base.py index 11120022c0..dbbc74c9fd 100644 --- a/plugins/train/trainer/_base.py +++ b/plugins/train/trainer/_base.py @@ -672,24 +672,23 @@ def _compile_masked(self, faces: list[np.ndarray], masks: np.ndarray) -> list[np list List of :class:`numpy.ndarray` faces with the opaque mask layer applied """ - orig_masks = 1 - np.rint(masks) + orig_masks = 1. - masks masks3: list[np.ndarray] | np.ndarray = [] if faces[-1].shape[-1] == 4: # Mask contained in alpha channel of predictions - pred_masks = [1 - np.rint(face[..., -1])[..., None] for face in faces[-2:]] + pred_masks = [1. - face[..., -1][..., None] for face in faces[-2:]] faces[-2:] = [face[..., :-1] for face in faces[-2:]] masks3 = [orig_masks, *pred_masks] else: masks3 = np.repeat(np.expand_dims(orig_masks, axis=0), 3, axis=0) retval: list[np.ndarray] = [] - alpha = 1.0 - self._mask_opacity - for previews, compiled_masks in zip(faces, masks3): - overlays = previews.copy() - overlays[np.where((compiled_masks == 1.).all(axis=3))] = self._mask_color - retval.append(np.array([cv2.addWeighted(img, alpha, ovl, self._mask_opacity, 0) - for img, ovl in zip(previews, overlays)])) - + overlays3 = np.ones_like(faces) * self._mask_color + for previews, overlays, compiled_masks in zip(faces, overlays3, masks3): + compiled_masks *= self._mask_opacity + overlays *= compiled_masks + previews *= (1. - compiled_masks) + retval.append(previews + overlays) logger.debug("masked shapes: %s", [faces.shape for faces in retval]) return retval