| import os | |
| from glob import glob | |
| from rembg import remove | |
| from argparse import ArgumentParser | |
| from PIL import Image | |
| if __name__ == '__main__': | |
| parser = ArgumentParser() | |
| parser.add_argument('--path', type=str, required=True, help='Path to input images') | |
| args = parser.parse_args() | |
| imgs = glob(os.path.join(args.path, '*.png')) + glob(os.path.join(args.path, '*.jpg')) | |
| for img in imgs: | |
| path = os.path.dirname(img) | |
| name = os.path.basename(img).split('.')[0] + '_rmbg.png' | |
| img_np = Image.open(img) | |
| img = remove(img_np) | |
| img.save(os.path.join(args.path, name)) |