haydpw commited on
Commit
f0aee21
·
1 Parent(s): cf104bc

fix process image

Browse files
Files changed (1) hide show
  1. utils/helpers.py +15 -6
utils/helpers.py CHANGED
@@ -15,20 +15,29 @@ def calculate_mask_area(mask: Image.Image, background=False) -> int:
15
  return non_zero_pixels
16
 
17
  def process_image(input_image: Image.Image) -> Image.Image:
18
- if input_image.mode != 'RGBA':
19
- input_image = input_image.convert('RGBA')
20
 
21
  data = np.array(input_image)
22
  # Split the image into its component channels
23
- r, g, b, a = data.T
24
 
25
  # Create a mask where all pixels that are black (0,0,0) will have 0 alpha
26
- black_areas = (r == 0) & (g == 0) & (b == 0)
 
 
 
 
 
 
 
 
 
 
 
27
 
28
  # Apply the mask to the alpha channel
29
- data[..., 3][black_areas] = 0
30
 
31
- return Image.fromarray(data)
32
 
33
  def combine_images(original_image: Image.Image, masks: list) -> Image.Image:
34
  combined = original_image.copy()
 
15
  return non_zero_pixels
16
 
17
  def process_image(input_image: Image.Image) -> Image.Image:
18
+
 
19
 
20
  data = np.array(input_image)
21
  # Split the image into its component channels
 
22
 
23
  # Create a mask where all pixels that are black (0,0,0) will have 0 alpha
24
+ black_areas = data == 0
25
+
26
+ rgba_image = Image.new('RGBA', input_image.size)
27
+ rgba_data = np.array(rgba_image)
28
+
29
+ # Copy the grayscale data to all RGB channels
30
+ rgba_data[..., 0] = data
31
+ rgba_data[..., 1] = data
32
+ rgba_data[..., 2] = data
33
+
34
+ # Set alpha channel to 255 (fully opaque) for all pixels
35
+ rgba_data[..., 3] = 255
36
 
37
  # Apply the mask to the alpha channel
38
+ rgba_data[..., 3][black_areas] = 0
39
 
40
+ return Image.fromarray(rgba_data)
41
 
42
  def combine_images(original_image: Image.Image, masks: list) -> Image.Image:
43
  combined = original_image.copy()