[Bugfix] Inpaint masked image

configs/_base_/datasets/dog_inpaint_multiple_mask.py

@@ -58,6 +58,7 @@
     dict(type="DumpImage", max_imgs=10, dump_dir="work_dirs/dump"),
     dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
     dict(type="GetMaskedImage"),
+    dict(type="DumpMaskedImage", max_imgs=10, dump_dir="work_dirs/dump"),
     dict(type="PackInputs",
          input_keys=["img", "mask", "masked_image", "text"]),
 ]

diffengine/datasets/transforms/__init__.py

@@ -1,5 +1,5 @@
 from .base import BaseTransform
-from .dump_image import DumpImage
+from .dump_image import DumpImage, DumpMaskedImage
 from .formatting import PackInputs
 from .loading import LoadMask
 from .processing import (
@@ -40,4 +40,5 @@
     "GetMaskedImage",
     "RandomChoice",
     "AddConstantCaption",
+    "DumpMaskedImage",
 ]

diffengine/datasets/transforms/dump_image.py

@@ -70,3 +70,53 @@ def __call__(self, results) -> dict:
                             mask.numpy().astype(np.uint8))
 
         return results
+
+
+@TRANSFORMS.register_module()
+class DumpMaskedImage:
+    """Dump Masked the image processed by the pipeline.
+
+    Args:
+    ----
+        max_imgs (int): Maximum value of output.
+        dump_dir (str): Dump output directory.
+    """
+
+    def __init__(self, max_imgs: int, dump_dir: str) -> None:
+        self.max_imgs = max_imgs
+        self.dump_dir = dump_dir
+        mmengine.mkdir_or_exist(self.dump_dir)
+        self.num_dumped_imgs = Value("i", 0)
+
+    def __call__(self, results) -> dict:
+        """Dump the input image to the specified directory.
+
+        No changes will be
+        made.
+
+        Args:
+        ----
+            results (dict): Result dict from loading pipeline.
+
+        Returns:
+        -------
+            results (dict): Result dict from loading pipeline. (same as input)
+        """
+        enable_dump = False
+        with self.num_dumped_imgs.get_lock():
+            if self.num_dumped_imgs.value < self.max_imgs:
+                self.num_dumped_imgs.value += 1
+                enable_dump = True
+                dump_id = self.num_dumped_imgs.value
+
+        if enable_dump:
+            masked_image = results["masked_image"]
+            masked_image = (masked_image / 2 + 0.5).clamp(0, 1)
+            if masked_image.shape[0] in [1, 3]:
+                masked_image = masked_image.permute(1, 2, 0) * 255
+            masked_image_out_file = osp.join(
+                self.dump_dir, f"{dump_id}_masked_image.png")
+            cv2.imwrite(masked_image_out_file,
+                        masked_image.numpy().astype(np.uint8)[..., ::-1])
+
+        return results

diffengine/datasets/transforms/processing.py

@@ -685,7 +685,8 @@ def transform(self, results: dict) -> dict | tuple[list, list] | None:
         ----
             results (dict): The result dict.
         """
-        results[self.key] = results["img"] * results["mask"]
+        mask_threahold = 0.5
+        results[self.key] = results["img"] * (results["mask"] < mask_threahold)
         return results
 
 

tests/test_datasets/test_transforms/test_processing.py

@@ -490,8 +490,8 @@ def test_register(self):
     def test_transform(self):
         img_path = osp.join(osp.dirname(__file__), "../../testdata/color.jpg")
         img = torch.Tensor(np.array(Image.open(img_path)))
-        mask = np.ones((img.shape[0], img.shape[1], 1))
-        mask[:10, :10] = 0
+        mask = np.zeros((img.shape[0], img.shape[1], 1))
+        mask[:10, :10] = 1
         mask = torch.Tensor(mask)
         data = {"img": img, "mask": mask}