Adding documentation

workflow/RemoveBG.json

@@ -1 +1 @@
-{"version": "0", "type": "workflow", "nodes": [{"id": "2", "position": {"x": 187.5, "y": 52.5}, "type": "step", "data": {"label": "HuggingfaceDataset", "name": "HuggingfaceDataset", "parameters": {"dataset_id": {"type": "string", "description": "The dataset ID from Huggingface", "value": "microsoft/cats_vs_dogs"}, "split": {"type": "string", "default": "train", "description": "The split of the dataset to use", "value": "train"}, "shuffle": {"type": "boolean", "default": false, "description": "Whether to shuffle the dataset", "value": true}, "log_data": {"type": "boolean", "default": true, "description": "Whether to log the outputs as JSON to the node UI", "value": false}, "image_columns": {"type": "list[string]", "description": "The columns in the dataset that contain images. This is to let Graphbook know how to display the images in the UI.", "required": false, "value": ["image"]}, "kwargs": {"type": "dict", "description": "Additional keyword arguments to pass to the dataset", "required": false, "value": []}}, "inputs": [], "outputs": ["out"], "category": "Huggingface", "isCollapsed": false}}, {"id": "3", "type": "resource", "position": {"x": 187.5, "y": 282.5}, "data": {"name": "RMBGModel", "parameters": {"model_name": {"type": "string", "default": "briaai/RMBG-1.4", "required": false, "description": "The name of the RMBG model.", "value": "briaai/RMBG-1.4"}, "use_cuda": {"type": "boolean", "default": true, "required": false, "description": "Whether to use CUDA acceleration.", "value": false}}, "category": "BackgroundRemoval", "label": "RMBGModel", "isCollapsed": false}}, {"id": "4", "type": "step", "position": {"x": 444.5, "y": 52.5}, "data": {"name": "RemoveBackground", "parameters": {"batch_size": {"type": "number", "description": "The batch size for background removal.", "default": 8, "value": 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", "value": "/media/sam/shared/__DELETE"}}, "inputs": ["in"], "outputs": ["out"], "category": "BackgroundRemoval", "label": "RemoveBackground", "isCollapsed": false}}], "edges": [{"source": "2", "sourceHandle": "out", "target": "4", "targetHandle": "in", "data": {}, "id": "reactflow__edge-2out-4in"}, {"source": "3", "sourceHandle": "resource", "target": "4", "targetHandle": "model", "data": {}, "id": "reactflow__edge-3resource-4model"}]}
+{"version": "0", "type": "workflow", "nodes": [{"id": "2", "position": {"x": 187.5, "y": 52.5}, "type": "step", "data": {"label": "HuggingfaceDataset", "name": "HuggingfaceDataset", "parameters": {"dataset_id": {"type": "string", "description": "The dataset ID from Huggingface", "value": "microsoft/cats_vs_dogs"}, "split": {"type": "string", "default": "train", "description": "The split of the dataset to use", "value": "train"}, "shuffle": {"type": "boolean", "default": false, "description": "Whether to shuffle the dataset", "value": true}, "log_data": {"type": "boolean", "default": true, "description": "Whether to log the outputs as JSON to the node UI", "value": false}, "image_columns": {"type": "list[string]", "description": "The columns in the dataset that contain images. This is to let Graphbook know how to display the images in the UI.", "required": false, "value": ["image"]}, "kwargs": {"type": "dict", "description": "Additional keyword arguments to pass to the dataset", "required": false, "value": []}}, "inputs": [], "outputs": ["out"], "category": "Huggingface", "isCollapsed": false}}, {"id": "3", "type": "resource", "position": {"x": 187.5, "y": 282.5}, "data": {"name": "RMBGModel", "parameters": {"model_name": {"type": "string", "default": "briaai/RMBG-1.4", "required": false, "description": "The name of the RMBG model.", "value": "briaai/RMBG-1.4"}, "use_cuda": {"type": "boolean", "default": true, "required": false, "description": "Whether to use CUDA acceleration.", "value": false}}, "category": "BackgroundRemoval", "label": "RMBGModel", "isCollapsed": false}}, {"id": "4", "type": "step", "position": {"x": 444.5, "y": 52.5}, "data": {"name": "RemoveBackground", "parameters": {"batch_size": {"type": "number", "description": "The batch size for background removal.", "default": 8, "value": 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", "value": "/tmp/output"}}, "inputs": ["in"], "outputs": ["out"], "category": "BackgroundRemoval", "label": "RemoveBackground", "isCollapsed": false}}], "edges": [{"source": "2", "sourceHandle": "out", "target": "4", "targetHandle": "in", "data": {}, "id": "reactflow__edge-2out-4in"}, {"source": "3", "sourceHandle": "resource", "target": "4", "targetHandle": "model", "data": {}, "id": "reactflow__edge-3resource-4model"}]}

workflow/custom_nodes/model.py

@@ -16,6 +16,14 @@ from transformers import AutoModelForImageSegmentation
     default=True,
 )
 def rmbg_model(self):
+    """
+    Loads the background removal model from the Hugging Face model hub.
+    By default, we are using [briaai/RMBG-1.4](https://huggingface.co/briaai/RMBG-1.4).
+    
+    Args:
+        model_name (str): The name of the RMBG model.
+        use_cuda (bool): Whether to use CUDA acceleration.
+    """
     model = AutoModelForImageSegmentation.from_pretrained(
         self.model_name, trust_remote_code=True
     )

workflow/custom_nodes/rmbg.py

@@ -8,6 +8,14 @@ import os.path as osp
 
 
 class RemoveBackground(steps.BatchStep):
+    """
+    Executes background removal on a batch of images using the provided model.
+    
+    Args:
+        batch_size (int): The batch size (number of images) given to the background removal model at a time
+        model (resource): The model to use for background removal. Use the resource `BackgroundRemoval/RMBGModel` to load the model
+        output_dir (str): The directory to save the output images to
+    """
     RequiresInput = True
     Outputs = ["out"]
     Parameters = {

workflow/docs/RemoveBG.md

@@ -0,0 +1,15 @@
+<div align="center">
+    <img src="https://github.com/graphbookai/graphbook/blob/main/docs/_static/graphbook.png?raw=true" alt="Graphbook Logo" height="64">
+    <h1 style="margin: 0 0 10px 0">RMBG Graphbook</h1>
+</div>
+
+This is a [Graphbook](https://github.com/graphbookai/graphbook) app that allows you to remove the background from any image dataset from Hugging Face.
+
+* Model used: [briaai/RMBG-1.4](https://huggingface.co/briaai/RMBG-1.4)
+* Dataset used: [microsoft/cats_vs_dogs](https://huggingface.co/datasets/microsoft/cats_vs_dogs)
+
+You may swap datasets using the 🤗 Hugging Face extension or by simply changing the `dataset_id` in the HuggingfaceDataset node.
+
+* [Repo](https://github.com/graphbookai/graphbook) 
+* [Docs](https://docs.graphbook.ai/)
+* [Examples](https://github.com/graphbookai/graphbook-examples)