Improved Remove Background types (#2291)

* Improved Remove Background types * Minor fix
chaiNNer-org · Oct 30, 2023 · 8e8be5d · 8e8be5d
1 parent c41f915
commit 8e8be5d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/backend/src/packages/chaiNNer_onnx/onnx/processing/remove_background.py b/backend/src/packages/chaiNNer_onnx/onnx/processing/remove_background.py
@@ -25,7 +25,7 @@
         Currently supports u2net models from the rembg project (links found in readme).""",
     icon="ONNX",
     inputs=[
-        ImageInput(),
+        ImageInput(channels=[3, 4]),
         OnnxRemBgModelInput(),
         BoolInput("Post-process Mask", default=False),
         BoolInput("Alpha Matting", default=False),
@@ -38,7 +38,16 @@
     outputs=[
         ImageOutput(
             "Image",
-            image_type="""removeBackground(Input1, Input0)""",
+            image_type="""
+                let image = Input0;
+                let model = Input1;
+
+                Image {
+                    width: image.width,
+                    height: image.height * model.scaleHeight,
+                }
+            """,
+            channels=4,
         ),
         ImageOutput("Mask", image_type=navi.Image(size_as="Input0"), channels=1),
     ],

diff --git a/src/common/types/chainner-scope.ts b/src/common/types/chainner-scope.ts
@@ -122,14 +122,6 @@ def convenientUpscale(model: PyTorchModel | NcnnNetwork, image: Image) {
     }
 }
 
-def removeBackground(model: OnnxRemBgModel, image: Image) {
-    Image {
-        width: image.width,
-        height: image.height * model.scaleHeight,
-        channels: 4,
-    }
-}
-
 struct SplitFilePath {
     dir: Directory,
     basename: string,