Add mask dilation/erosion option for training

deepfakes · Mar 26, 2024 · 4502824 · 4502824
1 parent 2b08b8d
commit 4502824
Show file tree

Hide file tree

Showing 6 changed files with 156 additions and 78 deletions.
diff --git a/lib/align/detected_face.py b/lib/align/detected_face.py
diff --git a/lib/training/cache.py b/lib/training/cache.py
@@ -428,8 +428,10 @@ def _get_face_mask(self, filename: str, detected_face: DetectedFace) -> np.ndarr
                 f"The masks that exist for this face are: {list(detected_face.mask)}")
 
         mask = detected_face.mask[str(self._config["mask_type"])]
+        assert isinstance(self._config["mask_dilation"], float)
         assert isinstance(self._config["mask_blur_kernel"], int)
         assert isinstance(self._config["mask_threshold"], int)
+        mask.set_dilation(self._config["mask_dilation"])
         mask.set_blur_and_threshold(blur_kernel=self._config["mask_blur_kernel"],
                                     threshold=self._config["mask_threshold"])
 
@@ -467,7 +469,7 @@ def _get_localized_mask(self,
         assert isinstance(multiplier, int)
         if not self._config["penalized_mask_loss"] or multiplier <= 1:
             return None
-        mask = detected_face.get_landmark_mask(area, self._size // 16, self._size // 32)
+        mask = detected_face.get_landmark_mask(area, self._size // 16, 2.5)
         logger.trace("Caching localized '%s' mask for: %s %s",  # type: ignore
                      area, filename, mask.shape)
         return mask

diff --git a/locales/plugins.train._config.pot b/locales/plugins.train._config.pot
@@ -8,7 +8,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2023-08-20 14:54+0100\n"
+"POT-Creation-Date: 2024-03-26 17:37+0000\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <[email protected]>\n"
@@ -539,8 +539,9 @@ msgid ""
 "attention on the core face area."
 msgstr ""
 
-#: plugins/train/_config.py:600 plugins/train/_config.py:642
-#: plugins/train/_config.py:656 plugins/train/_config.py:665
+#: plugins/train/_config.py:600 plugins/train/_config.py:643
+#: plugins/train/_config.py:656 plugins/train/_config.py:671
+#: plugins/train/_config.py:680
 msgid "mask"
 msgstr ""
 
@@ -582,7 +583,14 @@ msgid ""
 "performance."
 msgstr ""
 
-#: plugins/train/_config.py:644
+#: plugins/train/_config.py:645
+msgid ""
+"Dilate or erode the mask. Negative values erode the mask (make it smaller). "
+"Positive values dilate the mask (make it larger). The value given is a "
+"percentage of the total mask size."
+msgstr ""
+
+#: plugins/train/_config.py:658
 msgid ""
 "Apply gaussian blur to the mask input. This has the effect of smoothing the "
 "edges of the mask, which can help with poorly calculated masks and give less "
@@ -592,13 +600,13 @@ msgid ""
 "number."
 msgstr ""
 
-#: plugins/train/_config.py:658
+#: plugins/train/_config.py:673
 msgid ""
 "Sets pixels that are near white to white and near black to black. Set to 0 "
 "for off."
 msgstr ""
 
-#: plugins/train/_config.py:667
+#: plugins/train/_config.py:682
 msgid ""
 "Dedicate a portion of the model to learning how to duplicate the input mask. "
 "Increases VRAM usage in exchange for learning a quick ability to try to "

diff --git a/locales/ru/LC_MESSAGES/plugins.train._config.mo b/locales/ru/LC_MESSAGES/plugins.train._config.mo
diff --git a/locales/ru/LC_MESSAGES/plugins.train._config.po b/locales/ru/LC_MESSAGES/plugins.train._config.po
@@ -7,15 +7,15 @@ msgid ""
 msgstr ""
 "Project-Id-Version: \n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2023-08-20 14:54+0100\n"
-"PO-Revision-Date: 2023-08-20 14:58+0100\n"
+"POT-Creation-Date: 2024-03-26 17:37+0000\n"
+"PO-Revision-Date: 2024-03-26 17:40+0000\n"
 "Last-Translator: \n"
 "Language-Team: \n"
 "Language: ru_RU\n"
 "MIME-Version: 1.0\n"
 "Content-Type: text/plain; charset=UTF-8\n"
 "Content-Transfer-Encoding: 8bit\n"
-"X-Generator: Poedit 3.3.2\n"
+"X-Generator: Poedit 3.4.2\n"
 
 #: plugins/train/_config.py:17
 msgid ""
@@ -851,8 +851,9 @@ msgstr ""
 "время как область лица с маской является приоритетной. Может повысить общее "
 "качество за счет концентрации внимания на основной области лица."
 
-#: plugins/train/_config.py:600 plugins/train/_config.py:642
-#: plugins/train/_config.py:656 plugins/train/_config.py:665
+#: plugins/train/_config.py:600 plugins/train/_config.py:643
+#: plugins/train/_config.py:656 plugins/train/_config.py:671
+#: plugins/train/_config.py:680
 msgid "mask"
 msgstr "маска"
 
@@ -928,7 +929,16 @@ msgstr ""
 "сообщества и для дальнейшего описания нуждается в тестировании. Профильные "
 "лица могут иметь низкую производительность."
 
-#: plugins/train/_config.py:644
+#: plugins/train/_config.py:645
+msgid ""
+"Dilate or erode the mask. Negative values erode the mask (make it smaller). "
+"Positive values dilate the mask (make it larger). The value given is a "
+"percentage of the total mask size."
+msgstr ""
+"Расширяет или сужает маску. Отрицательные значения сужают маску (делают её "
+"меньше). Положительные значения расширяют маску (делают её больше). "
+
+#: plugins/train/_config.py:658
 msgid ""
 "Apply gaussian blur to the mask input. This has the effect of smoothing the "
 "edges of the mask, which can help with poorly calculated masks and give less "
@@ -944,15 +954,15 @@ msgstr ""
 "должно быть нечетным, если передано четное число, то оно будет округлено до "
 "следующего нечетного числа."
 
-#: plugins/train/_config.py:658
+#: plugins/train/_config.py:673
 msgid ""
 "Sets pixels that are near white to white and near black to black. Set to 0 "
 "for off."
 msgstr ""
 "Устанавливает пиксели, которые почти белые - в белые и которые почти черные "
 "- в черные. Установите 0, чтобы выключить."
 
-#: plugins/train/_config.py:667
+#: plugins/train/_config.py:682
 msgid ""
 "Dedicate a portion of the model to learning how to duplicate the input mask. "
 "Increases VRAM usage in exchange for learning a quick ability to try to "

diff --git a/plugins/train/_config.py b/plugins/train/_config.py
@@ -632,13 +632,27 @@ def _set_loss(self) -> None:
                 "faces. The mask model has been trained by community members and will need "
                 "testing for further description. Profile faces may result in sub-par "
                 "performance."))
+        self.add_item(
+            section=section,
+            title="mask_dilation",
+            datatype=float,
+            min_max=(-5.0, 5.0),
+            rounding=1,
+            default=0,
+            fixed=False,
+            group=_("mask"),
+            info=_(
+                "Dilate or erode the mask. Negative values erode the mask (make it smaller). "
+                "Positive values dilate the mask (make it larger). The value given is a "
+                "percentage of the total mask size."))
         self.add_item(
             section=section,
             title="mask_blur_kernel",
             datatype=int,
             min_max=(0, 9),
             rounding=1,
             default=3,
+            fixed=False,
             group=_("mask"),
             info=_(
                 "Apply gaussian blur to the mask input. This has the effect of smoothing the "
@@ -653,6 +667,7 @@ def _set_loss(self) -> None:
             default=4,
             min_max=(0, 50),
             rounding=1,
+            fixed=False,
             group=_("mask"),
             info=_(
                 "Sets pixels that are near white to white and near black to black. Set to 0 for "