From c8e67b67320f5d090b758303e675c1e5586575a5 Mon Sep 17 00:00:00 2001 From: Artem Kotov Date: Mon, 29 May 2023 20:39:24 +0400 Subject: [PATCH] improve filename matching for mask we should not rely that mask filename will be of the same extension as the image filename so better pattern matching is added --- modules/img2img.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modules/img2img.py b/modules/img2img.py index d704bf90..bc79ea1f 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import numpy as np from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, UnidentifiedImageError @@ -53,7 +54,11 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args): if is_inpaint_batch: # try to find corresponding mask for an image using simple filename matching - mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image)) + path = Path(os.path.join(inpaint_mask_dir, os.path.basename(image))) + mask_image_path = list(path.parent.glob(f"**/{path.stem}*")) + if len(mask_image_path) > 0: + mask_image_path = str(mask_image_path[0]) + # if not found use first one ("same mask for all images" use-case) if mask_image_path not in inpaint_masks: mask_image_path = inpaint_masks[0]