Remove background from images using rembg

remove_bg.py

import logging

import rembg
from tqdm.contrib.concurrent import process_map
import os
import pathlib


def filter_files(files: list[pathlib.Path], target_suffix: str):
    files_map = set(files)
    return [file for file in files if file.suffix == '.jpg' and not file.with_suffix(target_suffix) in files_map]


def get_list_of_files(input_dir: str, target_suffix: str):
    files = [pathlib.Path(f"{input_dir}/{f}") for f in os.listdir(input_dir)]
    only_missing = filter_files(files, target_suffix=target_suffix)
    logging.error(f"{len(only_missing)=} {len(files)=}")
    return only_missing


def remove_background(image_path: pathlib.Path):
    input_data = image_path.open('rb').read()
    output = rembg.remove(input_data)
    with open(image_path.with_suffix('.png'), 'wb') as o:
        o.write(output)


def remove_backgrounds(input_dir: str, target_suffix: str):
    inputs = get_list_of_files(input_dir, target_suffix=target_suffix)
    process_map(remove_background, inputs, max_workers=3, chunksize=1)


if __name__ == "__main__":
    remove_backgrounds('../zigbee2mqtt.io/public/images/devices', target_suffix='.png')