not-lain commited on
Commit
e546fea
β€’
1 Parent(s): 770ca42

🌘wπŸŒ–

Browse files
Files changed (3) hide show
  1. README.md +1 -1
  2. app.py +47 -0
  3. requirements.txt +16 -0
README.md CHANGED
@@ -1,6 +1,6 @@
1
  ---
2
  title: Background Removal
3
- emoji: πŸƒ
4
  colorFrom: purple
5
  colorTo: indigo
6
  sdk: gradio
 
1
  ---
2
  title: Background Removal
3
+ emoji: 🌘wπŸŒ–
4
  colorFrom: purple
5
  colorTo: indigo
6
  sdk: gradio
app.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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('zhengpeng7/BiRefNet', trust_remote_code=True,device="auto",torch_dtype=torch.float16)
12
+
13
+ transform_image = transforms.Compose([
14
+ transforms.Resize((1024, 1024)),
15
+ transforms.ToTensor(),
16
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
17
+ ])
18
+
19
+
20
+
21
+ @spaces.GPU
22
+ def fn(image):
23
+ im = load_img(image)
24
+ im = im.convert('RGB')
25
+ image = load_img(im)
26
+ input_images = transform_image(image).unsqueeze(0).to('cuda')
27
+ # Prediction
28
+ with torch.no_grad():
29
+ preds = birefnet(input_images)[-1].sigmoid().cpu()
30
+ pred = preds[0].squeeze()
31
+ pred_pil = transforms.ToPILImage()(pred)
32
+ out = (pred_pil , im)
33
+ return out
34
+
35
+ slider1 = ImageSlider(label="birefnet", type="pil")
36
+ slider2 = ImageSlider(label="RMBG", type="pil")
37
+ image = gr.Image(label="Upload an image")
38
+ text = gr.Textbox(label="Paste an image URL")
39
+
40
+
41
+ tab1 = gr.Interface(fn,inputs= image, outputs= slider1, api_name="image")
42
+ tab2 = gr.Interface(fn,inputs= text, outputs= slider2, api_name="text")
43
+
44
+ demo = gr.TabbedInterface([tab1,tab2],["image","text"],title="RMBG with image slider")
45
+
46
+ if __name__ == "__main__":
47
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ accelerate
3
+ opencv-python
4
+ spaces
5
+ pillow
6
+ numpy
7
+ timm
8
+ kornia
9
+ prettytable
10
+ typing
11
+ scikit-image
12
+ huggingface_hub
13
+ transformers>=4.39.1
14
+ gradio
15
+ gradio_imageslider
16
+ loadimg>=0.1.1