Add colour picker

app.py

@@ -1,17 +1,12 @@
 import gradio as gr
-from gradio_imageslider import ImageSlider
 from loadimg import load_img
-import spaces
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
+from PIL import Image, ImageColor
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
-
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
-birefnet.to("cuda")
+birefnet = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -20,44 +15,35 @@ transform_image = transforms.Compose(
     ]
 )
 
-
-@spaces.GPU
-def fn(image):
+def fn(image, color):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    origin = im.copy()
     image = load_img(im)
-    input_images = transform_image(image).unsqueeze(0).to("cuda")
-    # Prediction
+    input_images = transform_image(image).unsqueeze(0)
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
     image.putalpha(mask)
-    return (image, origin)
-
-
-slider1 = ImageSlider(label="birefnet", type="pil")
-slider2 = ImageSlider(label="birefnet", type="pil")
+    color = ImageColor.getcolor(color, "RGBA")
+    new_image = Image.new("RGBA", image.size, color)
+    new_image.paste(image, (0, 0), image)
+    new_image = new_image.convert("RGB")
+    return new_image
+
+outimage1 = gr.Image(format="png", label="Result")
+outimage2 = gr.Image(format="png", label="Result")
 image = gr.Image(label="Upload an image")
 text = gr.Textbox(label="Paste an image URL")
-
-
+color1 = gr.ColorPicker(value="#FFFFFF", label="Color")
+color2 = gr.ColorPicker(value="#FFFFFF", label="Color")
 chameleon = load_img("chameleon.jpg", output_type="pil")
-
 url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
-tab1 = gr.Interface(
-    fn, inputs=image, outputs=slider1, examples=[chameleon], api_name="image"
-)
-
-tab2 = gr.Interface(fn, inputs=text, outputs=slider2, examples=[url], api_name="text")
-
-
-demo = gr.TabbedInterface(
-    [tab1, tab2], ["image", "text"], title="birefnet for background removal"
-)
+tab1 = gr.Interface(fn, inputs=[image, color1], outputs=outimage1, examples=[[chameleon, "#FFFFFF"]], api_name="image")
+tab2 = gr.Interface(fn, inputs=[text, color2], outputs=outimage2, examples=[[url, "#FFFFFF"]], api_name="text")
+demo = gr.TabbedInterface([tab1, tab2], ["image", "text"], title="birefnet for background coloring")
 
 if __name__ == "__main__":
     demo.launch()