commit d29a5883425567f81d7147a535ee71e607148b44 Author: kwaroran Date: Mon Jul 31 11:13:08 2023 +0900 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6769e21 --- /dev/null +++ b/.gitignore @@ -0,0 +1,160 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6f7ebc9 --- /dev/null +++ b/__init__.py @@ -0,0 +1,80 @@ +import huggingface_hub +import torch +import onnxruntime as rt +import numpy as np +import cv2 + +def get_mask(img:torch.Tensor, s=1024): + img = (img / 255).astype(np.float32) + h, w = h0, w0 = img.shape[:-1] + h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s) + ph, pw = s - h, s - w + img_input = np.zeros([s, s, 3], dtype=np.float32) + img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h)) + img_input = np.transpose(img_input, (2, 0, 1)) + img_input = img_input[np.newaxis, :] + mask = rmbg_model.run(None, {'img': img_input})[0][0] + mask = np.transpose(mask, (1, 2, 0)) + mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] + mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis] + return mask + +# Declare Execution Providers +providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] + +# Download and host the model +model_path = huggingface_hub.hf_hub_download( + "skytnt/anime-seg", "isnetis.onnx") +rmbg_model = rt.InferenceSession(model_path, providers=providers) + +def rmbg_fn(img): + mask = get_mask(img) + img = (mask * img + 255 * (1 - mask)).astype(np.uint8) + mask = (mask * 255).astype(np.uint8) + img = np.concatenate([img, mask], axis=2, dtype=np.uint8) + mask = mask.repeat(3, axis=2) + return img + +class RemoveImageBackgroundARB: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "arb_remover" + CATEGORY = "image" + + def arb_remover(self, image:torch.Tensor): + npa = image2nparray(image) + print(npa.ndim) + rmb = rmbg_fn(npa) + image = nparray2image(rmb) + return (image,) + +def image2nparray(image:torch.Tensor): + narray:np.array = np.clip(255. * image.cpu().numpy().squeeze(),0, 255).astype(np.uint8) + if narray.shape[-1] == 4: + narray = narray[..., [2, 1, 0, 3]] # For RGBA + else: + narray = narray[..., [2, 1, 0]] # For RGB + return narray + +def nparray2image(narray:np.array): + print(f"narray shape: {narray.shape}") + if narray.shape[-1] == 4: + narray = narray[..., [2, 1, 0, 3]] + else: + narray = narray[..., [2, 1, 0]] + tensor = torch.from_numpy(narray/255.).float().unsqueeze(0) + return tensor + +NODE_CLASS_MAPPINGS = { + "Remove Image Background (ARB)": RemoveImageBackgroundARB +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..466c8b6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +opencv-python +onnxruntime +onnx +numpy +Pillow \ No newline at end of file