Commit c8e67b67 authored by Artem Kotov's avatar Artem Kotov

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
parent 20ae71fa
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): ...@@ -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]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment