url test

app.py

@@ -20,42 +20,52 @@ def resize_image(image):
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
+def get_url_image(url):
+    headers = {'User-Agent': 'gradio-app'}
+    response = requests.get(url, headers=headers)
+    return BytesIO(response.content)
+
 def load_image(image_source):
     if isinstance(image_source, str):  # Check if input is a URL
-        response = requests.get(image_source)
-        image = Image.open(BytesIO(response.content))
+        print(f"Loading image from URL: {image_source}")
+        image = Image.open(get_url_image(image_source))
     else:
+        print("Loading image from file upload")
         image = Image.fromarray(image_source)
     return image
 
 def process(image_source):
-    # Load and prepare input
-    orig_image = load_image(image_source)
-    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)
-    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()
+    try:
+        # Load and prepare input
+        orig_image = load_image(image_source)
+        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)
+        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)
-    # Post-process
-    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)
-    # Image to PIL
-    im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-    pil_im = Image.fromarray(np.squeeze(im_array))
-    # Paste the mask on the original image
-    new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
-    new_im.paste(orig_image, mask=pil_im)
+        # Inference
+        result = net(im_tensor)
+        # Post-process
+        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)
+        # Image to PIL
+        im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
+        pil_im = Image.fromarray(np.squeeze(im_array))
+        # Paste the mask on the original image
+        new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
+        new_im.paste(orig_image, mask=pil_im)
 
-    return new_im
+        return new_im
+    except Exception as e:
+        print(f"Error during processing: {e}")
+        return None
 
 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>