RemBG
Build-Deploy-Actions
Details
Build-Deploy-Actions
Details
This commit is contained in:
commit
049b836f04
|
@ -0,0 +1,27 @@
|
||||||
|
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.model filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||||
|
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.zstandard filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
@ -0,0 +1,47 @@
|
||||||
|
name: Build
|
||||||
|
run-name: ${{ github.actor }} is upgrade release 🚀
|
||||||
|
on: [push]
|
||||||
|
env:
|
||||||
|
REPOSITORY: ${{ github.repository }}
|
||||||
|
COMMIT_ID: ${{ github.sha }}
|
||||||
|
jobs:
|
||||||
|
Build-Deploy-Actions:
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
steps:
|
||||||
|
- run: echo "🎉 The job was automatically triggered by a ${{ github.event_name }} event."
|
||||||
|
- run: echo "🐧 This job is now running on a ${{ runner.os }} server hosted by Gitea!"
|
||||||
|
- run: echo "🔎 The name of your branch is ${{ github.ref }} and your repository is ${{ github.repository }}."
|
||||||
|
- name: Check out repository code
|
||||||
|
uses: actions/checkout@v3
|
||||||
|
-
|
||||||
|
name: Setup Git LFS
|
||||||
|
run: |
|
||||||
|
git lfs install
|
||||||
|
git lfs fetch
|
||||||
|
git lfs checkout
|
||||||
|
- name: List files in the repository
|
||||||
|
run: |
|
||||||
|
ls ${{ github.workspace }}
|
||||||
|
-
|
||||||
|
name: Docker Image Info
|
||||||
|
id: image-info
|
||||||
|
run: |
|
||||||
|
echo "::set-output name=image_name::$(echo $REPOSITORY | tr '[:upper:]' '[:lower:]')"
|
||||||
|
echo "::set-output name=image_tag::${COMMIT_ID:0:10}"
|
||||||
|
-
|
||||||
|
name: Login to Docker Hub
|
||||||
|
uses: docker/login-action@v2
|
||||||
|
with:
|
||||||
|
registry: artifacts.iflytek.com
|
||||||
|
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||||
|
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||||
|
- name: Set up Docker Buildx
|
||||||
|
uses: docker/setup-buildx-action@v2
|
||||||
|
-
|
||||||
|
name: Build and push
|
||||||
|
run: |
|
||||||
|
docker version
|
||||||
|
docker buildx build -t artifacts.iflytek.com/docker-private/atp/${{ steps.image-info.outputs.image_name }}:${{ steps.image-info.outputs.image_tag }} . --file ${{ github.workspace }}/Dockerfile --load
|
||||||
|
docker push artifacts.iflytek.com/docker-private/atp/${{ steps.image-info.outputs.image_name }}:${{ steps.image-info.outputs.image_tag }}
|
||||||
|
docker rmi artifacts.iflytek.com/docker-private/atp/${{ steps.image-info.outputs.image_name }}:${{ steps.image-info.outputs.image_tag }}
|
||||||
|
- run: echo "🍏 This job's status is ${{ job.status }}."
|
|
@ -0,0 +1,13 @@
|
||||||
|
#FROM python:3.8.13
|
||||||
|
FROM artifacts.iflytek.com/docker-private/atp/base_image_for_ailab:0.0.1
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
COPY . /app
|
||||||
|
|
||||||
|
RUN pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple
|
||||||
|
RUN pip install -r requirements.txt
|
||||||
|
|
||||||
|
RUN mkdir -p /root/.u2net && mv u2net.onnx /root/.u2net/
|
||||||
|
|
||||||
|
CMD ["python", "app.py"]
|
|
@ -0,0 +1,21 @@
|
||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2020 Daniel Gatis
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
|
@ -0,0 +1,12 @@
|
||||||
|
---
|
||||||
|
title: Rembg
|
||||||
|
emoji: 👀
|
||||||
|
colorFrom: pink
|
||||||
|
colorTo: indigo
|
||||||
|
sdk: gradio
|
||||||
|
sdk_version: 3.0.20
|
||||||
|
app_file: app.py
|
||||||
|
pinned: false
|
||||||
|
---
|
||||||
|
|
||||||
|
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
@ -0,0 +1,67 @@
|
||||||
|
## Modified from Akhaliq Hugging Face Demo
|
||||||
|
## https://huggingface.co/akhaliq
|
||||||
|
|
||||||
|
import gradio as gr
|
||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
def inference(file, af, mask):
|
||||||
|
im = cv2.imread(file, cv2.IMREAD_COLOR)
|
||||||
|
cv2.imwrite(os.path.join("input.png"), im)
|
||||||
|
|
||||||
|
model = "u2net"
|
||||||
|
from rembg import remove
|
||||||
|
from rembg.session_base import BaseSession
|
||||||
|
from rembg.session_factory import new_session
|
||||||
|
|
||||||
|
input_path = 'input.png'
|
||||||
|
output_path = 'output.png'
|
||||||
|
|
||||||
|
with open(input_path, 'rb') as i:
|
||||||
|
with open(output_path, 'wb') as o:
|
||||||
|
input = i.read()
|
||||||
|
sessions: dict[str, BaseSession] = {}
|
||||||
|
output = remove(
|
||||||
|
input,
|
||||||
|
session=sessions.setdefault(
|
||||||
|
model, new_session(model)
|
||||||
|
),
|
||||||
|
alpha_matting_erode_size = af,
|
||||||
|
only_mask = (True if mask == "Mask only" else False)
|
||||||
|
)
|
||||||
|
o.write(output)
|
||||||
|
return os.path.join("output.png")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
gr.Interface(
|
||||||
|
inference,
|
||||||
|
[
|
||||||
|
gr.inputs.Image(type="filepath", label="Input"),
|
||||||
|
gr.inputs.Slider(10, 25, default=10, label="Alpha matting"),
|
||||||
|
gr.inputs.Radio(
|
||||||
|
[
|
||||||
|
"Default",
|
||||||
|
"Mask only"
|
||||||
|
],
|
||||||
|
type="value",
|
||||||
|
default="Default",
|
||||||
|
label="Choices"
|
||||||
|
),
|
||||||
|
#gr.inputs.Dropdown([
|
||||||
|
# "u2net",
|
||||||
|
# "u2netp",
|
||||||
|
# "u2net_human_seg",
|
||||||
|
# "u2net_cloth_seg",
|
||||||
|
# "silueta"
|
||||||
|
# ],
|
||||||
|
# type="value",
|
||||||
|
# default="u2net",
|
||||||
|
# label="Models"
|
||||||
|
#),
|
||||||
|
],
|
||||||
|
# gr.outputs.Image(type="filepath", label="Output"),
|
||||||
|
gr.Image(),
|
||||||
|
examples=[["lion.png", 10, "Default", "u2net"], ["girl.jpg", 10, "Default", "u2net"]],
|
||||||
|
enable_queue=True
|
||||||
|
).launch(server_name="0.0.0.0" , share=True)
|
|
@ -0,0 +1,6 @@
|
||||||
|
from . import _version
|
||||||
|
|
||||||
|
__version__ = _version.get_versions()["version"]
|
||||||
|
|
||||||
|
from .bg import remove
|
||||||
|
from .session_factory import new_session
|
|
@ -0,0 +1,677 @@
|
||||||
|
# This file helps to compute a version number in source trees obtained from
|
||||||
|
# git-archive tarball (such as those provided by githubs download-from-tag
|
||||||
|
# feature). Distribution tarballs (built by setup.py sdist) and build
|
||||||
|
# directories (produced by setup.py build) will contain a much shorter file
|
||||||
|
# that just contains the computed version number.
|
||||||
|
|
||||||
|
# This file is released into the public domain. Generated by
|
||||||
|
# versioneer-0.21 (https://github.com/python-versioneer/python-versioneer)
|
||||||
|
|
||||||
|
"""Git implementation of _version.py."""
|
||||||
|
|
||||||
|
import errno
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
from typing import Callable, Dict
|
||||||
|
|
||||||
|
|
||||||
|
def get_keywords():
|
||||||
|
"""Get the keywords needed to look up the version information."""
|
||||||
|
# these strings will be replaced by git during git-archive.
|
||||||
|
# setup.py/versioneer.py will grep for the variable names, so they must
|
||||||
|
# each be defined on a line of their own. _version.py will just call
|
||||||
|
# get_keywords().
|
||||||
|
git_refnames = " (HEAD -> main)"
|
||||||
|
git_full = "d62227d5866e2178e88f06074917484a4424082e"
|
||||||
|
git_date = "2022-12-10 11:51:49 -0300"
|
||||||
|
keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
|
||||||
|
return keywords
|
||||||
|
|
||||||
|
|
||||||
|
class VersioneerConfig:
|
||||||
|
"""Container for Versioneer configuration parameters."""
|
||||||
|
|
||||||
|
|
||||||
|
def get_config():
|
||||||
|
"""Create, populate and return the VersioneerConfig() object."""
|
||||||
|
# these strings are filled in when 'setup.py versioneer' creates
|
||||||
|
# _version.py
|
||||||
|
cfg = VersioneerConfig()
|
||||||
|
cfg.VCS = "git"
|
||||||
|
cfg.style = "pep440"
|
||||||
|
cfg.tag_prefix = "v"
|
||||||
|
cfg.parentdir_prefix = "rembg-"
|
||||||
|
cfg.versionfile_source = "rembg/_version.py"
|
||||||
|
cfg.verbose = False
|
||||||
|
return cfg
|
||||||
|
|
||||||
|
|
||||||
|
class NotThisMethod(Exception):
|
||||||
|
"""Exception raised if a method is not valid for the current scenario."""
|
||||||
|
|
||||||
|
|
||||||
|
LONG_VERSION_PY: Dict[str, str] = {}
|
||||||
|
HANDLERS: Dict[str, Dict[str, Callable]] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def register_vcs_handler(vcs, method): # decorator
|
||||||
|
"""Create decorator to mark a method as the handler of a VCS."""
|
||||||
|
|
||||||
|
def decorate(f):
|
||||||
|
"""Store f in HANDLERS[vcs][method]."""
|
||||||
|
if vcs not in HANDLERS:
|
||||||
|
HANDLERS[vcs] = {}
|
||||||
|
HANDLERS[vcs][method] = f
|
||||||
|
return f
|
||||||
|
|
||||||
|
return decorate
|
||||||
|
|
||||||
|
|
||||||
|
def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None):
|
||||||
|
"""Call the given command(s)."""
|
||||||
|
assert isinstance(commands, list)
|
||||||
|
process = None
|
||||||
|
for command in commands:
|
||||||
|
try:
|
||||||
|
dispcmd = str([command] + args)
|
||||||
|
# remember shell=False, so use git.cmd on windows, not just git
|
||||||
|
process = subprocess.Popen(
|
||||||
|
[command] + args,
|
||||||
|
cwd=cwd,
|
||||||
|
env=env,
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=(subprocess.PIPE if hide_stderr else None),
|
||||||
|
)
|
||||||
|
break
|
||||||
|
except OSError:
|
||||||
|
e = sys.exc_info()[1]
|
||||||
|
if e.errno == errno.ENOENT:
|
||||||
|
continue
|
||||||
|
if verbose:
|
||||||
|
print("unable to run %s" % dispcmd)
|
||||||
|
print(e)
|
||||||
|
return None, None
|
||||||
|
else:
|
||||||
|
if verbose:
|
||||||
|
print("unable to find command, tried %s" % (commands,))
|
||||||
|
return None, None
|
||||||
|
stdout = process.communicate()[0].strip().decode()
|
||||||
|
if process.returncode != 0:
|
||||||
|
if verbose:
|
||||||
|
print("unable to run %s (error)" % dispcmd)
|
||||||
|
print("stdout was %s" % stdout)
|
||||||
|
return None, process.returncode
|
||||||
|
return stdout, process.returncode
|
||||||
|
|
||||||
|
|
||||||
|
def versions_from_parentdir(parentdir_prefix, root, verbose):
|
||||||
|
"""Try to determine the version from the parent directory name.
|
||||||
|
|
||||||
|
Source tarballs conventionally unpack into a directory that includes both
|
||||||
|
the project name and a version string. We will also support searching up
|
||||||
|
two directory levels for an appropriately named parent directory
|
||||||
|
"""
|
||||||
|
rootdirs = []
|
||||||
|
|
||||||
|
for _ in range(3):
|
||||||
|
dirname = os.path.basename(root)
|
||||||
|
if dirname.startswith(parentdir_prefix):
|
||||||
|
return {
|
||||||
|
"version": dirname[len(parentdir_prefix) :],
|
||||||
|
"full-revisionid": None,
|
||||||
|
"dirty": False,
|
||||||
|
"error": None,
|
||||||
|
"date": None,
|
||||||
|
}
|
||||||
|
rootdirs.append(root)
|
||||||
|
root = os.path.dirname(root) # up a level
|
||||||
|
|
||||||
|
if verbose:
|
||||||
|
print(
|
||||||
|
"Tried directories %s but none started with prefix %s"
|
||||||
|
% (str(rootdirs), parentdir_prefix)
|
||||||
|
)
|
||||||
|
raise NotThisMethod("rootdir doesn't start with parentdir_prefix")
|
||||||
|
|
||||||
|
|
||||||
|
@register_vcs_handler("git", "get_keywords")
|
||||||
|
def git_get_keywords(versionfile_abs):
|
||||||
|
"""Extract version information from the given file."""
|
||||||
|
# the code embedded in _version.py can just fetch the value of these
|
||||||
|
# keywords. When used from setup.py, we don't want to import _version.py,
|
||||||
|
# so we do it with a regexp instead. This function is not used from
|
||||||
|
# _version.py.
|
||||||
|
keywords = {}
|
||||||
|
try:
|
||||||
|
with open(versionfile_abs, "r") as fobj:
|
||||||
|
for line in fobj:
|
||||||
|
if line.strip().startswith("git_refnames ="):
|
||||||
|
mo = re.search(r'=\s*"(.*)"', line)
|
||||||
|
if mo:
|
||||||
|
keywords["refnames"] = mo.group(1)
|
||||||
|
if line.strip().startswith("git_full ="):
|
||||||
|
mo = re.search(r'=\s*"(.*)"', line)
|
||||||
|
if mo:
|
||||||
|
keywords["full"] = mo.group(1)
|
||||||
|
if line.strip().startswith("git_date ="):
|
||||||
|
mo = re.search(r'=\s*"(.*)"', line)
|
||||||
|
if mo:
|
||||||
|
keywords["date"] = mo.group(1)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
return keywords
|
||||||
|
|
||||||
|
|
||||||
|
@register_vcs_handler("git", "keywords")
|
||||||
|
def git_versions_from_keywords(keywords, tag_prefix, verbose):
|
||||||
|
"""Get version information from git keywords."""
|
||||||
|
if "refnames" not in keywords:
|
||||||
|
raise NotThisMethod("Short version file found")
|
||||||
|
date = keywords.get("date")
|
||||||
|
if date is not None:
|
||||||
|
# Use only the last line. Previous lines may contain GPG signature
|
||||||
|
# information.
|
||||||
|
date = date.splitlines()[-1]
|
||||||
|
|
||||||
|
# git-2.2.0 added "%cI", which expands to an ISO-8601 -compliant
|
||||||
|
# datestamp. However we prefer "%ci" (which expands to an "ISO-8601
|
||||||
|
# -like" string, which we must then edit to make compliant), because
|
||||||
|
# it's been around since git-1.5.3, and it's too difficult to
|
||||||
|
# discover which version we're using, or to work around using an
|
||||||
|
# older one.
|
||||||
|
date = date.strip().replace(" ", "T", 1).replace(" ", "", 1)
|
||||||
|
refnames = keywords["refnames"].strip()
|
||||||
|
if refnames.startswith("$Format"):
|
||||||
|
if verbose:
|
||||||
|
print("keywords are unexpanded, not using")
|
||||||
|
raise NotThisMethod("unexpanded keywords, not a git-archive tarball")
|
||||||
|
refs = {r.strip() for r in refnames.strip("()").split(",")}
|
||||||
|
# starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of
|
||||||
|
# just "foo-1.0". If we see a "tag: " prefix, prefer those.
|
||||||
|
TAG = "tag: "
|
||||||
|
tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)}
|
||||||
|
if not tags:
|
||||||
|
# Either we're using git < 1.8.3, or there really are no tags. We use
|
||||||
|
# a heuristic: assume all version tags have a digit. The old git %d
|
||||||
|
# expansion behaves like git log --decorate=short and strips out the
|
||||||
|
# refs/heads/ and refs/tags/ prefixes that would let us distinguish
|
||||||
|
# between branches and tags. By ignoring refnames without digits, we
|
||||||
|
# filter out many common branch names like "release" and
|
||||||
|
# "stabilization", as well as "HEAD" and "master".
|
||||||
|
tags = {r for r in refs if re.search(r"\d", r)}
|
||||||
|
if verbose:
|
||||||
|
print("discarding '%s', no digits" % ",".join(refs - tags))
|
||||||
|
if verbose:
|
||||||
|
print("likely tags: %s" % ",".join(sorted(tags)))
|
||||||
|
for ref in sorted(tags):
|
||||||
|
# sorting will prefer e.g. "2.0" over "2.0rc1"
|
||||||
|
if ref.startswith(tag_prefix):
|
||||||
|
r = ref[len(tag_prefix) :]
|
||||||
|
# Filter out refs that exactly match prefix or that don't start
|
||||||
|
# with a number once the prefix is stripped (mostly a concern
|
||||||
|
# when prefix is '')
|
||||||
|
if not re.match(r"\d", r):
|
||||||
|
continue
|
||||||
|
if verbose:
|
||||||
|
print("picking %s" % r)
|
||||||
|
return {
|
||||||
|
"version": r,
|
||||||
|
"full-revisionid": keywords["full"].strip(),
|
||||||
|
"dirty": False,
|
||||||
|
"error": None,
|
||||||
|
"date": date,
|
||||||
|
}
|
||||||
|
# no suitable tags, so version is "0+unknown", but full hex is still there
|
||||||
|
if verbose:
|
||||||
|
print("no suitable tags, using unknown + full revision id")
|
||||||
|
return {
|
||||||
|
"version": "0+unknown",
|
||||||
|
"full-revisionid": keywords["full"].strip(),
|
||||||
|
"dirty": False,
|
||||||
|
"error": "no suitable tags",
|
||||||
|
"date": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@register_vcs_handler("git", "pieces_from_vcs")
|
||||||
|
def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
|
||||||
|
"""Get version from 'git describe' in the root of the source tree.
|
||||||
|
|
||||||
|
This only gets called if the git-archive 'subst' keywords were *not*
|
||||||
|
expanded, and _version.py hasn't already been rewritten with a short
|
||||||
|
version string, meaning we're inside a checked out source tree.
|
||||||
|
"""
|
||||||
|
GITS = ["git"]
|
||||||
|
TAG_PREFIX_REGEX = "*"
|
||||||
|
if sys.platform == "win32":
|
||||||
|
GITS = ["git.cmd", "git.exe"]
|
||||||
|
TAG_PREFIX_REGEX = r"\*"
|
||||||
|
|
||||||
|
_, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=True)
|
||||||
|
if rc != 0:
|
||||||
|
if verbose:
|
||||||
|
print("Directory %s not under git control" % root)
|
||||||
|
raise NotThisMethod("'git rev-parse --git-dir' returned error")
|
||||||
|
|
||||||
|
# if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty]
|
||||||
|
# if there isn't one, this yields HEX[-dirty] (no NUM)
|
||||||
|
describe_out, rc = runner(
|
||||||
|
GITS,
|
||||||
|
[
|
||||||
|
"describe",
|
||||||
|
"--tags",
|
||||||
|
"--dirty",
|
||||||
|
"--always",
|
||||||
|
"--long",
|
||||||
|
"--match",
|
||||||
|
"%s%s" % (tag_prefix, TAG_PREFIX_REGEX),
|
||||||
|
],
|
||||||
|
cwd=root,
|
||||||
|
)
|
||||||
|
# --long was added in git-1.5.5
|
||||||
|
if describe_out is None:
|
||||||
|
raise NotThisMethod("'git describe' failed")
|
||||||
|
describe_out = describe_out.strip()
|
||||||
|
full_out, rc = runner(GITS, ["rev-parse", "HEAD"], cwd=root)
|
||||||
|
if full_out is None:
|
||||||
|
raise NotThisMethod("'git rev-parse' failed")
|
||||||
|
full_out = full_out.strip()
|
||||||
|
|
||||||
|
pieces = {}
|
||||||
|
pieces["long"] = full_out
|
||||||
|
pieces["short"] = full_out[:7] # maybe improved later
|
||||||
|
pieces["error"] = None
|
||||||
|
|
||||||
|
branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root)
|
||||||
|
# --abbrev-ref was added in git-1.6.3
|
||||||
|
if rc != 0 or branch_name is None:
|
||||||
|
raise NotThisMethod("'git rev-parse --abbrev-ref' returned error")
|
||||||
|
branch_name = branch_name.strip()
|
||||||
|
|
||||||
|
if branch_name == "HEAD":
|
||||||
|
# If we aren't exactly on a branch, pick a branch which represents
|
||||||
|
# the current commit. If all else fails, we are on a branchless
|
||||||
|
# commit.
|
||||||
|
branches, rc = runner(GITS, ["branch", "--contains"], cwd=root)
|
||||||
|
# --contains was added in git-1.5.4
|
||||||
|
if rc != 0 or branches is None:
|
||||||
|
raise NotThisMethod("'git branch --contains' returned error")
|
||||||
|
branches = branches.split("\n")
|
||||||
|
|
||||||
|
# Remove the first line if we're running detached
|
||||||
|
if "(" in branches[0]:
|
||||||
|
branches.pop(0)
|
||||||
|
|
||||||
|
# Strip off the leading "* " from the list of branches.
|
||||||
|
branches = [branch[2:] for branch in branches]
|
||||||
|
if "master" in branches:
|
||||||
|
branch_name = "master"
|
||||||
|
elif not branches:
|
||||||
|
branch_name = None
|
||||||
|
else:
|
||||||
|
# Pick the first branch that is returned. Good or bad.
|
||||||
|
branch_name = branches[0]
|
||||||
|
|
||||||
|
pieces["branch"] = branch_name
|
||||||
|
|
||||||
|
# parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty]
|
||||||
|
# TAG might have hyphens.
|
||||||
|
git_describe = describe_out
|
||||||
|
|
||||||
|
# look for -dirty suffix
|
||||||
|
dirty = git_describe.endswith("-dirty")
|
||||||
|
pieces["dirty"] = dirty
|
||||||
|
if dirty:
|
||||||
|
git_describe = git_describe[: git_describe.rindex("-dirty")]
|
||||||
|
|
||||||
|
# now we have TAG-NUM-gHEX or HEX
|
||||||
|
|
||||||
|
if "-" in git_describe:
|
||||||
|
# TAG-NUM-gHEX
|
||||||
|
mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe)
|
||||||
|
if not mo:
|
||||||
|
# unparsable. Maybe git-describe is misbehaving?
|
||||||
|
pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out
|
||||||
|
return pieces
|
||||||
|
|
||||||
|
# tag
|
||||||
|
full_tag = mo.group(1)
|
||||||
|
if not full_tag.startswith(tag_prefix):
|
||||||
|
if verbose:
|
||||||
|
fmt = "tag '%s' doesn't start with prefix '%s'"
|
||||||
|
print(fmt % (full_tag, tag_prefix))
|
||||||
|
pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % (
|
||||||
|
full_tag,
|
||||||
|
tag_prefix,
|
||||||
|
)
|
||||||
|
return pieces
|
||||||
|
pieces["closest-tag"] = full_tag[len(tag_prefix) :]
|
||||||
|
|
||||||
|
# distance: number of commits since tag
|
||||||
|
pieces["distance"] = int(mo.group(2))
|
||||||
|
|
||||||
|
# commit: short hex revision ID
|
||||||
|
pieces["short"] = mo.group(3)
|
||||||
|
|
||||||
|
else:
|
||||||
|
# HEX: no tags
|
||||||
|
pieces["closest-tag"] = None
|
||||||
|
count_out, rc = runner(GITS, ["rev-list", "HEAD", "--count"], cwd=root)
|
||||||
|
pieces["distance"] = int(count_out) # total number of commits
|
||||||
|
|
||||||
|
# commit date: see ISO-8601 comment in git_versions_from_keywords()
|
||||||
|
date = runner(GITS, ["show", "-s", "--format=%ci", "HEAD"], cwd=root)[0].strip()
|
||||||
|
# Use only the last line. Previous lines may contain GPG signature
|
||||||
|
# information.
|
||||||
|
date = date.splitlines()[-1]
|
||||||
|
pieces["date"] = date.strip().replace(" ", "T", 1).replace(" ", "", 1)
|
||||||
|
|
||||||
|
return pieces
|
||||||
|
|
||||||
|
|
||||||
|
def plus_or_dot(pieces):
|
||||||
|
"""Return a + if we don't already have one, else return a ."""
|
||||||
|
if "+" in pieces.get("closest-tag", ""):
|
||||||
|
return "."
|
||||||
|
return "+"
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440(pieces):
|
||||||
|
"""Build up version string, with post-release "local version identifier".
|
||||||
|
|
||||||
|
Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you
|
||||||
|
get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty]
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"] or pieces["dirty"]:
|
||||||
|
rendered += plus_or_dot(pieces)
|
||||||
|
rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440_branch(pieces):
|
||||||
|
"""TAG[[.dev0]+DISTANCE.gHEX[.dirty]] .
|
||||||
|
|
||||||
|
The ".dev0" means not master branch. Note that .dev0 sorts backwards
|
||||||
|
(a feature branch will appear "older" than the master branch).
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty]
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"] or pieces["dirty"]:
|
||||||
|
if pieces["branch"] != "master":
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += plus_or_dot(pieces)
|
||||||
|
rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0"
|
||||||
|
if pieces["branch"] != "master":
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def pep440_split_post(ver):
|
||||||
|
"""Split pep440 version string at the post-release segment.
|
||||||
|
|
||||||
|
Returns the release segments before the post-release and the
|
||||||
|
post-release version number (or -1 if no post-release segment is present).
|
||||||
|
"""
|
||||||
|
vc = str.split(ver, ".post")
|
||||||
|
return vc[0], int(vc[1] or 0) if len(vc) == 2 else None
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440_pre(pieces):
|
||||||
|
"""TAG[.postN.devDISTANCE] -- No -dirty.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. 0.post0.devDISTANCE
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
if pieces["distance"]:
|
||||||
|
# update the post release segment
|
||||||
|
tag_version, post_version = pep440_split_post(pieces["closest-tag"])
|
||||||
|
rendered = tag_version
|
||||||
|
if post_version is not None:
|
||||||
|
rendered += ".post%d.dev%d" % (post_version + 1, pieces["distance"])
|
||||||
|
else:
|
||||||
|
rendered += ".post0.dev%d" % (pieces["distance"])
|
||||||
|
else:
|
||||||
|
# no commits, use the tag as the version
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0.post0.dev%d" % pieces["distance"]
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440_post(pieces):
|
||||||
|
"""TAG[.postDISTANCE[.dev0]+gHEX] .
|
||||||
|
|
||||||
|
The ".dev0" means dirty. Note that .dev0 sorts backwards
|
||||||
|
(a dirty tree will appear "older" than the corresponding clean one),
|
||||||
|
but you shouldn't be releasing software with -dirty anyways.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. 0.postDISTANCE[.dev0]
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"] or pieces["dirty"]:
|
||||||
|
rendered += ".post%d" % pieces["distance"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += plus_or_dot(pieces)
|
||||||
|
rendered += "g%s" % pieces["short"]
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0.post%d" % pieces["distance"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += "+g%s" % pieces["short"]
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440_post_branch(pieces):
|
||||||
|
"""TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] .
|
||||||
|
|
||||||
|
The ".dev0" means not master branch.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty]
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"] or pieces["dirty"]:
|
||||||
|
rendered += ".post%d" % pieces["distance"]
|
||||||
|
if pieces["branch"] != "master":
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += plus_or_dot(pieces)
|
||||||
|
rendered += "g%s" % pieces["short"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0.post%d" % pieces["distance"]
|
||||||
|
if pieces["branch"] != "master":
|
||||||
|
rendered += ".dev0"
|
||||||
|
rendered += "+g%s" % pieces["short"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dirty"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_pep440_old(pieces):
|
||||||
|
"""TAG[.postDISTANCE[.dev0]] .
|
||||||
|
|
||||||
|
The ".dev0" means dirty.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. 0.postDISTANCE[.dev0]
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"] or pieces["dirty"]:
|
||||||
|
rendered += ".post%d" % pieces["distance"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dev0"
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = "0.post%d" % pieces["distance"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += ".dev0"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_git_describe(pieces):
|
||||||
|
"""TAG[-DISTANCE-gHEX][-dirty].
|
||||||
|
|
||||||
|
Like 'git describe --tags --dirty --always'.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. HEX[-dirty] (note: no 'g' prefix)
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
if pieces["distance"]:
|
||||||
|
rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = pieces["short"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += "-dirty"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render_git_describe_long(pieces):
|
||||||
|
"""TAG-DISTANCE-gHEX[-dirty].
|
||||||
|
|
||||||
|
Like 'git describe --tags --dirty --always -long'.
|
||||||
|
The distance/hash is unconditional.
|
||||||
|
|
||||||
|
Exceptions:
|
||||||
|
1: no tags. HEX[-dirty] (note: no 'g' prefix)
|
||||||
|
"""
|
||||||
|
if pieces["closest-tag"]:
|
||||||
|
rendered = pieces["closest-tag"]
|
||||||
|
rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
|
||||||
|
else:
|
||||||
|
# exception #1
|
||||||
|
rendered = pieces["short"]
|
||||||
|
if pieces["dirty"]:
|
||||||
|
rendered += "-dirty"
|
||||||
|
return rendered
|
||||||
|
|
||||||
|
|
||||||
|
def render(pieces, style):
|
||||||
|
"""Render the given version pieces into the requested style."""
|
||||||
|
if pieces["error"]:
|
||||||
|
return {
|
||||||
|
"version": "unknown",
|
||||||
|
"full-revisionid": pieces.get("long"),
|
||||||
|
"dirty": None,
|
||||||
|
"error": pieces["error"],
|
||||||
|
"date": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
if not style or style == "default":
|
||||||
|
style = "pep440" # the default
|
||||||
|
|
||||||
|
if style == "pep440":
|
||||||
|
rendered = render_pep440(pieces)
|
||||||
|
elif style == "pep440-branch":
|
||||||
|
rendered = render_pep440_branch(pieces)
|
||||||
|
elif style == "pep440-pre":
|
||||||
|
rendered = render_pep440_pre(pieces)
|
||||||
|
elif style == "pep440-post":
|
||||||
|
rendered = render_pep440_post(pieces)
|
||||||
|
elif style == "pep440-post-branch":
|
||||||
|
rendered = render_pep440_post_branch(pieces)
|
||||||
|
elif style == "pep440-old":
|
||||||
|
rendered = render_pep440_old(pieces)
|
||||||
|
elif style == "git-describe":
|
||||||
|
rendered = render_git_describe(pieces)
|
||||||
|
elif style == "git-describe-long":
|
||||||
|
rendered = render_git_describe_long(pieces)
|
||||||
|
else:
|
||||||
|
raise ValueError("unknown style '%s'" % style)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"version": rendered,
|
||||||
|
"full-revisionid": pieces["long"],
|
||||||
|
"dirty": pieces["dirty"],
|
||||||
|
"error": None,
|
||||||
|
"date": pieces.get("date"),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def get_versions():
|
||||||
|
"""Get version information or return default if unable to do so."""
|
||||||
|
# I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have
|
||||||
|
# __file__, we can work backwards from there to the root. Some
|
||||||
|
# py2exe/bbfreeze/non-CPython implementations don't do __file__, in which
|
||||||
|
# case we can only use expanded keywords.
|
||||||
|
|
||||||
|
cfg = get_config()
|
||||||
|
verbose = cfg.verbose
|
||||||
|
|
||||||
|
try:
|
||||||
|
return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose)
|
||||||
|
except NotThisMethod:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
root = os.path.realpath(__file__)
|
||||||
|
# versionfile_source is the relative path from the top of the source
|
||||||
|
# tree (where the .git directory might live) to this file. Invert
|
||||||
|
# this to find the root from __file__.
|
||||||
|
for _ in cfg.versionfile_source.split("/"):
|
||||||
|
root = os.path.dirname(root)
|
||||||
|
except NameError:
|
||||||
|
return {
|
||||||
|
"version": "0+unknown",
|
||||||
|
"full-revisionid": None,
|
||||||
|
"dirty": None,
|
||||||
|
"error": "unable to find root of source tree",
|
||||||
|
"date": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose)
|
||||||
|
return render(pieces, cfg.style)
|
||||||
|
except NotThisMethod:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
if cfg.parentdir_prefix:
|
||||||
|
return versions_from_parentdir(cfg.parentdir_prefix, root, verbose)
|
||||||
|
except NotThisMethod:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return {
|
||||||
|
"version": "0+unknown",
|
||||||
|
"full-revisionid": None,
|
||||||
|
"dirty": None,
|
||||||
|
"error": "unable to compute version",
|
||||||
|
"date": None,
|
||||||
|
}
|
|
@ -0,0 +1,176 @@
|
||||||
|
import io
|
||||||
|
from enum import Enum
|
||||||
|
from typing import List, Optional, Union
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from cv2 import (
|
||||||
|
BORDER_DEFAULT,
|
||||||
|
MORPH_ELLIPSE,
|
||||||
|
MORPH_OPEN,
|
||||||
|
GaussianBlur,
|
||||||
|
getStructuringElement,
|
||||||
|
morphologyEx,
|
||||||
|
)
|
||||||
|
from PIL import Image
|
||||||
|
from PIL.Image import Image as PILImage
|
||||||
|
from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
|
||||||
|
from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
|
||||||
|
from pymatting.util.util import stack_images
|
||||||
|
from scipy.ndimage import binary_erosion
|
||||||
|
|
||||||
|
from .session_base import BaseSession
|
||||||
|
from .session_factory import new_session
|
||||||
|
|
||||||
|
kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
|
||||||
|
|
||||||
|
|
||||||
|
class ReturnType(Enum):
|
||||||
|
BYTES = 0
|
||||||
|
PILLOW = 1
|
||||||
|
NDARRAY = 2
|
||||||
|
|
||||||
|
|
||||||
|
def alpha_matting_cutout(
|
||||||
|
img: PILImage,
|
||||||
|
mask: PILImage,
|
||||||
|
foreground_threshold: int,
|
||||||
|
background_threshold: int,
|
||||||
|
erode_structure_size: int,
|
||||||
|
) -> PILImage:
|
||||||
|
|
||||||
|
if img.mode == "RGBA" or img.mode == "CMYK":
|
||||||
|
img = img.convert("RGB")
|
||||||
|
|
||||||
|
img = np.asarray(img)
|
||||||
|
mask = np.asarray(mask)
|
||||||
|
|
||||||
|
is_foreground = mask > foreground_threshold
|
||||||
|
is_background = mask < background_threshold
|
||||||
|
|
||||||
|
structure = None
|
||||||
|
if erode_structure_size > 0:
|
||||||
|
structure = np.ones(
|
||||||
|
(erode_structure_size, erode_structure_size), dtype=np.uint8
|
||||||
|
)
|
||||||
|
|
||||||
|
is_foreground = binary_erosion(is_foreground, structure=structure)
|
||||||
|
is_background = binary_erosion(is_background, structure=structure, border_value=1)
|
||||||
|
|
||||||
|
trimap = np.full(mask.shape, dtype=np.uint8, fill_value=128)
|
||||||
|
trimap[is_foreground] = 255
|
||||||
|
trimap[is_background] = 0
|
||||||
|
|
||||||
|
img_normalized = img / 255.0
|
||||||
|
trimap_normalized = trimap / 255.0
|
||||||
|
|
||||||
|
alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
|
||||||
|
foreground = estimate_foreground_ml(img_normalized, alpha)
|
||||||
|
cutout = stack_images(foreground, alpha)
|
||||||
|
|
||||||
|
cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
|
||||||
|
cutout = Image.fromarray(cutout)
|
||||||
|
|
||||||
|
return cutout
|
||||||
|
|
||||||
|
|
||||||
|
def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
|
||||||
|
empty = Image.new("RGBA", (img.size), 0)
|
||||||
|
cutout = Image.composite(img, empty, mask)
|
||||||
|
return cutout
|
||||||
|
|
||||||
|
|
||||||
|
def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
|
||||||
|
pivot = imgs.pop(0)
|
||||||
|
for im in imgs:
|
||||||
|
pivot = get_concat_v(pivot, im)
|
||||||
|
return pivot
|
||||||
|
|
||||||
|
|
||||||
|
def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
|
||||||
|
dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
|
||||||
|
dst.paste(img1, (0, 0))
|
||||||
|
dst.paste(img2, (0, img1.height))
|
||||||
|
return dst
|
||||||
|
|
||||||
|
|
||||||
|
def post_process(mask: np.ndarray) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Post Process the mask for a smooth boundary by applying Morphological Operations
|
||||||
|
Research based on paper: https://www.sciencedirect.com/science/article/pii/S2352914821000757
|
||||||
|
args:
|
||||||
|
mask: Binary Numpy Mask
|
||||||
|
"""
|
||||||
|
mask = morphologyEx(mask, MORPH_OPEN, kernel)
|
||||||
|
mask = GaussianBlur(mask, (5, 5), sigmaX=2, sigmaY=2, borderType=BORDER_DEFAULT)
|
||||||
|
mask = np.where(mask < 127, 0, 255).astype(np.uint8) # convert again to binary
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
def remove(
|
||||||
|
data: Union[bytes, PILImage, np.ndarray],
|
||||||
|
alpha_matting: bool = False,
|
||||||
|
alpha_matting_foreground_threshold: int = 240,
|
||||||
|
alpha_matting_background_threshold: int = 10,
|
||||||
|
alpha_matting_erode_size: int = 10,
|
||||||
|
session: Optional[BaseSession] = None,
|
||||||
|
only_mask: bool = False,
|
||||||
|
post_process_mask: bool = False,
|
||||||
|
) -> Union[bytes, PILImage, np.ndarray]:
|
||||||
|
|
||||||
|
if isinstance(data, PILImage):
|
||||||
|
return_type = ReturnType.PILLOW
|
||||||
|
img = data
|
||||||
|
elif isinstance(data, bytes):
|
||||||
|
return_type = ReturnType.BYTES
|
||||||
|
img = Image.open(io.BytesIO(data))
|
||||||
|
elif isinstance(data, np.ndarray):
|
||||||
|
return_type = ReturnType.NDARRAY
|
||||||
|
img = Image.fromarray(data)
|
||||||
|
else:
|
||||||
|
raise ValueError("Input type {} is not supported.".format(type(data)))
|
||||||
|
|
||||||
|
if session is None:
|
||||||
|
session = new_session("u2net")
|
||||||
|
|
||||||
|
masks = session.predict(img)
|
||||||
|
cutouts = []
|
||||||
|
|
||||||
|
for mask in masks:
|
||||||
|
if post_process_mask:
|
||||||
|
mask = Image.fromarray(post_process(np.array(mask)))
|
||||||
|
|
||||||
|
if only_mask:
|
||||||
|
cutout = mask
|
||||||
|
|
||||||
|
elif alpha_matting:
|
||||||
|
try:
|
||||||
|
cutout = alpha_matting_cutout(
|
||||||
|
img,
|
||||||
|
mask,
|
||||||
|
alpha_matting_foreground_threshold,
|
||||||
|
alpha_matting_background_threshold,
|
||||||
|
alpha_matting_erode_size,
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
cutout = naive_cutout(img, mask)
|
||||||
|
|
||||||
|
else:
|
||||||
|
cutout = naive_cutout(img, mask)
|
||||||
|
|
||||||
|
cutouts.append(cutout)
|
||||||
|
|
||||||
|
cutout = img
|
||||||
|
if len(cutouts) > 0:
|
||||||
|
cutout = get_concat_v_multi(cutouts)
|
||||||
|
|
||||||
|
if ReturnType.PILLOW == return_type:
|
||||||
|
return cutout
|
||||||
|
|
||||||
|
if ReturnType.NDARRAY == return_type:
|
||||||
|
return np.asarray(cutout)
|
||||||
|
|
||||||
|
bio = io.BytesIO()
|
||||||
|
cutout.save(bio, "PNG")
|
||||||
|
bio.seek(0)
|
||||||
|
|
||||||
|
return bio.read()
|
|
@ -0,0 +1,440 @@
|
||||||
|
import pathlib
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from enum import Enum
|
||||||
|
from typing import IO, cast
|
||||||
|
|
||||||
|
import aiohttp
|
||||||
|
import click
|
||||||
|
import filetype
|
||||||
|
import uvicorn
|
||||||
|
from asyncer import asyncify
|
||||||
|
from fastapi import Depends, FastAPI, File, Form, Query
|
||||||
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
from starlette.responses import Response
|
||||||
|
from tqdm import tqdm
|
||||||
|
from watchdog.events import FileSystemEvent, FileSystemEventHandler
|
||||||
|
from watchdog.observers import Observer
|
||||||
|
|
||||||
|
from . import _version
|
||||||
|
from .bg import remove
|
||||||
|
from .session_base import BaseSession
|
||||||
|
from .session_factory import new_session
|
||||||
|
|
||||||
|
|
||||||
|
@click.group()
|
||||||
|
@click.version_option(version=_version.get_versions()["version"])
|
||||||
|
def main() -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@main.command(help="for a file as input")
|
||||||
|
@click.option(
|
||||||
|
"-m",
|
||||||
|
"--model",
|
||||||
|
default="u2net",
|
||||||
|
type=click.Choice(
|
||||||
|
["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
|
||||||
|
),
|
||||||
|
show_default=True,
|
||||||
|
show_choices=True,
|
||||||
|
help="model name",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-a",
|
||||||
|
"--alpha-matting",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="use alpha matting",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-af",
|
||||||
|
"--alpha-matting-foreground-threshold",
|
||||||
|
default=240,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="trimap fg threshold",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ab",
|
||||||
|
"--alpha-matting-background-threshold",
|
||||||
|
default=10,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="trimap bg threshold",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ae",
|
||||||
|
"--alpha-matting-erode-size",
|
||||||
|
default=10,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="erode size",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-om",
|
||||||
|
"--only-mask",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="output only the mask",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ppm",
|
||||||
|
"--post-process-mask",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="post process the mask",
|
||||||
|
)
|
||||||
|
@click.argument(
|
||||||
|
"input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
|
||||||
|
)
|
||||||
|
@click.argument(
|
||||||
|
"output",
|
||||||
|
default=(None if sys.stdin.isatty() else "-"),
|
||||||
|
type=click.File("wb", lazy=True),
|
||||||
|
)
|
||||||
|
def i(model: str, input: IO, output: IO, **kwargs) -> None:
|
||||||
|
output.write(remove(input.read(), session=new_session(model), **kwargs))
|
||||||
|
|
||||||
|
|
||||||
|
@main.command(help="for a folder as input")
|
||||||
|
@click.option(
|
||||||
|
"-m",
|
||||||
|
"--model",
|
||||||
|
default="u2net",
|
||||||
|
type=click.Choice(
|
||||||
|
["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
|
||||||
|
),
|
||||||
|
show_default=True,
|
||||||
|
show_choices=True,
|
||||||
|
help="model name",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-a",
|
||||||
|
"--alpha-matting",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="use alpha matting",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-af",
|
||||||
|
"--alpha-matting-foreground-threshold",
|
||||||
|
default=240,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="trimap fg threshold",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ab",
|
||||||
|
"--alpha-matting-background-threshold",
|
||||||
|
default=10,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="trimap bg threshold",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ae",
|
||||||
|
"--alpha-matting-erode-size",
|
||||||
|
default=10,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="erode size",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-om",
|
||||||
|
"--only-mask",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="output only the mask",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-ppm",
|
||||||
|
"--post-process-mask",
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="post process the mask",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-w",
|
||||||
|
"--watch",
|
||||||
|
default=False,
|
||||||
|
is_flag=True,
|
||||||
|
show_default=True,
|
||||||
|
help="watches a folder for changes",
|
||||||
|
)
|
||||||
|
@click.argument(
|
||||||
|
"input",
|
||||||
|
type=click.Path(
|
||||||
|
exists=True,
|
||||||
|
path_type=pathlib.Path,
|
||||||
|
file_okay=False,
|
||||||
|
dir_okay=True,
|
||||||
|
readable=True,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
@click.argument(
|
||||||
|
"output",
|
||||||
|
type=click.Path(
|
||||||
|
exists=False,
|
||||||
|
path_type=pathlib.Path,
|
||||||
|
file_okay=False,
|
||||||
|
dir_okay=True,
|
||||||
|
writable=True,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
def p(
|
||||||
|
model: str, input: pathlib.Path, output: pathlib.Path, watch: bool, **kwargs
|
||||||
|
) -> None:
|
||||||
|
session = new_session(model)
|
||||||
|
|
||||||
|
def process(each_input: pathlib.Path) -> None:
|
||||||
|
try:
|
||||||
|
mimetype = filetype.guess(each_input)
|
||||||
|
if mimetype is None:
|
||||||
|
return
|
||||||
|
if mimetype.mime.find("image") < 0:
|
||||||
|
return
|
||||||
|
|
||||||
|
each_output = (output / each_input.name).with_suffix(".png")
|
||||||
|
each_output.parents[0].mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
if not each_output.exists():
|
||||||
|
each_output.write_bytes(
|
||||||
|
cast(
|
||||||
|
bytes,
|
||||||
|
remove(each_input.read_bytes(), session=session, **kwargs),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if watch:
|
||||||
|
print(
|
||||||
|
f"processed: {each_input.absolute()} -> {each_output.absolute()}"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
|
||||||
|
inputs = list(input.glob("**/*"))
|
||||||
|
if not watch:
|
||||||
|
inputs = tqdm(inputs)
|
||||||
|
|
||||||
|
for each_input in inputs:
|
||||||
|
if not each_input.is_dir():
|
||||||
|
process(each_input)
|
||||||
|
|
||||||
|
if watch:
|
||||||
|
observer = Observer()
|
||||||
|
|
||||||
|
class EventHandler(FileSystemEventHandler):
|
||||||
|
def on_any_event(self, event: FileSystemEvent) -> None:
|
||||||
|
if not (
|
||||||
|
event.is_directory or event.event_type in ["deleted", "closed"]
|
||||||
|
):
|
||||||
|
process(pathlib.Path(event.src_path))
|
||||||
|
|
||||||
|
event_handler = EventHandler()
|
||||||
|
observer.schedule(event_handler, input, recursive=False)
|
||||||
|
observer.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
time.sleep(1)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
observer.stop()
|
||||||
|
observer.join()
|
||||||
|
|
||||||
|
|
||||||
|
@main.command(help="for a http server")
|
||||||
|
@click.option(
|
||||||
|
"-p",
|
||||||
|
"--port",
|
||||||
|
default=5000,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="port",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-l",
|
||||||
|
"--log_level",
|
||||||
|
default="info",
|
||||||
|
type=str,
|
||||||
|
show_default=True,
|
||||||
|
help="log level",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-t",
|
||||||
|
"--threads",
|
||||||
|
default=None,
|
||||||
|
type=int,
|
||||||
|
show_default=True,
|
||||||
|
help="number of worker threads",
|
||||||
|
)
|
||||||
|
def s(port: int, log_level: str, threads: int) -> None:
|
||||||
|
sessions: dict[str, BaseSession] = {}
|
||||||
|
tags_metadata = [
|
||||||
|
{
|
||||||
|
"name": "Background Removal",
|
||||||
|
"description": "Endpoints that perform background removal with different image sources.",
|
||||||
|
"externalDocs": {
|
||||||
|
"description": "GitHub Source",
|
||||||
|
"url": "https://github.com/danielgatis/rembg",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
]
|
||||||
|
app = FastAPI(
|
||||||
|
title="Rembg",
|
||||||
|
description="Rembg is a tool to remove images background. That is it.",
|
||||||
|
version=_version.get_versions()["version"],
|
||||||
|
contact={
|
||||||
|
"name": "Daniel Gatis",
|
||||||
|
"url": "https://github.com/danielgatis",
|
||||||
|
"email": "danielgatis@gmail.com",
|
||||||
|
},
|
||||||
|
license_info={
|
||||||
|
"name": "MIT License",
|
||||||
|
"url": "https://github.com/danielgatis/rembg/blob/main/LICENSE.txt",
|
||||||
|
},
|
||||||
|
openapi_tags=tags_metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
app.add_middleware(
|
||||||
|
CORSMiddleware,
|
||||||
|
allow_credentials=True,
|
||||||
|
allow_origins=["*"],
|
||||||
|
allow_methods=["*"],
|
||||||
|
allow_headers=["*"],
|
||||||
|
)
|
||||||
|
|
||||||
|
class ModelType(str, Enum):
|
||||||
|
u2net = "u2net"
|
||||||
|
u2netp = "u2netp"
|
||||||
|
u2net_human_seg = "u2net_human_seg"
|
||||||
|
u2net_cloth_seg = "u2net_cloth_seg"
|
||||||
|
silueta = "silueta"
|
||||||
|
|
||||||
|
class CommonQueryParams:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
model: ModelType = Query(
|
||||||
|
default=ModelType.u2net,
|
||||||
|
description="Model to use when processing image",
|
||||||
|
),
|
||||||
|
a: bool = Query(default=False, description="Enable Alpha Matting"),
|
||||||
|
af: int = Query(
|
||||||
|
default=240,
|
||||||
|
ge=0,
|
||||||
|
le=255,
|
||||||
|
description="Alpha Matting (Foreground Threshold)",
|
||||||
|
),
|
||||||
|
ab: int = Query(
|
||||||
|
default=10,
|
||||||
|
ge=0,
|
||||||
|
le=255,
|
||||||
|
description="Alpha Matting (Background Threshold)",
|
||||||
|
),
|
||||||
|
ae: int = Query(
|
||||||
|
default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
|
||||||
|
),
|
||||||
|
om: bool = Query(default=False, description="Only Mask"),
|
||||||
|
ppm: bool = Query(default=False, description="Post Process Mask"),
|
||||||
|
):
|
||||||
|
self.model = model
|
||||||
|
self.a = a
|
||||||
|
self.af = af
|
||||||
|
self.ab = ab
|
||||||
|
self.ae = ae
|
||||||
|
self.om = om
|
||||||
|
self.ppm = ppm
|
||||||
|
|
||||||
|
class CommonQueryPostParams:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
model: ModelType = Form(
|
||||||
|
default=ModelType.u2net,
|
||||||
|
description="Model to use when processing image",
|
||||||
|
),
|
||||||
|
a: bool = Form(default=False, description="Enable Alpha Matting"),
|
||||||
|
af: int = Form(
|
||||||
|
default=240,
|
||||||
|
ge=0,
|
||||||
|
le=255,
|
||||||
|
description="Alpha Matting (Foreground Threshold)",
|
||||||
|
),
|
||||||
|
ab: int = Form(
|
||||||
|
default=10,
|
||||||
|
ge=0,
|
||||||
|
le=255,
|
||||||
|
description="Alpha Matting (Background Threshold)",
|
||||||
|
),
|
||||||
|
ae: int = Form(
|
||||||
|
default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
|
||||||
|
),
|
||||||
|
om: bool = Form(default=False, description="Only Mask"),
|
||||||
|
ppm: bool = Form(default=False, description="Post Process Mask"),
|
||||||
|
):
|
||||||
|
self.model = model
|
||||||
|
self.a = a
|
||||||
|
self.af = af
|
||||||
|
self.ab = ab
|
||||||
|
self.ae = ae
|
||||||
|
self.om = om
|
||||||
|
self.ppm = ppm
|
||||||
|
|
||||||
|
def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
|
||||||
|
return Response(
|
||||||
|
remove(
|
||||||
|
content,
|
||||||
|
session=sessions.setdefault(
|
||||||
|
commons.model.value, new_session(commons.model.value)
|
||||||
|
),
|
||||||
|
alpha_matting=commons.a,
|
||||||
|
alpha_matting_foreground_threshold=commons.af,
|
||||||
|
alpha_matting_background_threshold=commons.ab,
|
||||||
|
alpha_matting_erode_size=commons.ae,
|
||||||
|
only_mask=commons.om,
|
||||||
|
post_process_mask=commons.ppm,
|
||||||
|
),
|
||||||
|
media_type="image/png",
|
||||||
|
)
|
||||||
|
|
||||||
|
@app.on_event("startup")
|
||||||
|
def startup():
|
||||||
|
if threads is not None:
|
||||||
|
from anyio import CapacityLimiter
|
||||||
|
from anyio.lowlevel import RunVar
|
||||||
|
|
||||||
|
RunVar("_default_thread_limiter").set(CapacityLimiter(threads))
|
||||||
|
|
||||||
|
@app.get(
|
||||||
|
path="/",
|
||||||
|
tags=["Background Removal"],
|
||||||
|
summary="Remove from URL",
|
||||||
|
description="Removes the background from an image obtained by retrieving an URL.",
|
||||||
|
)
|
||||||
|
async def get_index(
|
||||||
|
url: str = Query(
|
||||||
|
default=..., description="URL of the image that has to be processed."
|
||||||
|
),
|
||||||
|
commons: CommonQueryParams = Depends(),
|
||||||
|
):
|
||||||
|
async with aiohttp.ClientSession() as session:
|
||||||
|
async with session.get(url) as response:
|
||||||
|
file = await response.read()
|
||||||
|
return await asyncify(im_without_bg)(file, commons)
|
||||||
|
|
||||||
|
@app.post(
|
||||||
|
path="/",
|
||||||
|
tags=["Background Removal"],
|
||||||
|
summary="Remove from Stream",
|
||||||
|
description="Removes the background from an image sent within the request itself.",
|
||||||
|
)
|
||||||
|
async def post_index(
|
||||||
|
file: bytes = File(
|
||||||
|
default=...,
|
||||||
|
description="Image file (byte stream) that has to be processed.",
|
||||||
|
),
|
||||||
|
commons: CommonQueryPostParams = Depends(),
|
||||||
|
):
|
||||||
|
return await asyncify(im_without_bg)(file, commons)
|
||||||
|
|
||||||
|
uvicorn.run(app, host="0.0.0.0", port=port, log_level=log_level)
|
|
@ -0,0 +1,40 @@
|
||||||
|
from typing import Dict, List, Tuple
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import onnxruntime as ort
|
||||||
|
from PIL import Image
|
||||||
|
from PIL.Image import Image as PILImage
|
||||||
|
|
||||||
|
|
||||||
|
class BaseSession:
|
||||||
|
def __init__(self, model_name: str, inner_session: ort.InferenceSession):
|
||||||
|
self.model_name = model_name
|
||||||
|
self.inner_session = inner_session
|
||||||
|
|
||||||
|
def normalize(
|
||||||
|
self,
|
||||||
|
img: PILImage,
|
||||||
|
mean: Tuple[float, float, float],
|
||||||
|
std: Tuple[float, float, float],
|
||||||
|
size: Tuple[int, int],
|
||||||
|
) -> Dict[str, np.ndarray]:
|
||||||
|
im = img.convert("RGB").resize(size, Image.LANCZOS)
|
||||||
|
|
||||||
|
im_ary = np.array(im)
|
||||||
|
im_ary = im_ary / np.max(im_ary)
|
||||||
|
|
||||||
|
tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
|
||||||
|
tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0]
|
||||||
|
tmpImg[:, :, 1] = (im_ary[:, :, 1] - mean[1]) / std[1]
|
||||||
|
tmpImg[:, :, 2] = (im_ary[:, :, 2] - mean[2]) / std[2]
|
||||||
|
|
||||||
|
tmpImg = tmpImg.transpose((2, 0, 1))
|
||||||
|
|
||||||
|
return {
|
||||||
|
self.inner_session.get_inputs()[0]
|
||||||
|
.name: np.expand_dims(tmpImg, 0)
|
||||||
|
.astype(np.float32)
|
||||||
|
}
|
||||||
|
|
||||||
|
def predict(self, img: PILImage) -> List[PILImage]:
|
||||||
|
raise NotImplementedError
|
|
@ -0,0 +1,88 @@
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
from PIL.Image import Image as PILImage
|
||||||
|
from scipy.special import log_softmax
|
||||||
|
|
||||||
|
from .session_base import BaseSession
|
||||||
|
|
||||||
|
pallete1 = [
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
]
|
||||||
|
|
||||||
|
pallete2 = [
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
]
|
||||||
|
|
||||||
|
pallete3 = [
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
255,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class ClothSession(BaseSession):
|
||||||
|
def predict(self, img: PILImage) -> List[PILImage]:
|
||||||
|
ort_outs = self.inner_session.run(
|
||||||
|
None, self.normalize(img, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (768, 768))
|
||||||
|
)
|
||||||
|
|
||||||
|
pred = ort_outs
|
||||||
|
pred = log_softmax(pred[0], 1)
|
||||||
|
pred = np.argmax(pred, axis=1, keepdims=True)
|
||||||
|
pred = np.squeeze(pred, 0)
|
||||||
|
pred = np.squeeze(pred, 0)
|
||||||
|
|
||||||
|
mask = Image.fromarray(pred.astype("uint8"), mode="L")
|
||||||
|
mask = mask.resize(img.size, Image.LANCZOS)
|
||||||
|
|
||||||
|
masks = []
|
||||||
|
|
||||||
|
mask1 = mask.copy()
|
||||||
|
mask1.putpalette(pallete1)
|
||||||
|
mask1 = mask1.convert("RGB").convert("L")
|
||||||
|
masks.append(mask1)
|
||||||
|
|
||||||
|
mask2 = mask.copy()
|
||||||
|
mask2.putpalette(pallete2)
|
||||||
|
mask2 = mask2.convert("RGB").convert("L")
|
||||||
|
masks.append(mask2)
|
||||||
|
|
||||||
|
mask3 = mask.copy()
|
||||||
|
mask3.putpalette(pallete3)
|
||||||
|
mask3 = mask3.convert("RGB").convert("L")
|
||||||
|
masks.append(mask3)
|
||||||
|
|
||||||
|
return masks
|
|
@ -0,0 +1,71 @@
|
||||||
|
import hashlib
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from contextlib import redirect_stdout
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Type
|
||||||
|
|
||||||
|
import onnxruntime as ort
|
||||||
|
import pooch
|
||||||
|
|
||||||
|
from .session_base import BaseSession
|
||||||
|
from .session_cloth import ClothSession
|
||||||
|
from .session_simple import SimpleSession
|
||||||
|
|
||||||
|
|
||||||
|
def new_session(model_name: str = "u2net") -> BaseSession:
|
||||||
|
session_class: Type[BaseSession]
|
||||||
|
md5 = "60024c5c889badc19c04ad937298a77b"
|
||||||
|
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
|
||||||
|
session_class = SimpleSession
|
||||||
|
|
||||||
|
if model_name == "u2netp":
|
||||||
|
md5 = "8e83ca70e441ab06c318d82300c84806"
|
||||||
|
url = (
|
||||||
|
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
|
||||||
|
)
|
||||||
|
session_class = SimpleSession
|
||||||
|
elif model_name == "u2net_human_seg":
|
||||||
|
md5 = "c09ddc2e0104f800e3e1bb4652583d1f"
|
||||||
|
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx"
|
||||||
|
session_class = SimpleSession
|
||||||
|
elif model_name == "u2net_cloth_seg":
|
||||||
|
md5 = "2434d1f3cb744e0e49386c906e5a08bb"
|
||||||
|
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx"
|
||||||
|
session_class = ClothSession
|
||||||
|
elif model_name == "silueta":
|
||||||
|
md5 = "55e59e0d8062d2f5d013f4725ee84782"
|
||||||
|
url = (
|
||||||
|
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx"
|
||||||
|
)
|
||||||
|
session_class = SimpleSession
|
||||||
|
|
||||||
|
u2net_home = os.getenv(
|
||||||
|
"U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
|
||||||
|
)
|
||||||
|
|
||||||
|
fname = f"{model_name}.onnx"
|
||||||
|
path = Path(u2net_home).expanduser()
|
||||||
|
full_path = Path(u2net_home).expanduser() / fname
|
||||||
|
|
||||||
|
pooch.retrieve(
|
||||||
|
url,
|
||||||
|
f"md5:{md5}",
|
||||||
|
fname=fname,
|
||||||
|
path=Path(u2net_home).expanduser(),
|
||||||
|
progressbar=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
sess_opts = ort.SessionOptions()
|
||||||
|
|
||||||
|
if "OMP_NUM_THREADS" in os.environ:
|
||||||
|
sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
|
||||||
|
|
||||||
|
return session_class(
|
||||||
|
model_name,
|
||||||
|
ort.InferenceSession(
|
||||||
|
str(full_path),
|
||||||
|
providers=ort.get_available_providers(),
|
||||||
|
sess_options=sess_opts,
|
||||||
|
),
|
||||||
|
)
|
|
@ -0,0 +1,30 @@
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
from PIL.Image import Image as PILImage
|
||||||
|
|
||||||
|
from .session_base import BaseSession
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleSession(BaseSession):
|
||||||
|
def predict(self, img: PILImage) -> List[PILImage]:
|
||||||
|
ort_outs = self.inner_session.run(
|
||||||
|
None,
|
||||||
|
self.normalize(
|
||||||
|
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
pred = ort_outs[0][:, 0, :, :]
|
||||||
|
|
||||||
|
ma = np.max(pred)
|
||||||
|
mi = np.min(pred)
|
||||||
|
|
||||||
|
pred = (pred - mi) / (ma - mi)
|
||||||
|
pred = np.squeeze(pred)
|
||||||
|
|
||||||
|
mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
|
||||||
|
mask = mask.resize(img.size, Image.LANCZOS)
|
||||||
|
|
||||||
|
return [mask]
|
|
@ -0,0 +1,18 @@
|
||||||
|
aiohttp==3.8.1
|
||||||
|
asyncer==0.0.2
|
||||||
|
click==8.1.3
|
||||||
|
fastapi==0.87.0
|
||||||
|
filetype==1.2.0
|
||||||
|
pooch==1.6.0
|
||||||
|
imagehash==4.3.1
|
||||||
|
numpy==1.23.5
|
||||||
|
onnxruntime==1.13.1
|
||||||
|
opencv-python-headless==4.6.0.66
|
||||||
|
pillow==9.3.0
|
||||||
|
pymatting==1.1.8
|
||||||
|
python-multipart==0.0.5
|
||||||
|
scikit-image==0.19.3
|
||||||
|
scipy==1.9.3
|
||||||
|
tqdm==4.64.1
|
||||||
|
uvicorn==0.20.0
|
||||||
|
watchdog==2.1.9
|
Binary file not shown.
Loading…
Reference in New Issue