Update app.py

app.py

@@ -23,8 +23,9 @@ def resize_image(image):
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-# Background removal process
-def process(image):
+# Background removal process with progress bar
+def process(image, progress=gr.Progress()):
+    progress(0.1)  # Start with 10% for loading
     # Prepare the input
     orig_image = Image.fromarray(image)
     w, h = orig_im_size = orig_image.size
@@ -41,6 +42,8 @@ def process(image):
     # Inference with the model
     result = net(im_tensor)
     
+    progress(0.5)  # Progress 50% during inference
+    
     # Post-process the result
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
     ma = torch.max(result)
@@ -60,7 +63,8 @@ def process(image):
     new_im.save(temp_file, format='PNG')
     temp_file.close()  # Ensure the file is closed before Gradio uses it
 
-    return temp_file.name  # Return the path to the temporary file for downloading
+    progress(1.0)  # Completion of the process
+    return orig_image, new_im, temp_file.name  # Return both the original and processed images for display
 
 # Gradio interface setup
 gr.Markdown("## BRIA RMBG 1.4")
@@ -71,20 +75,4 @@ gr.HTML('''<p style="margin-bottom: 10px; font-size: 94%">
 
 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,  # 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(share=False)
+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 li