dmitriitochilkin commited on
Commit
23e4ec1
1 Parent(s): 78efa10

support rgba pngs

Browse files
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -35,7 +35,7 @@ def check_input_image(input_image):
35
  raise gr.Error("No image uploaded!")
36
 
37
 
38
- def preprocess(image_path, do_remove_background, foreground_ratio):
39
  def fill_background(image):
40
  image = np.array(image).astype(np.float32) / 255.0
41
  image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
@@ -43,11 +43,12 @@ def preprocess(image_path, do_remove_background, foreground_ratio):
43
  return image
44
 
45
  if do_remove_background:
46
- image = remove_background(Image.open(image_path), rembg_session)
 
47
  image = resize_foreground(image, foreground_ratio)
48
  image = fill_background(image)
49
  else:
50
- image = Image.open(image_path)
51
  if image.mode == "RGBA":
52
  image = fill_background(image)
53
  return image
@@ -74,8 +75,9 @@ with gr.Blocks() as demo:
74
  with gr.Row():
75
  input_image = gr.Image(
76
  label="Input Image",
 
77
  sources="upload",
78
- type="filepath",
79
  elem_id="content_image",
80
  )
81
  processed_image = gr.Image(label="Processed Image", interactive=False)
 
35
  raise gr.Error("No image uploaded!")
36
 
37
 
38
+ def preprocess(input_image, do_remove_background, foreground_ratio):
39
  def fill_background(image):
40
  image = np.array(image).astype(np.float32) / 255.0
41
  image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
 
43
  return image
44
 
45
  if do_remove_background:
46
+ image = input_image.convert("RGB")
47
+ image = remove_background(image, rembg_session)
48
  image = resize_foreground(image, foreground_ratio)
49
  image = fill_background(image)
50
  else:
51
+ image = input_image
52
  if image.mode == "RGBA":
53
  image = fill_background(image)
54
  return image
 
75
  with gr.Row():
76
  input_image = gr.Image(
77
  label="Input Image",
78
+ image_mode="RGBA",
79
  sources="upload",
80
+ type="pil",
81
  elem_id="content_image",
82
  )
83
  processed_image = gr.Image(label="Processed Image", interactive=False)