Masking images using rembg

[ghost]

Resizes the image for faster processing.

from argparse import ArgumentParser
from concurrent.futures import ProcessPoolExecutor
from glob import glob
import os
from pyvips import Image
from rembg import remove, new_session

def maskImage(image_path, suffix, max_res, out_dir):
  img = Image.new_from_file(image_path, access='sequential')
  
  image_name, image_ext = os.path.splitext(os.path.basename(image_path))
  if suffix == 0:
    filename = image_name + '_mask.png'
  elif suffix == 1:
    filename = image_name + '.mask.png'
  elif suffix == 2:
    filename = image_name + image_ext + '.png'
  elif suffix == 3:
    filename = image_name + image_ext + '.mask.png'
  elif suffix == 4:
    filename = image_name + '_mask.png'
  
  img_max, img_ratio = max(img.width, img.height), img.width / img.height
  ratio = 0
  while not ratio:
    scale = max_res / img_max
    if img_ratio == round(img.width * scale) / round(img.height * scale):
      ratio = 1
    else:
      max_res -= 1
  
  img = Image.new_from_array(remove(img.resize(scale, kernel='linear').numpy(), only_mask=True, post_process_mask=True, session=new_session('u2net')))
  img.resize(1 / scale, kernel='nearest').write_to_file(os.path.join(out_dir, filename), Q=100)

if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('-i', '--input', type=str, required=True, help='path to the input directory')
  parser.add_argument('-o', '--output', type=str, required=True, help='path to the output directory')
  parser.add_argument('-s', '--suffix', type=int, required=True, help='filename suffix (0 - OpenMVG, 1 - OpenMVS, 2 - COLMAP, 3 - RealityCapture, 4 - 3DF Zephyr)')
  parser.add_argument('--max_resolution', type=int, default=1024, help='max image resolution')
  args = parser.parse_args()
  
  os.makedirs(args.output, exist_ok = True)
  
  image_list = []
  {image_list.extend(glob(os.path.join(args.input, f'*.{file_extension}'))) for file_extension in ['jpg', 'jpeg', 'png']}
  
  with ProcessPoolExecutor(max_workers=int(os.cpu_count() / 2)) as executor:
    for image_path in image_list:
      executor.submit(maskImage, image_path, args.suffix, args.max_resolution, args.output)