fix match_histograms

2024-01-08 23:38:18 +08:00
parent be2f721c9d
commit 7c1f83e71d
2 changed files with 39 additions and 0 deletions
--- a/iopaint/model/base.py
+++ b/iopaint/model/base.py
@@ -211,6 +211,9 @@ class InpaintModel:

    def _match_histograms(self, source, reference, mask):
        transformed_channels = []
+        if len(mask.shape) == 3:
+            mask = mask[:, :, -1]
+
        for channel in range(source.shape[-1]):
            source_channel = source[:, :, channel]
            reference_channel = reference[:, :, channel]