Update app.py

app.py

@@ -3,59 +3,66 @@ import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 import gradio as gr
-from briarmbg import BriaRMBG
 from PIL import Image
+from briarmbg import BriaRMBG
+import io
 
 
+# Load the model
 net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 net.to(device)
-net.eval()    
+net.eval()
+
 
 def resize_image(image):
     image = image.convert('RGB')
     model_input_size = (1024, 1024)
-    image = image.resize(model_input_size, Image.BILINEAR)
-    return image
+    return image.resize(model_input_size, Image.BILINEAR)
+
 
 def process(image):
     orig_image = Image.fromarray(image)
     w, h = orig_image.size
     image = resize_image(orig_image)
     im_np = np.array(image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = torch.unsqueeze(im_tensor, 0) / 255.0
     im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
     if torch.cuda.is_available():
         im_tensor = im_tensor.cuda()
 
+    # Inference
     result = net(im_tensor)
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
-    result = (result - result.min()) / (result.max() - result.min())
-    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-
-    # Fix: Convert to grayscale image
-    pil_mask = Image.fromarray(result_array.squeeze(), mode="L")
+    result_array = ((result - result.min()) / (result.max() - result.min()) * 255).cpu().data.numpy().astype(np.uint8)
 
+    # Add mask to original image
+    pil_mask = Image.fromarray(np.squeeze(result_array))
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
-    return new_im
 
-gr.Markdown("## BRIA RMBG 1.4")
-gr.HTML('''
-  <p style="margin-bottom: 10px; font-size: 94%">
-    This is a demo for BRIA RMBG 1.4 that uses
-    <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone. 
-  </p>
-''')
+    # Convert to bytes for download
+    buffer = io.BytesIO()
+    new_im.save(buffer, format="PNG")
+    buffer.seek(0)
+
+    return new_im, buffer
 
-title = "Background Removal"
-description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br> 
-For test upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>. To purchase a commercial license, simply click <a href='https://go.bria.ai/3ZCBTLH' target='_blank'><b>Here</b></a>. <br>
-"""
 
-examples = [['./input.jpg']]
+def process_with_download(image):
+    new_image, buffer = process(image)
+    return new_image, ("background_removed.png", buffer)
 
-demo = gr.Interface(fn=process, inputs="image", outputs="image", examples=examples, title=title, description=description)
+
+demo = gr.Interface(
+    fn=process_with_download,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(label="Processed Image"), gr.File(label="Download")],
+    title="Background Removal with BRIA RMBG 1.4",
+    description="Upload an image to remove the background and download the result."
+)
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()
+