Training preview: Correctly display blur/kernel amount on mask

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