StableDiffusion/WebUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

improve filename matching for mask

Browse Source 
we should not rely that mask filename will be of the same extension
as the image filename so better pattern matching is added
This commit is contained in:
Artem Kotov 2023-05-29 20:39:24 +04:00
parent 20ae71faa8
commit c8e67b6732
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
modules/img2img.py
View File

@ -1,4 +1,5 @@
import os import os
from pathlib import Path
import numpy as np import numpy as np
from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, UnidentifiedImageError 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: if is_inpaint_batch:
# try to find corresponding mask for an image using simple filename matching # 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 not found use first one ("same mask for all images" use-case)
if mask_image_path not in inpaint_masks: if mask_image_path not in inpaint_masks:
mask_image_path = inpaint_masks[0] mask_image_path = inpaint_masks[0]