From 9a91de5fe7dc81833339c948d89811dafd2fe24d Mon Sep 17 00:00:00 2001
From: Jason Dogariu <jdogariu@gmail.com>
Date: Sun, 2 Apr 2023 19:54:25 -0400
Subject: [PATCH] Simple background removal (rembg) node

---
 __init__.py      | 39 +++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 2 files changed, 40 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 requirements.txt

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..9f203d8
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,39 @@
+from rembg import remove
+from PIL import Image
+import torch
+import numpy as np
+
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# Convert PIL to Tensor
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+class ImageRemoveBackgroundRembg:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "remove_background"
+    CATEGORY = "image"
+
+    def remove_background(self, image):
+        image = pil2tensor(remove(tensor2pil(image)))
+        return (image,)
+
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "Image Remove Background (rembg)": ImageRemoveBackgroundRembg
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..ed2ce86
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+rembg