Merge pull request #114 from OPHoperHPO/development

✂️ CarveKit v4.1.0
OPHoperHPO · Nov 5, 2022 · d274ff5 · d274ff5
2 parents f241215 + 992af42
commit d274ff5
Show file tree

Hide file tree

Showing 80 changed files with 4,017 additions and 1,277 deletions.
diff --git a/.flake8 b/.flake8
@@ -0,0 +1,5 @@
+[flake8]
+ignore = E203, E266, E501, W503, F403, F401
+max-line-length = 79
+max-complexity = 50
+select = B,C,E,F,W,T4,B9
diff --git a/.gitignore b/.gitignore
@@ -4,3 +4,4 @@ __pycache__
 /tests/tests_temp/
 /.pytest_cache/
 /carvekit.egg-info/
+venv
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,10 @@
+repos:
+-   repo: https://github.com/psf/black
+    rev: 22.10.0
+    hooks:
+    - id: black
+      language_version: python3.10
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: 3.9.2
+    hooks:
+    - id: flake8
diff --git a/Dockerfile.cpu b/Dockerfile.cpu
@@ -5,7 +5,7 @@ WORKDIR /app
 RUN pip3 install --no-cache-dir tqdm==4.64.0 requests==2.27.1
 RUN mkdir -p ./carvekit/utils/
 RUN mkdir -p ./carvekit/ml/files
-RUN touch ./carvekit/__init__.py
+COPY ./carvekit/__init__.py ./carvekit/__init__.py
 RUN touch ./carvekit/ml/__init__.py
 RUN touch ./carvekit/utils/__init__.py
 COPY ./carvekit/utils/download_models.py ./carvekit/utils/download_models.py
@@ -18,7 +18,7 @@ FROM python:3.10.4
 WORKDIR /app
 
 RUN apt-get update && apt-get -y install libgl1  # Install cv2 dep.
-COPY --from=builder /root/.carvekit /root/.carvekit
+COPY --from=builder /root/.cache/carvekit /root/.cache/carvekit
 
 # Install requirements
 COPY requirements.txt ./
@@ -34,15 +34,19 @@ RUN pip3 install -e ./
 
 ENV CARVEKIT_PORT '5000'
 ENV CARVEKIT_HOST '0.0.0.0'
-ENV CARVEKIT_SEGMENTATION_NETWORK 'u2net'
+ENV CARVEKIT_SEGMENTATION_NETWORK 'tracer_b7'
 ENV CARVEKIT_PREPROCESSING_METHOD 'none'
 ENV CARVEKIT_POSTPROCESSING_METHOD 'fba'
 ENV CARVEKIT_DEVICE 'cpu'
 ENV CARVEKIT_BATCH_SIZE_SEG '5'
 ENV CARVEKIT_BATCH_SIZE_MATTING '1'
-ENV CARVEKIT_SEG_MASK_SIZE '320'
+ENV CARVEKIT_SEG_MASK_SIZE '640'
 ENV CARVEKIT_MATTING_MASK_SIZE '2048'
 ENV CARVEKIT_AUTH_ENABLE '1'
+ENV CARVEKIT_FP16 '0'
+ENV CARVEKIT_TRIMAP_PROB_THRESHOLD=231
+ENV CARVEKIT_TRIMAP_DILATION=30
+ENV CARVEKIT_TRIMAP_EROSION=5
 
 # Tokens will be generated automatically every time the container is restarted if ENV is not set.
 

diff --git a/Dockerfile.cuda b/Dockerfile.cuda
@@ -5,7 +5,7 @@ WORKDIR /app
 RUN pip3 install --no-cache-dir tqdm==4.64.0 requests==2.27.1
 RUN mkdir -p ./carvekit/utils/
 RUN mkdir -p ./carvekit/ml/files
-RUN touch ./carvekit/__init__.py
+COPY ./carvekit/__init__.py ./carvekit/__init__.py
 RUN touch ./carvekit/ml/__init__.py
 RUN touch ./carvekit/utils/__init__.py
 COPY ./carvekit/utils/download_models.py ./carvekit/utils/download_models.py
@@ -18,7 +18,7 @@ FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
 WORKDIR /app
 
 RUN apt-get update && apt-get -y install libgl1 libglib2.0-0  # Install cv2 dep.
-COPY --from=builder /root/.carvekit /root/.carvekit
+COPY --from=builder /root/.cache/carvekit /root/.cache/carvekit
 
 # Install requirements
 COPY requirements.txt ./
@@ -34,15 +34,19 @@ RUN pip3 install -e ./
 
 ENV CARVEKIT_PORT '5000'
 ENV CARVEKIT_HOST '0.0.0.0'
-ENV CARVEKIT_SEGMENTATION_NETWORK 'u2net'
+ENV CARVEKIT_SEGMENTATION_NETWORK 'tracer_b7'
 ENV CARVEKIT_PREPROCESSING_METHOD 'none'
 ENV CARVEKIT_POSTPROCESSING_METHOD 'fba'
 ENV CARVEKIT_DEVICE 'cuda'
 ENV CARVEKIT_BATCH_SIZE_SEG '5'
 ENV CARVEKIT_BATCH_SIZE_MATTING '1'
-ENV CARVEKIT_SEG_MASK_SIZE '320'
+ENV CARVEKIT_SEG_MASK_SIZE '640'
 ENV CARVEKIT_MATTING_MASK_SIZE '2048'
 ENV CARVEKIT_AUTH_ENABLE '1'
+ENV CARVEKIT_FP16 '0'
+ENV CARVEKIT_TRIMAP_PROB_THRESHOLD=231
+ENV CARVEKIT_TRIMAP_DILATION=30
+ENV CARVEKIT_TRIMAP_EROSION=5
 
 # Tokens will be generated automatically every time the container is restarted if ENV is not set.
 

diff --git a/README.md b/README.md
@@ -5,7 +5,7 @@
 <p align="center">
 <img src="https://github.com/OPHoperHPO/image-background-remove-tool/actions/workflows/master_docker.yaml/badge.svg">
 <img src="https://github.com/OPHoperHPO/image-background-remove-tool/actions/workflows/master.yml/badge.svg">
-<a src="https://colab.research.google.com/github/OPHoperHPO/image-background-remove-tool/blob/master/docs/other/carvekit_try.ipynb">
+<a href="https://colab.research.google.com/github/OPHoperHPO/image-background-remove-tool/blob/master/docs/other/carvekit_try.ipynb">
 <img src="https://camo.githubusercontent.com/52feade06f2fecbf006889a904d221e6a730c194/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667"></a>
 
 </p>
@@ -24,11 +24,11 @@
 ## 📄 Description:  
 Automated high-quality background removal framework for an image using neural networks.
 
-
 ## 🎆 Features:  
 - High Quality
 - Batch Processing
 - NVIDIA CUDA and CPU processing
+- FP16 inference: Fast inference with low memory usage
 - Easy inference
 - 100% remove.bg compatible FastAPI HTTP API 
 - Removes background from hairs
@@ -42,11 +42,25 @@ It can be briefly described as
 3. Using machine learning technology, the background of the image is removed
 4. Image post-processing to improve the quality of the processed image
 ## 🎓 Implemented Neural Networks:
-* [U^2-net](https://github.com/NathanUA/U-2-Net)
-* [BASNet](https://github.com/NathanUA/BASNet)
-* [DeepLabV3](https://github.com/tensorflow/models/tree/master/research/deeplab)
-
-
+|        Networks         |                   Target                    |             Accuracy             |
+|:-----------------------:|:-------------------------------------------:|:--------------------------------:|
+| **Tracer-B7** (default) |     **General** (objects, animals, etc)     | **90%** (mean F1-Score, DUTS-TE) |
+|         U^2-net         | **Hairs** (hairs, people, animals, objects) |  80.4% (mean F1-Score, DUTS-TE)  |
+|         BASNet          |        **General** (people, objects)        |  80.3% (mean F1-Score, DUTS-TE)  |
+|        DeepLabV3        |         People, Animals, Cars, etc          |  67.4% (mean IoU, COCO val2017)  |
+
+### Recommended parameters for different models
+|  Networks   | Segmentation mask  size | Trimap parameters (dilation, erosion) |
+|:-----------:|:-----------------------:|:-------------------------------------:|
+| `tracer_b7` |           640           |                (30, 5)                |
+|   `u2net`   |           320           |                (30, 5)                |
+|  `basnet`   |           320           |                (30, 5)                |
+| `deeplabv3` |          1024           |               (40, 20)                |
+
+> ### Notes: 
+> 1. The final quality may depend on the resolution of your image, the type of scene or object.
+> 2. Use **U2-Net for hairs** and **Tracer-B7 for general images** and correct parameters. \
+> It is very important for final quality! Example images was taken by using U2-Net and FBA post-processing.
 ## 🖼️ Image pre-processing and post-processing methods:
 ### 🔍 Preprocessing methods:
 * `none` - No preprocessing methods used.
@@ -69,12 +83,21 @@ It can be briefly described as
 import torch
 from carvekit.api.high import HiInterface
 
-interface = HiInterface(batch_size_seg=5, batch_size_matting=1,
-                               device='cuda' if torch.cuda.is_available() else 'cpu',
-                               seg_mask_size=320, matting_mask_size=2048)
-images_without_background = interface(['./tests/data/cat.jpg'])                               
+# Check doc strings for more information
+interface = HiInterface(object_type="hairs-like",  # Can be "object" or "hairs-like".
+                        batch_size_seg=5,
+                        batch_size_matting=1,
+                        device='cuda' if torch.cuda.is_available() else 'cpu',
+                        seg_mask_size=640,  # Use 640 for Tracer B7 and 320 for U2Net
+                        matting_mask_size=2048,
+                        trimap_prob_threshold=231,
+                        trimap_dilation=30,
+                        trimap_erosion_iters=5,
+                        fp16=False)
+images_without_background = interface(['./tests/data/cat.jpg'])
 cat_wo_bg = images_without_background[0]
 cat_wo_bg.save('2.png')
+
 
 ```
 
@@ -84,12 +107,13 @@ import PIL.Image
 
 from carvekit.api.interface import Interface
 from carvekit.ml.wrap.fba_matting import FBAMatting
-from carvekit.ml.wrap.u2net import U2NET
+from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
 from carvekit.pipelines.postprocessing import MattingMethod
 from carvekit.pipelines.preprocessing import PreprocessingStub
 from carvekit.trimap.generator import TrimapGenerator
 
-u2net = U2NET(device='cpu',
+# Check doc strings for more information
+seg_net = TracerUniversalB7(device='cpu',
               batch_size=1)
 
 fba = FBAMatting(device='cpu',
@@ -106,7 +130,7 @@ postprocessing = MattingMethod(matting_module=fba,
 
 interface = Interface(pre_pipe=preprocessing,
                       post_pipe=postprocessing,
-                      seg_pipe=u2net)
+                      seg_pipe=seg_net)
 
 image = PIL.Image.open('tests/data/cat.jpg')
 cat_wo_bg = interface([image])[0]
@@ -129,7 +153,7 @@ Options:
   -o ./2.png                   Path to output file or dir
   --pre none                   Preprocessing method
   --post fba                   Postprocessing method.
-  --net u2net                  Segmentation Network
+  --net tracer_b7              Segmentation Network. Check README for more info.
   --recursive                  Enables recursive search for images in a folder
   --batch_size 10              Batch Size for list of images to be loaded to
                                RAM
@@ -140,20 +164,34 @@ Options:
   --batch_size_mat 1           Batch size for list of images to be processed
                                by matting network
 
-  --seg_mask_size 320          The size of the input image for the
-                               segmentation neural network.
+  --seg_mask_size 640          The size of the input image for the
+                               segmentation neural network. Use 640 for Tracer B7 and 320 for U2Net
 
   --matting_mask_size 2048     The size of the input image for the matting
                                neural network.
+  --trimap_dilation 30       The size of the offset radius from the
+                                  object mask in pixels when forming an
+                                  unknown area
+  --trimap_erosion 5        The number of iterations of erosion that the
+                                  object's mask will be subjected to before
+                                  forming an unknown area
+  --trimap_prob_threshold 231
+                                  Probability threshold at which the
+                                  prob_filter and prob_as_unknown_area
+                                  operations will be applied
 
   --device cpu                 Processing Device.
+  --fp16                       Enables mixed precision processing. Use only with CUDA. CPU support is experimental!
   --help                       Show this message and exit.
 
+
 ````
 ## 📦 Running the Framework / FastAPI HTTP API server via Docker:
+Using the API via docker is a **fast** and non-complex way to have a working API.
+> **Our docker images are available on [Docker Hub](https://hub.docker.com/r/anodev/carvekit).** \
+> Version tags are the same as the releases of the project with suffixes `-cpu` and `-cuda` for CPU and CUDA versions respectively.
+
 
-Using the API via docker is a **fast** and non-complex way to have a working API.\
-**This HTTP API is 100% compatible with remove.bg API clients.** 
 <p align="center"> 
 <img src="docs/imgs/screenshot/frontend.png"> 
 <img src="docs/imgs/screenshot/docs_fastapi.png"> 
@@ -188,21 +226,20 @@ See `docker-compose.<device>.yml` for more information. \
 1. Run `docker-compose -f docker-compose.cpu.yml run carvekit_api pytest`  # For testing on CPU
 2. Run `docker-compose -f docker-compose.cuda.yml run carvekit_api pytest`  # For testing on GPU
 
-
 ## 👪 Credits: [More info](docs/CREDITS.md)
 
 ## 💵 Support
   You can thank me for developing this project and buy me a small cup of coffee ☕
 
-| Blockchain |            Cryptocurrency           |          Network          |                                              Wallet                                             |
+| Blockchain |           Cryptocurrency            |          Network          |                                             Wallet                                              |
 |:----------:|:-----------------------------------:|:-------------------------:|:-----------------------------------------------------------------------------------------------:|
-|  Ethereum  | ETH / USDT / USDC / BNB / Dogecoin  |          Mainnet          |                            0x7Ab1B8015020242D2a9bC48F09b2F34b994bc2F8                           |
-|  Ethereum  |  ETH / USDT / USDC / BNB / Dogecoin | BSC (Binance Smart Chain) |                            0x7Ab1B8015020242D2a9bC48F09b2F34b994bc2F8                           |
-|   Bitcoin  |                 BTC                 |             -             |                            bc1qmf4qedujhhvcsg8kxpg5zzc2s3jvqssmu7mmhq                           |
-|    ZCash   |                 ZEC                 |             -             |                               t1d7b9WxdboGFrcVVHG2ZuwWBgWEKhNUbtm                               |
-|    Tron    |                 TRX                 |             -             |                                TH12CADSqSTcNZPvG77GVmYKAe4nrrJB5X                               |
+|  Ethereum  | ETH / USDT / USDC / BNB / Dogecoin  |          Mainnet          |                           0x7Ab1B8015020242D2a9bC48F09b2F34b994bc2F8                            |
+|  Ethereum  | ETH / USDT / USDC / BNB / Dogecoin  | BSC (Binance Smart Chain) |                           0x7Ab1B8015020242D2a9bC48F09b2F34b994bc2F8                            |
+|  Bitcoin   |                 BTC                 |             -             |                           bc1qmf4qedujhhvcsg8kxpg5zzc2s3jvqssmu7mmhq                            |
+|   ZCash    |                 ZEC                 |             -             |                               t1d7b9WxdboGFrcVVHG2ZuwWBgWEKhNUbtm                               |
+|    Tron    |                 TRX                 |             -             |                               TH12CADSqSTcNZPvG77GVmYKAe4nrrJB5X                                |
 |   Monero   |                 XMR                 |          Mainnet          | 48w2pDYgPtPenwqgnNneEUC9Qt1EE6eD5MucLvU3FGpY3SABudDa4ce5bT1t32oBwchysRCUimCkZVsD1HQRBbxVLF9GTh3 |
-|     TON    |                 TON                 |             -             |                         EQCznqTdfOKI3L06QX-3Q802tBL0ecSWIKfkSjU-qsoy0CWE                        |
+|    TON     |                 TON                 |             -             |                        EQCznqTdfOKI3L06QX-3Q802tBL0ecSWIKfkSjU-qsoy0CWE                         |
 ## 📧 __Feedback__
 I will be glad to receive feedback on the project and suggestions for integration.
 

diff --git a/carvekit/__init__.py b/carvekit/__init__.py
@@ -1 +1 @@
-version = "4.0.8"
+version = "4.1.0"