AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-05 03:50:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix dockerfile

Browse Source
This commit is contained in:
Daniel Gatis 2023-06-02 08:54:59 -03:00
parent 0766402109
commit fd0c871aac
2 changed files with 10 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
Dockerfile
View File

@ -1,30 +1,11 @@
FROM nvidia/cuda:11.6.0-runtime-ubuntu18.04 FROM python:3.10-slim
ENV DEBIAN_FRONTEND noninteractive
RUN rm /etc/apt/sources.list.d/cuda.list || true
RUN rm /etc/apt/sources.list.d/nvidia-ml.list || true
RUN apt-key del 7fa2af80
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
RUN apt update -y
RUN apt upgrade -y
RUN apt install -y curl wget software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt install -y python3.9 python3.9-distutils
RUN curl https://bootstrap.pypa.io/get-pip.py | python3.9
WORKDIR /rembg WORKDIR /rembg
COPY . . COPY . .
RUN python3.9 -m pip install .[gpu] RUN pip install --upgrade pip
RUN python -m pip install .
RUN mkdir -p ~/.u2net RUN python -c 'from rembg.bg import download_models; download_models()'
RUN wget https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx -O ~/.u2net/u2netp.onnx
RUN wget https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx -O ~/.u2net/u2net.onnx
RUN wget https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx -O ~/.u2net/u2net_human_seg.onnx
RUN wget https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx -O ~/.u2net/u2net_cloth_seg.onnx
EXPOSE 5000 EXPOSE 5000
ENTRYPOINT ["rembg"] ENTRYPOINT ["rembg"]

6
rembg/bg.py
View File

@ -19,6 +19,7 @@ from pymatting.util.util import stack_images
from scipy.ndimage import binary_erosion from scipy.ndimage import binary_erosion
from .session_factory import new_session from .session_factory import new_session
from .sessions import sessions_class
from .sessions.base import BaseSession from .sessions.base import BaseSession
kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3)) kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
@ -117,6 +118,11 @@ def fix_image_orientation(img: PILImage) -> PILImage:
return ImageOps.exif_transpose(img) return ImageOps.exif_transpose(img)
def download_models() -> None:
for session in sessions_class:
session.download_models()
def remove( def remove(
data: Union[bytes, PILImage, np.ndarray], data: Union[bytes, PILImage, np.ndarray],
alpha_matting: bool = False, alpha_matting: bool = False,