ZhengPeng7 commited on
Commit
54516d1
1 Parent(s): 9f09c5a

Fix an error in loading image in tab url.

Browse files
Files changed (1) hide show
  1. app.py +13 -5
app.py CHANGED
@@ -13,6 +13,10 @@ from gradio_imageslider import ImageSlider
13
  from transformers import AutoModelForImageSegmentation
14
  from torchvision import transforms
15
 
 
 
 
 
16
  torch.set_float32_matmul_precision('high')
17
  torch.jit.script = lambda f: f
18
 
@@ -55,11 +59,11 @@ birefnet.eval()
55
 
56
 
57
  @spaces.GPU
58
- def predict(
59
- image: np.ndarray,
60
- resolution: str,
61
- weights_file: Optional[str]
62
- ) -> Tuple[np.ndarray, np.ndarray]:
63
  global birefnet
64
  # Load BiRefNet with chosen weights
65
  _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
@@ -138,6 +142,10 @@ tab_text = gr.Interface(
138
  outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
139
  examples=examples_url,
140
  api_name="text",
 
 
 
 
141
  )
142
 
143
  demo = gr.TabbedInterface(
 
13
  from transformers import AutoModelForImageSegmentation
14
  from torchvision import transforms
15
 
16
+ import requests
17
+ from io import BytesIO
18
+
19
+
20
  torch.set_float32_matmul_precision('high')
21
  torch.jit.script = lambda f: f
22
 
 
59
 
60
 
61
  @spaces.GPU
62
+ def predict(image, resolution, weights_file):
63
+ if isinstance(image, str):
64
+ response = requests.get(image)
65
+ image_data = BytesIO(response.content)
66
+ image = np.array(Image.open(image_data))
67
  global birefnet
68
  # Load BiRefNet with chosen weights
69
  _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
 
142
  outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
143
  examples=examples_url,
144
  api_name="text",
145
+ description=('Upload a URL, our model will extract a highly accurate segmentation of the subject in it.\n)'
146
+ ' The resolution used in our training was `1024x1024`, thus the suggested resolution to obtain good results!\n'
147
+ ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n'
148
+ ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.'),
149
  )
150
 
151
  demo = gr.TabbedInterface(