Update app.py

app.py

@@ -4,63 +4,32 @@ import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 import gradio as gr
 from briarmbg import BriaRMBG
-import PIL
 from PIL import Image
 import tempfile
 
-
-# Load the pre-trained model
+# 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()
 
-# Resize the input image for model compatibility
+# Resize input image
 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)
 
-# Background removal process with progress bar
+# Process the image
 def process(image, progress=gr.Progress()):
-    progress(0.1)  # Start with 10% for loading
-    # Prepare the input
+    progress(0.2)  # 20% progress for loading
     orig_image = Image.fromarray(image)
-    w, h = orig_im_size = orig_image.size
+    w, h = orig_image.size
+
+    # Resize image
     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)
-    im_tensor = torch.divide(im_tensor, 255.0)
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(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 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)
-    mi = torch.min(result)
-    result = (result - mi) / (ma - mi)  # Normalize result
-
-    # 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)
-
-    # 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
-
-    progress(1.0)  # Completion of the process
-    return new_im, temp_file.name  # Return the processed image and d
+        im_tensor