Update app.py

app.py

@@ -5,69 +5,71 @@ from torchvision.transforms.functional import normalize
 import gradio as gr
 from briarmbg import BriaRMBG
 from PIL import Image
-import tempfile
+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()
 
-# Resize input image
+
 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
 
-# Process the image
-def process(image, progress=gr.Progress()):
-    progress(0.2)  # 20% progress for loading
-    orig_image = Image.fromarray(image)
-    w, h = orig_image.size
 
-    # Resize image
+def process(image):
+    orig_image = Image.fromarray(image)
+    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).unsqueeze(0) / 255.0
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
+    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
     result = net(im_tensor)
-    progress(0.7)  # 70% progress during inference
-    
-    # Post-process
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
-    result = (result - result.min()) / (result.max() - result.min())
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)
 
-    # Convert to PIL image with alpha mask
     result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-    pil_mask = Image.fromarray(result_array)
+    pil_mask = Image.fromarray(result_array, mode="L")
+
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
 
-    # Save the image for download
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-    new_im.save(temp_file.name, format='PNG')
-
-    progress(1.0)  # 100% complete
-    return new_im, temp_file.name
-
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("## BRIA RMBG 1.4 - Background Remover")
-    gr.HTML("<p>Upload your image to remove the background.</p>")
-    
-    with gr.Row():
-        input_image = gr.Image(type="numpy", label="Upload Image")
-        output_image = gr.Image(type="pil", label="Processed Image")
-    
-    download_button = gr.File(label="Download Processed Image")
-    
-    input_image.change(fn=process, inputs=input_image, outputs=[output_image, download_button])
-
-# Run the app on port 7860
-demo.launch(server_name="0.0.0.0", server_port=7860)
+    output_buffer = io.BytesIO()
+    new_im.save(output_buffer, format="PNG")
+    output_buffer.seek(0)
+
+    return new_im, output_buffer
+
+
+def process_with_download(image):
+    new_im, output_buffer = process(image)
+    return new_im, ("output.png", output_buffer, "image/png")
+
+
+gr.Markdown("## BRIA RMBG 1.4 - Background Remover")
+
+demo = gr.Interface(
+    fn=process_with_download,
+    inputs=gr.Image(type="numpy"),
+    outputs=[
+        gr.Image(type="pil", label="Result"),
+        gr.File(label="Download Image")
+    ],
+    live=True,
+    title="BRIA RMBG Background Remover",
+    description="Upload an image to remove the background automatically.",
+)
 
+if __name__ == "__main__":
+    demo.launch()