Update app.py

app.py

@@ -8,12 +8,15 @@ import PIL
 from PIL import Image
 import tempfile
 import os
+import time
 
 # Load the pre-trained model
+print("Loading 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()
+print(f"Model loaded on {device}")
 
 # Resize the input image for model compatibility
 def resize_image(image):
@@ -23,8 +26,14 @@ def resize_image(image):
     return image
 
 # Background removal process
-def process(image):
+def process(image, progress=gr.Progress()):
+    if image is None:
+        return None, None
+    
+    progress(0, desc="Starting processing...")
+    
     # Prepare the input
+    progress(0.1, desc="Preparing image...")
     orig_image = Image.fromarray(image)
     w, h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
@@ -33,6 +42,8 @@ def process(image):
     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])
+    
+    progress(0.3, desc="Processing with AI model...")
     if torch.cuda.is_available():
         im_tensor = im_tensor.cuda()
     
@@ -40,6 +51,7 @@ def process(image):
     with torch.no_grad():
         result = net(im_tensor)
     
+    progress(0.6, desc="Post-processing...")
     # Post-process the result
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
     ma = torch.max(result)
@@ -54,13 +66,18 @@ def process(image):
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
     
+    progress(0.8, desc="Preparing download...")
     # Save the processed image to a temporary file
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
     new_im.save(temp_file.name, format='PNG')
     temp_file.close()
     
-    # Return the path to the temporary file for downloading
-    return temp_file.name
+    # Convert to numpy array for display
+    output_array = np.array(new_im.convert("RGBA"))
+    
+    progress(1.0, desc="Done!")
+    # Return both the image for display and the path for download
+    return output_array, temp_file.name
 
 # Gradio interface setup
 title = "Background Removal"
@@ -69,15 +86,37 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 examples = [['./input.jpg']]
 
 # Create the Gradio interface
-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
-)
+with gr.Blocks() as demo:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="numpy", label="Upload Image")
+            process_btn = gr.Button("Remove Background", variant="primary")
+        
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="numpy", label="Result")
+            download_btn = gr.File(label="Download Image")
+    
+    # Set up example images
+    gr.Examples(examples, inputs=input_image)
+    
+    # Set up processing logic
+    process_btn.click(
+        fn=process,
+        inputs=input_image,
+        outputs=[output_image, download_btn],
+        show_progress="full"
+    )
+    
+    # Also process automatically when image is uploaded
+    input_image.change(
+        fn=process,
+        inputs=input_image,
+        outputs=[output_image, download_btn],
+        show_progress="full"
+    )
 
 if __name__ == "__main__":
     demo.launch(share=False)