gaur3009 commited on
Commit
e342ca5
·
verified ·
1 Parent(s): 0d98a52

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +62 -0
app.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from gradio_imageslider import ImageSlider
3
+ from loadimg import load_img
4
+ import spaces
5
+ from transformers import AutoModelForImageSegmentation
6
+ import torch
7
+ from torchvision import transforms
8
+
9
+ torch.set_float32_matmul_precision(["high", "highest"][0])
10
+
11
+ birefnet = AutoModelForImageSegmentation.from_pretrained(
12
+ "ZhengPeng7/BiRefNet", trust_remote_code=True
13
+ )
14
+ birefnet.to("cuda")
15
+ transform_image = transforms.Compose(
16
+ [
17
+ transforms.Resize((1024, 1024)),
18
+ transforms.ToTensor(),
19
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
20
+ ]
21
+ )
22
+
23
+
24
+ @spaces.GPU
25
+ def fn(image):
26
+ im = load_img(image, output_type="pil")
27
+ im = im.convert("RGB")
28
+ image_size = im.size
29
+ origin = im.copy()
30
+ image = load_img(im)
31
+ input_images = transform_image(image).unsqueeze(0).to("cuda")
32
+ # Prediction
33
+ with torch.no_grad():
34
+ preds = birefnet(input_images)[-1].sigmoid().cpu()
35
+ pred = preds[0].squeeze()
36
+ pred_pil = transforms.ToPILImage()(pred)
37
+ mask = pred_pil.resize(image_size)
38
+ image.putalpha(mask)
39
+ return (image, origin)
40
+
41
+
42
+ slider1 = ImageSlider(label="birefnet", type="pil")
43
+ slider2 = ImageSlider(label="birefnet", type="pil")
44
+ image = gr.Image(label="Upload an image")
45
+ text = gr.Textbox(label="Paste an image URL")
46
+
47
+
48
+ chameleon = load_img("chameleon.jpg", output_type="pil")
49
+
50
+ url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
51
+ tab1 = gr.Interface(
52
+ fn, inputs=image, outputs=slider1, examples=[chameleon], api_name="image"
53
+ )
54
+
55
+ tab2 = gr.Interface(fn, inputs=text, outputs=slider2, examples=[url], api_name="text")
56
+
57
+
58
+ demo = gr.TabbedInterface(
59
+ [tab1, tab2], ["image", "text"], title="birefnet for background removal"
60
+ )
61
+
62
+ if __name__ == "__main__":