add anime model

README.md

@@ -37,6 +37,15 @@ Rembg is a tool to remove images background.
   <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-3.out.png" width="100" />
 </p>
 
+<p style="display: flex;align-items: center;justify-content: center;">
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-1.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-1.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-2.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-2.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-3.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/anime-girl-3.out.png" width="100" />
+</p>
+
 **If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
 
 ## Sponsor

rembg/sessions/dis_anime.py

@@ -0,0 +1,49 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class DisSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)),
+        )
+
+        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]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-anime.onnx",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:6f184e756bb3bd901c8849220a83e38e",
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "isnet-anime"

tests/test_remove.py

@@ -11,6 +11,11 @@ here = Path(__file__).parent.resolve()
 def test_remove():
     kwargs = {
         "sam": {
+            "anime-girl-1" : {
+                "input_points": [[400, 165]],
+                "input_labels": [1],
+            },
+
             "car-1" : {
                 "input_points": [[250, 200]],
                 "input_labels": [1],
@@ -19,7 +24,7 @@ def test_remove():
             "cloth-1" : {
                 "input_points": [[370, 495]],
                 "input_labels": [1],
-            }
+            },
         }
     }
 
@@ -30,9 +35,10 @@ def test_remove():
         "u2net_cloth_seg",
         "silueta",
         "isnet-general-use",
+        "isnet-anime",
         "sam"
     ]:
-        for picture in ["car-1", "cloth-1"]:
+        for picture in ["anime-girl-1", "car-1", "cloth-1"]:
             image_path = Path(here / "fixtures" / f"{picture}.jpg")
             image = image_path.read_bytes()