Update app.py

app.py

@@ -1,68 +1,90 @@
-import numpy as np 
+import numpy as np
 import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 import gradio as gr
-from PIL import Image
 from briarmbg import BriaRMBG
-import io
+import PIL
+from PIL import Image
+import tempfile
+import os
 
 
-# Load the model
+# Load the pre-trained 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()    
 
+# Resize the input image for model compatibility
 def resize_image(image):
     image = image.convert('RGB')
     model_input_size = (1024, 1024)
-    return image.resize(model_input_size, Image.BILINEAR)
-
+    image = image.resize(model_input_size, Image.BILINEAR)
+    return image
 
+# Background removal process
 def process(image):
+    # Prepare the input
     orig_image = Image.fromarray(image)
-    w, h = orig_image.size
+    w, h = orig_im_size = 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)
-    im_tensor = torch.unsqueeze(im_tensor, 0) / 255.0
+    im_tensor = torch.unsqueeze(im_tensor, 0)
+    im_tensor = torch.divide(im_tensor, 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
+    # Inference with the model
     result = net(im_tensor)
+    
+    # Post-process the result
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
-    result_array = ((result - result.min()) / (result.max() - result.min()) * 255).cpu().data.numpy().astype(np.uint8)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)  # Normalize result
 
-    # Add mask to original image
+    # Convert the result to an image
+    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
     pil_mask = Image.fromarray(np.squeeze(result_array))
+
+    # Add the mask as alpha channel to the original image
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
 
-    # Convert to bytes for download
-    buffer = io.BytesIO()
-    new_im.save(buffer, format="PNG")
-    buffer.seek(0)
+    # Save the processed image to a temporary file
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    new_im.save(temp_file, format='PNG')
+    temp_file.close()  # Ensure the file is closed before Gradio uses it
 
-    return new_im, buffer
+    return temp_file.name  # Return the path to the temporary file for downloading
 
+# Gradio interface setup
+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 a backbone.
+</p>''')
 
-def process_with_download(image):
-    new_image, buffer = process(image)
-    return new_image, ("background_removed.png", 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 testing, 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'],]
 
+# Modify the interface to use live updates and file download
 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."
+    fn=process,  # The function to process the image
+    inputs=gr.Image(type="numpy"),  # Input type (image)
+    outputs=gr.File(label="Download Processed Image"),  # Output as a file (download button)
+    examples=examples,  # Example images for users to try
+    title=title,  # Title of the app
+    description=description,  # Description of the app
+    live=True  # Automatically processes when an image is uploaded
 )
 
 if __name__ == "__main__":
-    demo.launch()
-
+    demo.launch(share=False)