add more models

danielgatis · Apr 22, 2022 · 84a21b4 · 84a21b4
1 parent 738cf15
commit 84a21b4
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 15 deletions.
diff --git a/rembg/bg.py b/rembg/bg.py
@@ -120,7 +120,7 @@ def remove(
                     alpha_matting_background_threshold,
                     alpha_matting_erode_size,
                 )
-            except:
+            except ValueError:
                 cutout = naive_cutout(img, mask)
 
         else:

diff --git a/rembg/session_cloth.py b/rembg/session_cloth.py
@@ -6,28 +6,50 @@
 
 from .session_base import BaseSession
 
-# fmt: off
 pallete1 = [
-      0,   0,   0, # background
-    255, 255, 255, # upper body
-      0,   0,   0, # lower body
-      0,   0,   0, # full body
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
 ]
 
 pallete2 = [
-      0,   0,   0, # background
-      0,   0,   0, # upper body
-    255, 255, 255, # lower body
-      0,   0,   0, # full body
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
+    0,
+    0,
+    0,
 ]
 
 pallete3 = [
-      0,   0,   0, # background
-      0,   0,   0, # upper body
-      0,   0,   0, # lower body
-    255, 255, 255, # full body
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
 ]
-# fmt: on
 
 
 class ClothSession(BaseSession):