init

.gitignore

@@ -0,0 +1 @@
+/venv

app.py

@@ -0,0 +1,24 @@
+from transformers import pipeline
+from PIL import Image
+import gradio as gr
+import spaces
+
+@spaces.GPU
+def remove_background_bria(input_image):
+    print(f"Removing background using bria for image")
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=0)
+    input_image = Image.fromarray(input_image)
+    output_image = pipe(input_image)
+    return output_image
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Background Removal with BRIA")
+    with gr.Row():
+        input_image = gr.Image(label="Input Image")
+        output_image = gr.Image(label="Output Image")
+    
+    remove_button = gr.Button("Remove Background")
+    remove_button.click(fn=remove_background_bria, inputs=input_image, outputs=output_image)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,8 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.2.0+cu118
+torchvision==0.17.0
+gradio==4.41.0
+spaces==0.29.3
+transformers==4.42.4
+scikit-image==0.24.0
+