Daniel Gatis преди 2 години
родител
ревизия
c72eb1206c
променени са 3 файла, в които са добавени 100 реда и са изтрити 84 реда
  1. 95 75
      README.md
  2. 1 0
      rembg/__init__.py
  3. 4 9
      rembg/session_factory.py

+ 95 - 75
README.md

@@ -6,7 +6,7 @@
 [![License](https://img.shields.io/badge/License-MIT-blue.svg)](https://img.shields.io/badge/License-MIT-blue.svg)
 [![Hugging Face Spaces](https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/KenjieDec/RemBG)
 
-Rembg is a tool to remove images background. That is it.
+Rembg is a tool to remove images background.
 
 <p style="display: flex;align-items: center;justify-content: center;">
   <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-1.jpg" width="100" />
@@ -37,13 +37,13 @@ Rembg is a tool to remove images background. That is it.
 
 **If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
 
-### Requirements
+## Requirements
 
 ```
 python: >3.7, <3.11
 ```
 
-### Installation
+## Installation
 
 CPU support:
 
@@ -57,66 +57,97 @@ GPU support:
 pip install rembg[gpu]
 ```
 
-### Usage as a cli
+## Usage as a cli
+
+After the installation step you can use rembg just typing `rembg` in your terminal window.
+
+The `rembg` command has 3 subcommands, one for each input type:
+- `i` for files
+- `p` for folders
+- `s` for http server
+
+You can get help about the main command using:
+
+```
+rembg --help
+```
+
+As well, about all the subcommands using:
+
+```
+rembg <COMMAND> --help
+```
+
+### rembg `i`
+
+Used when input and output are files.
 
 Remove the background from a remote image
 
-```bash
+```
 curl -s http://input.png | rembg i > output.png
 ```
 
 Remove the background from a local file
 
-```bash
+```
 rembg i path/to/input.png path/to/output.png
 ```
 
-Remove the background from all images in a folder
+Remove the background specifying a model
 
-```bash
-rembg p path/to/input path/to/output
 ```
+rembg -m u2netp i path/to/input.png path/to/output.png
+```
+
+Remove the background returning only the mask
 
-### Usage as a server
+```
+rembg -om i path/to/input.png path/to/output.png
+```
 
-Start the server
 
-```bash
-rembg s
+Remove the background applying an alpha matting
+
+```
+rembg -a i path/to/input.png path/to/output.png
 ```
 
-And go to:
+### rembg `p`
+
+Used when input and output are folders.
+
+Remove the background from all images in a folder
 
 ```
-http://localhost:5000/docs
+rembg p path/to/input path/to/output
 ```
 
-Image with background:
+Same as before, but watching for new/changed files to process
 
 ```
-https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
+rembg p -w path/to/input path/to/output
 ```
 
-Image without background:
+### rembg `s`
+
+Used to start http server.
+
+To see the complete endpoints documentation, go to: `http://localhost:5000/docs`.
+
+Remove the background from an image url
 
 ```
-http://localhost:5000/?url=https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
+curl -s "http://localhost:5000/?url=http://input.png" -o output.png
 ```
 
-Also you can send the file as a FormData (multipart/form-data):
+Remove the background from an uploaded image
 
-```html
-<form
-    action="http://localhost:5000"
-    method="post"
-    enctype="multipart/form-data"
->
-    <input type="file" name="file" />
-    <input type="submit" value="upload" />
-</form>
+```
+curl -s -F file=@/path/to/input.jpg "http://localhost:5000"  -o output.png
 ```
 
-### Usage as a library
+## Usage as a library
 
 Input and output as bytes
 
@@ -161,27 +192,36 @@ output = remove(input)
 cv2.imwrite(output_path, output)
 ```
 
-### Usage as a docker
+How to iterate over files in a performatic way
 
-Try this:
+```python
+from pathlib import Path
+from rembg import remove, new_session
 
-```
-docker run -p 5000:5000 danielgatis/rembg s
-```
+session = new_session()
 
-Image with background:
+for file in Path('path/to/folder').glob('*.png'):
+    input_path = str(file)
+    output_path = str(file.parent / (file.stem + ".out.png"))
 
-```
-https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
+    with open(input_path, 'rb') as i:
+        with open(output_path, 'wb') as o:
+            input = i.read()
+            output = remove(input, session=session)
+            o.write(output)
 ```
 
-Image without background:
+## Usage as a docker
+
+Just replace the `rembg` command for `docker run danielgatis/rembg`.
+
+Try this:
 
 ```
-http://localhost:5000/?url=https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
+docker run danielgatis/rembg i path/to/input.png path/to/output.png
 ```
 
-### Models
+## Models
 
 All models are downloaded and saved in the user home folder in the `.u2net` directory.
 
@@ -191,55 +231,35 @@ The available models are:
 -   u2netp ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A lightweight version of u2net model.
 -   u2net_human_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation.
 -   u2net_cloth_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx), [source](https://github.com/levindabhi/cloth-segmentation)): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
--   silueta ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx), [source](https://github.com/xuebinqin/U-2-Net/issues/295)): Same as u2net but the size is reduced to 43Mb.
+-   silueta ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx), [source](https://github.com/xuebinqin/U-2-Net/issues/295)): Same as u2net but the size is reduced to 43Mb.
 
-#### How to train your own model
+### How to train your own model
 
 If You need more fine tunned models try this:
 https://github.com/danielgatis/rembg/issues/193#issuecomment-1055534289
 
-### Advance usage
-
-Sometimes it is possible to achieve better results by turning on alpha matting. Example:
-
-```bash
-curl -s http://input.png | rembg i -a -ae 15 > output.png
-```
-
-<table>
-    <thead>
-        <tr>
-            <td>Original</td>
-            <td>Without alpha matting</td>
-            <td>With alpha matting (-a -ae 15)</td>
-        </tr>
-    </thead>
-    <tbody>
-        <tr>
-            <td><img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/food-1.jpg"/></td>
-            <td><img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/food-1.out.jpg"/></td>
-            <td><img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/food-1.out.alpha.jpg"/></td>
-        </tr>
-    </tbody>
-</table>
 
-### In the cloud
+## Some video tutorials
 
-Please contact me at [email protected] if you need help to put it on the cloud.
+- https://www.youtube.com/watch?v=3xqwpXjxyMQ
+- https://www.youtube.com/watch?v=dFKRGXdkGJU
+- https://www.youtube.com/watch?v=Ai-BS_T7yjE
+- https://www.youtube.com/watch?v=dFKRGXdkGJU
+- https://www.youtube.com/watch?v=D7W-C0urVcQ
 
-### References
+## References
 
--   https://arxiv.org/pdf/2005.09007.pdf
--   https://github.com/NathanUA/U-2-Net
--   https://github.com/pymatting/pymatting
+- https://arxiv.org/pdf/2005.09007.pdf
+- https://github.com/NathanUA/U-2-Net
+- https://github.com/pymatting/pymatting
 
-### Buy me a coffee
+## Buy me a coffee
 
 Liked some of my work? Buy me a coffee (or more likely a beer)
 
 <a href="https://www.buymeacoffee.com/danielgatis" target="_blank"><img src="https://bmc-cdn.nyc3.digitaloceanspaces.com/BMC-button-images/custom_images/orange_img.png" alt="Buy Me A Coffee" style="height: auto !important;width: auto !important;"></a>
 
-### License
+## License
 
 Copyright (c) 2020-present [Daniel Gatis](https://github.com/danielgatis)
 

+ 1 - 0
rembg/__init__.py

@@ -3,3 +3,4 @@ from . import _version
 __version__ = _version.get_versions()["version"]
 
 from .bg import remove
+from .session_factory import new_session

+ 4 - 9
rembg/session_factory.py

@@ -13,8 +13,11 @@ from .session_cloth import ClothSession
 from .session_simple import SimpleSession
 
 
-def new_session(model_name: str) -> BaseSession:
+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"
@@ -22,10 +25,6 @@ def new_session(model_name: str) -> BaseSession:
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
         )
         session_class = SimpleSession
-    elif model_name == "u2net":
-        md5 = "60024c5c889badc19c04ad937298a77b"
-        url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.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"
@@ -40,10 +39,6 @@ def new_session(model_name: str) -> BaseSession:
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx"
         )
         session_class = SimpleSession
-    else:
-        assert AssertionError(
-            "Choose between u2net, u2netp, u2net_human_seg or u2net_cloth_seg"
-        )
 
     u2net_home = os.getenv(
         "U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")