not-lain commited on
Commit
13a0890
·
verified ·
1 Parent(s): 6438ac6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -4,6 +4,7 @@ import spaces
4
  from transformers import AutoModelForImageSegmentation
5
  import torch
6
  from torchvision import transforms
 
7
 
8
  torch.set_float32_matmul_precision(["high", "highest"][0])
9
 
@@ -20,7 +21,7 @@ transform_image = transforms.Compose(
20
  ]
21
  )
22
 
23
- def fn(image):
24
  """
25
  Remove the background from an image and return both the transparent version and the original.
26
 
@@ -43,7 +44,7 @@ def fn(image):
43
  return (processed_image, origin)
44
 
45
  @spaces.GPU
46
- def process(image):
47
  """
48
  Apply BiRefNet-based image segmentation to remove the background.
49
 
@@ -67,7 +68,7 @@ def process(image):
67
  image.putalpha(mask)
68
  return image
69
 
70
- def process_file(f):
71
  """
72
  Load an image file from disk, remove the background, and save the output as a transparent PNG.
73
 
 
4
  from transformers import AutoModelForImageSegmentation
5
  import torch
6
  from torchvision import transforms
7
+ from typing import Union, Tuple
8
 
9
  torch.set_float32_matmul_precision(["high", "highest"][0])
10
 
 
21
  ]
22
  )
23
 
24
+ def fn(image: Union[PILImage.Image, str]) -> Tuple[PILImage.Image, PILImage.Image]:
25
  """
26
  Remove the background from an image and return both the transparent version and the original.
27
 
 
44
  return (processed_image, origin)
45
 
46
  @spaces.GPU
47
+ def process(image: PILImage.Image) -> PILImage.Image:
48
  """
49
  Apply BiRefNet-based image segmentation to remove the background.
50
 
 
68
  image.putalpha(mask)
69
  return image
70
 
71
+ def process_file(f: str) -> str:
72
  """
73
  Load an image file from disk, remove the background, and save the output as a transparent PNG.
74