Spaces:
mashroo
/
Running on Zero

YoussefAnso commited on
Commit
882d814
·
1 Parent(s): 05a5508

Enhance image preprocessing in app.py with background removal and color conversion

Browse files

- Added a new function to convert hex color strings to RGB tuples.
- Updated the preprocess_image function to handle RGBA images, perform background removal, resize content, and apply a solid background color.
- Improved error handling for background removal failures, ensuring user feedback on model availability.

Files changed (1) hide show
  1. app.py +24 -6
app.py CHANGED
@@ -98,19 +98,37 @@ def add_background(image, bg_color=(255, 255, 255)):
98
  return Image.alpha_composite(background, image)
99
 
100
 
 
 
 
 
 
 
101
  def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
102
  """
103
- input image is a pil image in RGBA, return RGB image
 
 
104
  """
105
- print(background_choice)
106
- if background_choice == "Alpha as mask":
107
- background = Image.new("RGBA", image.size, (0, 0, 0, 0))
108
- image = Image.alpha_composite(background, image)
109
- else:
110
  image = remove_background(image)
111
  if image is None:
112
  raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
113
 
 
 
 
 
 
 
 
 
 
 
 
114
  @spaces.GPU
115
  def gen_image(input_image, seed, scale, step):
116
  global pipeline, model, args
 
98
  return Image.alpha_composite(background, image)
99
 
100
 
101
+ def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
102
+ """Converts a hex color string to an RGB tuple."""
103
+ hex_color = hex_color.lstrip('#')
104
+ return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
105
+
106
+
107
  def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
108
  """
109
+ Preprocesses the input image by optionally removing the background, resizing,
110
+ and adding a new solid background.
111
+ Returns an RGB PIL Image.
112
  """
113
+ if image.mode != 'RGBA':
114
+ image = image.convert('RGBA')
115
+
116
+ if background_choice == "Auto Remove background":
 
117
  image = remove_background(image)
118
  if image is None:
119
  raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
120
 
121
+ # Resize the content of the image
122
+ image = do_resize_content(image, foreground_ratio)
123
+
124
+ # Add a solid background color
125
+ rgb_background = hex_to_rgb(backgroud_color)
126
+ image_with_bg = add_background(image, rgb_background)
127
+
128
+ # Convert to RGB and return
129
+ return image_with_bg.convert("RGB")
130
+
131
+
132
  @spaces.GPU
133
  def gen_image(input_image, seed, scale, step):
134
  global pipeline, model, args