Adding initial files

Dockerfile

@@ -0,0 +1,8 @@
+FROM rsamf/graphbook:10.0.0
+
+RUN pip install graphbook_huggingface
+
+COPY config.yaml .
+COPY workflow workflow
+
+CMD ["python", "-m", "graphbook.main", "--web_dir", "web/dist", "--isolate_users", "--no_sample"] 

README.md

@@ -1,5 +1,5 @@
 ---
-title: Rmbg Graphbook
+title: RMBG Graphbook
 emoji: 🏢
 colorFrom: red
 colorTo: pink
@@ -7,6 +7,10 @@ sdk: docker
 pinned: false
 license: mit
 short_description: An ML app to remove the background from image datasets
+app_port: 8005
+thumbnail: https://cdn.prod.website-files.com/6620137e5938f28a7e4eef8a/662013b554ec4c6bf1306445_Graphbook_Logo_Final_Borderless-512.png
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This is a Graphbook app that allows users to remove the background from any image dataset from Hugging Face. You may swap different datasets using the Hugging Face extension.
+
+Check out our documentation at https://docs.graphbook.ai/.

config.yaml

@@ -0,0 +1,2 @@
+plugins:
+  - "graphbook_huggingface"

workflow/BackgroundRemoval.json

@@ -0,0 +1 @@
+{"version": "0", "type": "workflow", "nodes": [{"id": "0", "type": "step", "position": {"x": 253.6139923086571, "y": 59.52037501520965}, "data": {"name": "RemoveBackground", "parameters": {"model": {"type": "resource"}, "batch_size": {"type": "number", "default": 8, "value": 8}, "item_key": {"type": "string", "default": "image", "value": "image"}, "output_dir": {"type": "string", "value": "/media/sam/shared/alternative/pokemon_masks"}}, "inputs": ["in"], "outputs": ["out"], "category": "Custom", "label": "RemoveBackground", "key": 12, "isCollapsed": false}, "width": 174, "height": 175, "selected": false, "positionAbsolute": {"x": 671.1501153288792, "y": 288.1274724574072}, "dragging": true, "parentId": "2", "hidden": false}, {"id": "2", "type": "group", "position": {"x": 417.53612302022213, "y": 228.60709744219753}, "width": 487, "height": 306, "data": {"label": "Group", "exports": {"inputs": [{"id": "1", "name": "in", "type": "step"}, {"id": "0", "name": "resource", "type": "resource"}], "outputs": [{"id": "0", "name": "out", "type": "step"}]}, "isCollapsed": false}, "style": {"width": 487, "height": 306}, "resizing": false, "selected": false, "positionAbsolute": {"x": 417.53612302022213, "y": 228.60709744219753}, "dragging": true}, {"id": "3", "type": "export", "position": {"x": 171.28302237599206, "y": 307.94365582341436}, "data": {"label": "split_fn", "exportType": "input", "isResource": true, "isCollapsed": false, "isEditing": false}, "width": 150, "height": 24, "selected": false, "positionAbsolute": {"x": 171.28302237599206, "y": 307.94365582341436}, "dragging": true}, {"id": "4", "type": "export", "position": {"x": 169.3283191390251, "y": 276.5582983666894}, "data": {"label": "Input", "exportType": "input", "isResource": false, "isCollapsed": false, "isEditing": false}, "width": 150, "height": 24, "selected": false, "positionAbsolute": {"x": 169.3283191390251, "y": 276.5582983666894}, "dragging": true}, {"id": "5", "type": "export", "position": {"x": 980.069924178094, "y": 281.42455113963734}, "data": {"label": "Output", "exportType": "output", "isResource": false, "isCollapsed": false}, "width": 150, "height": 24, "positionAbsolute": {"x": 980.069924178094, "y": 281.42455113963734}}, {"id": "6", "type": "resource", "position": {"x": 65.7467508136707, "y": 199.57544189414847}, "data": {"name": "RMBGModel", "parameters": {"model_name": {"type": "string", "description": "The name of the image processor.", "value": "briaai/RMBG-1.4"}}, "category": "Custom", "label": "RMBGModel", "key": 19, "isCollapsed": true}, "width": 150, "height": 24, "selected": false, "positionAbsolute": {"x": 483.28287383389284, "y": 428.182539336346}, "dragging": true, "parentId": "2", "hidden": false}, {"id": "7", "type": "step", "position": {"x": 78.30458080504991, "y": 20.95433445659424}, "data": {"name": "Split", "parameters": {"split_fn": {"type": "resource"}}, "inputs": ["in"], "outputs": ["A", "B"], "category": "Filtering", "label": "Split", "key": 2, "isCollapsed": false}, "width": 150, "height": 120, "selected": false, "positionAbsolute": {"x": 495.84070382527204, "y": 249.56143189879177}, "dragging": true, "parentId": "2", "hidden": false}], "edges": [{"source": "4", "sourceHandle": "in", "target": "2", "targetHandle": "1", "data": {}, "id": "reactflow__edge-4in-21"}, {"source": "6", "sourceHandle": "resource", "target": "0", "targetHandle": "model", "data": {}, "id": "reactflow__edge-6resource-0model"}, {"source": "0", "sourceHandle": "out", "target": "2", "targetHandle": "0_inner", "data": {}, "id": "reactflow__edge-0out-20_inner"}, {"source": "2", "sourceHandle": "0", "target": "5", "targetHandle": "out", "data": {}, "id": "reactflow__edge-20-5out"}, {"source": "3", "sourceHandle": "in", "target": "2", "targetHandle": "0", "data": {}, "id": "reactflow__edge-3in-20"}, {"source": "2", "sourceHandle": "1_inner", "target": "7", "targetHandle": "in", "data": {}, "id": "reactflow__edge-21_inner-7in"}, {"source": "2", "sourceHandle": "0_inner", "target": "7", "targetHandle": "split_fn", "data": {}, "id": "reactflow__edge-20_inner-7split_fn"}, {"source": "7", "sourceHandle": "A", "target": "0", "targetHandle": "in", "data": {}, "id": "reactflow__edge-7A-0in"}]}

workflow/custom_nodes/model.py

@@ -0,0 +1,24 @@
+from graphbook import param, resource
+from transformers import AutoModelForImageSegmentation
+
+
+@resource("BackgroundRemoval/RMBGModel")
+@param(
+    "model_name",
+    "string",
+    description="The name of the RMBG model.",
+    default="briaai/RMBG-1.4",
+)
+@param(
+    "use_cuda",
+    "boolean",
+    description="Whether to use CUDA acceleration.",
+    default=True,
+)
+def rmbg_model(self):
+    model = AutoModelForImageSegmentation.from_pretrained(
+        self.model_name, trust_remote_code=True
+    )
+    if self.use_cuda:
+        return model.to("cuda")
+    return model

workflow/custom_nodes/rmbg.py

@@ -0,0 +1,111 @@
+from graphbook import steps
+from transformers import AutoModelForImageSegmentation
+import torchvision.transforms.functional as F
+import torch.nn.functional
+import torch
+import os
+import os.path as osp
+
+
+class RemoveBackground(steps.BatchStep):
+    RequiresInput = True
+    Outputs = ["out"]
+    Parameters = {
+        "batch_size": {
+            "type": "number",
+            "description": "The batch size for background removal.",
+            "default": 8,
+        },
+        "model": {
+            "type": "resource",
+            "description": "The model to use for background removal.",
+        },
+        "output_dir": {
+            "type": "string",
+            "description": "The directory to save the output images.",
+            "default": "/tmp/output",
+        },
+    }
+    Category = "BackgroundRemoval"
+
+    def __init__(self, batch_size: int, model: AutoModelForImageSegmentation, output_dir: str):
+        super().__init__(batch_size, "image")
+        self.model = model
+        self.output_dir = output_dir
+        if self.output_dir:
+            os.makedirs(self.output_dir, exist_ok=True)
+
+    def load_fn(self, item: dict) -> torch.Tensor:
+        return load_image_as_tensor(item)
+
+    def dump_fn(self, t: torch.Tensor, output_path: str):
+        save_image_to_disk(t, output_path)
+
+    @torch.no_grad()
+    def on_item_batch(self, tensors, items, notes):
+        def get_output_path(note):
+            label = ""
+            if note["labels"] == 0:
+                label = "cat"
+            else:
+                label = "dog"
+            return osp.join(self.output_dir, label, f"{note['image']['shm_id']}.jpg")
+
+        og_sizes = [t.shape[1:] for t in tensors]
+
+        images = [
+            F.normalize(
+                torch.nn.functional.interpolate(
+                    torch.unsqueeze(image, 0), size=[1024, 1024], mode="bilinear"
+                ),
+                [0.5, 0.5, 0.5],
+                [1.0, 1.0, 1.0],
+            )
+            for image in tensors
+        ]
+        images = torch.stack(images).to("cuda")
+        images = torch.squeeze(images, 1)
+        tup = self.model(images)
+        result = tup[0][0]
+        ma = torch.max(result)
+        mi = torch.min(result)
+        result = (result - mi) / (ma - mi)
+        resized = [
+            torch.squeeze(
+                torch.nn.functional.interpolate(
+                    torch.unsqueeze(image, 0), size=og_size, mode="bilinear"
+                ),
+                0,
+            ).cpu()
+            for image, og_size in zip(result, og_sizes)
+        ]
+        paths = [
+            get_output_path(note) for note in notes
+        ]
+        removed_bg = list(zip(resized, paths))
+        for path, note in zip(paths, notes):
+            masks = note["masks"]
+            if masks is None:
+                masks = []
+            masks.append({"value": path, "type": "image"})
+            note["masks"] = masks
+
+        return removed_bg
+
+
+def load_image_as_tensor(item: dict) -> torch.Tensor:
+    im = item["value"]
+    image = F.to_tensor(im)
+    if image.shape[0] == 1:
+        image = image.repeat(3, 1, 1)
+    elif image.shape[0] == 4:
+        image = image[:3]
+
+    return image
+
+
+def save_image_to_disk(t: torch.Tensor, output_path: str):
+    dir = osp.dirname(output_path)
+    os.makedirs(dir, exist_ok=True)
+    img = F.to_pil_image(t)
+    img.save(output_path)