Update app.py

app.py

@@ -1,78 +1,71 @@
-import gradio as gr
-import cv2
-import torch
-import numpy as np
-from torchvision import transforms
-
-title = "移除背景 Demo"
-description = "自動去除背景圖片."
-article = "<p style='text-align: center'><a href='https://news.machinelearning.sg/posts/beautiful_profile_pics_remove_background_image_with_deeplabv3/'>Blog</a> | <a href='https://github.com/eugenesiow/practical-ml'>Github Repo</a></p>"
-
-
-def make_transparent_foreground(pic, mask):
-    # split the image into channels
-    b, g, r = cv2.split(np.array(pic).astype('uint8'))
-    # add an alpha channel with and fill all with transparent pixels (max 255)
-    a = np.ones(mask.shape, dtype='uint8') * 255
-    # merge the alpha channel back
-    alpha_im = cv2.merge([b, g, r, a], 4)
-    # create a transparent background
-    bg = np.zeros(alpha_im.shape)
-    # setup the new mask
-    new_mask = np.stack([mask, mask, mask, mask], axis=2)
-    # copy only the foreground color pixels from the original image where mask is set
-    foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8)
-
-    return foreground
-
-
-def remove_background(input_image):
-    preprocess = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-
-    input_tensor = preprocess(input_image)
-    input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model
-
-    # move the input and model to GPU for speed if available
-    if torch.cuda.is_available():
-        input_batch = input_batch.to('cuda')
-        model.to('cuda')
-
-    with torch.no_grad():
-        output = model(input_batch)['out'][0]
-    output_predictions = output.argmax(0)
-
-    # create a binary (black and white) mask of the profile foreground
-    mask = output_predictions.byte().cpu().numpy()
-    background = np.zeros(mask.shape)
-    bin_mask = np.where(mask, 255, background).astype(np.uint8)
-
-    foreground = make_transparent_foreground(input_image, bin_mask)
-
-    return foreground, bin_mask
-
-
-def inference(img):
-    foreground, _ = remove_background(img)
-    return foreground
-
-
-torch.hub.download_url_to_file('https://pbs.twimg.com/profile_images/691700243809718272/z7XZUARB_400x400.jpg',
-                               'capture.png')
-torch.hub.download_url_to_file('https://hai.stanford.edu/sites/default/files/styles/person_medium/public/2020-03/hai_1512feifei.png?itok=INFuLABp',
-                               'girl1.png')
-model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True)
-model.eval()
-
-gr.Interface(
-    inference,
-    gr.Image(type="pil", label="Input"),
-    gr.Image(type="pil", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[['removebg/girl1.png'],['removebg/girl2.png'],['removebg/girl3.png'],['removebg/gonfu1.jpg'],['removebg/angel.png']],
-    #enable_queue=True,
-).launch(debug=False)
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+from PIL import Image
+import requests
+from io import BytesIO
+import gradio as gr
+
+title = "移除背景 Demo"
+description = "上傳圖片 ,自動去除背景."
+
+# Set up CUDA if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_float32_matmul_precision("high")
+
+# Load the model
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to(device)
+
+# Define image transformations
+transform_image = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+
+def load_img(image_path_or_url):
+    if image_path_or_url.startswith('http'):
+        response = requests.get(image_path_or_url)
+        img = Image.open(BytesIO(response.content))
+    else:
+        img = Image.open(image_path_or_url)
+    return img.convert("RGB")
+
+def process(image):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+
+    # Create a new image with transparency
+    transparent_image = Image.new("RGBA", image.size)
+    transparent_image.paste(image, (0, 0))
+    transparent_image.putalpha(mask)  # Apply mask to the new image
+
+    return transparent_image  # Return the new transparent image
+
+def remove_background_gradio(image):
+    processed_img = process(image)
+    return processed_img
+
+
+# Create the Gradio interface with drag-and-drop and paste functionality
+demo = gr.Interface(
+    fn=remove_background_gradio,
+    inputs = gr.Image(type="pil"),  # Remove 'source' argument
+    outputs = gr.Image(type="pil"),
+    title = title,
+    description = description,
+  #  examples=[['girl1.png'],['girl2.png'],['girl3.png'],['gonfu1.jpg'],['removebg/angel.png']],
+)
+
+demo.launch(share=True)  # Launch the interface and get a shareable link