ZhengPeng7 hysts HF staff commited on
Commit
e797135
1 Parent(s): b0a9abf

ZeroGPU and ImageSlider (#1)

Browse files

- ZeroGPU and ImageSlider (dd31e60100ef70bff8cad77af8ea1f4460a9320d)


Co-authored-by: hysts <[email protected]>

Files changed (3) hide show
  1. app.py +12 -10
  2. config.py +3 -1
  3. requirements.txt +1 -2
app.py CHANGED
@@ -6,6 +6,10 @@ from PIL import Image
6
  import torch
7
  from torchvision import transforms
8
  import gradio as gr
 
 
 
 
9
 
10
  from models.baseline import BiRefNet
11
  from config import Config
@@ -35,20 +39,22 @@ class ImagePreprocessor():
35
  return image
36
 
37
 
38
- model = BiRefNet(bb_pretrained=False).to(device)
39
  state_dict = './BiRefNet_ep580.pth'
40
  if os.path.exists(state_dict):
41
- birefnet_dict = torch.load(state_dict, map_location=device)
42
  unwanted_prefix = '_orig_mod.'
43
  for k, v in list(birefnet_dict.items()):
44
  if k.startswith(unwanted_prefix):
45
  birefnet_dict[k[len(unwanted_prefix):]] = birefnet_dict.pop(k)
46
  model.load_state_dict(birefnet_dict)
 
47
  model.eval()
48
 
49
 
50
  # def predict(image_1, image_2):
51
  # images = [image_1, image_2]
 
52
  def predict(image, resolution='1024x1024'):
53
  # Image is a RGB numpy array.
54
  resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
@@ -74,17 +80,13 @@ def predict(image, resolution='1024x1024'):
74
  image_preds.append(
75
  cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
76
  )
77
- return image_preds[:] if len(images) > 1 else image_preds[0]
78
 
 
79
 
80
- examples = [[_] for _ in glob('materials/examples/*')][:]
81
 
82
- N = 1
83
- ipt = [gr.Image() for _ in range(N)]
84
- opt = [gr.Image() for _ in range(N)]
85
 
86
  # Add the option of resolution in a text box.
87
- ipt += [gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")]
88
  for idx_example, example in enumerate(examples):
89
  examples[idx_example].append('1024x1024')
90
  examples.append(examples[-1].copy())
@@ -92,8 +94,8 @@ examples[-1][1] = '512x512'
92
 
93
  demo = gr.Interface(
94
  fn=predict,
95
- inputs=ipt,
96
- outputs=opt,
97
  examples=examples,
98
  title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
99
  description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)'
 
6
  import torch
7
  from torchvision import transforms
8
  import gradio as gr
9
+ import spaces
10
+ from gradio_imageslider import ImageSlider
11
+
12
+ torch.jit.script = lambda f: f
13
 
14
  from models.baseline import BiRefNet
15
  from config import Config
 
39
  return image
40
 
41
 
42
+ model = BiRefNet(bb_pretrained=False)
43
  state_dict = './BiRefNet_ep580.pth'
44
  if os.path.exists(state_dict):
45
+ birefnet_dict = torch.load(state_dict, map_location="cpu")
46
  unwanted_prefix = '_orig_mod.'
47
  for k, v in list(birefnet_dict.items()):
48
  if k.startswith(unwanted_prefix):
49
  birefnet_dict[k[len(unwanted_prefix):]] = birefnet_dict.pop(k)
50
  model.load_state_dict(birefnet_dict)
51
+ model = model.to(device)
52
  model.eval()
53
 
54
 
55
  # def predict(image_1, image_2):
56
  # images = [image_1, image_2]
57
+ @spaces.GPU
58
  def predict(image, resolution='1024x1024'):
59
  # Image is a RGB numpy array.
60
  resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
 
80
  image_preds.append(
81
  cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
82
  )
 
83
 
84
+ return image, image_preds[0]
85
 
 
86
 
87
+ examples = [[_] for _ in glob('materials/examples/*')][:]
 
 
88
 
89
  # Add the option of resolution in a text box.
 
90
  for idx_example, example in enumerate(examples):
91
  examples[idx_example].append('1024x1024')
92
  examples.append(examples[-1].copy())
 
94
 
95
  demo = gr.Interface(
96
  fn=predict,
97
+ inputs=['image', gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")],
98
+ outputs=ImageSlider(),
99
  examples=examples,
100
  title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
101
  description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)'
config.py CHANGED
@@ -1,6 +1,8 @@
1
  import os
2
  import math
3
 
 
 
4
 
5
  class Config():
6
  def __init__(self) -> None:
@@ -97,7 +99,7 @@ class Config():
97
  self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
98
 
99
  # others
100
- self.device = [0, 'cpu'][1] # .to(0) = .to('cuda:0')
101
 
102
  self.batch_size_valid = 1
103
  self.rand_seed = 7
 
1
  import os
2
  import math
3
 
4
+ import torch
5
+
6
 
7
  class Config():
8
  def __init__(self) -> None:
 
99
  self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
100
 
101
  # others
102
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
103
 
104
  self.batch_size_valid = 1
105
  self.rand_seed = 7
requirements.txt CHANGED
@@ -1,6 +1,4 @@
1
- --extra-index-url https://download.pytorch.org/whl/cu118
2
  torch==2.0.1
3
- --extra-index-url https://download.pytorch.org/whl/cu118
4
  torchvision==0.15.2
5
  opencv-python==4.9.0.80
6
  tqdm==4.66.2
@@ -9,3 +7,4 @@ prettytable==3.10.0
9
  scipy==1.12.0
10
  scikit-image==0.22.0
11
  kornia==0.7.1
 
 
 
1
  torch==2.0.1
 
2
  torchvision==0.15.2
3
  opencv-python==4.9.0.80
4
  tqdm==4.66.2
 
7
  scipy==1.12.0
8
  scikit-image==0.22.0
9
  kornia==0.7.1
10
+ gradio_imageslider==0.0.18