jhj0517 commited on
Commit
60def5b
1 Parent(s): 06525bf

Add pixelize & solid color filter function

Browse files
Files changed (1) hide show
  1. modules/mask_utils.py +53 -0
modules/mask_utils.py CHANGED
@@ -6,6 +6,8 @@ from pytoshop import layers
6
  from pytoshop.enums import BlendMode
7
  from pytoshop.core import PsdFile
8
 
 
 
9
 
10
  def decode_to_mask(seg: np.ndarray[np.bool_] | np.ndarray[np.uint8]) -> np.ndarray[np.uint8]:
11
 
@@ -102,6 +104,57 @@ def create_mask_combined_images(
102
  return [enhanced, "Masked"]
103
 
104
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
105
  def insert_psd_layer(
106
  psd: PsdFile,
107
  image_data: np.ndarray,
 
6
  from pytoshop.enums import BlendMode
7
  from pytoshop.core import PsdFile
8
 
9
+ from modules.constants import DEFAULT_COLOR, DEFAULT_PIXEL_SIZE
10
+
11
 
12
  def decode_to_mask(seg: np.ndarray[np.bool_] | np.ndarray[np.uint8]) -> np.ndarray[np.uint8]:
13
 
 
104
  return [enhanced, "Masked"]
105
 
106
 
107
+ def create_mask_pixelized_image(
108
+ image: np.ndarray,
109
+ masks: List,
110
+ pixel_size: int = DEFAULT_PIXEL_SIZE
111
+ ):
112
+ final_result = image.copy()
113
+
114
+ def pixelize(img: np.ndarray, mask: np.ndarray[np.uint8], pixel_size: int):
115
+ h, w = img.shape[:2]
116
+ temp = cv2.resize(img, (w // pixel_size, h // pixel_size), interpolation=cv2.INTER_LINEAR)
117
+
118
+ pixelated = cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
119
+
120
+ return np.where(mask[:, :, np.newaxis] > 0, pixelated, img)
121
+
122
+ for info in masks:
123
+ rle = info['segmentation']
124
+ mask = decode_to_mask(rle)
125
+
126
+ pixelated_segment = pixelize(final_result, mask, pixel_size)
127
+
128
+ final_result = np.where(mask[:, :, np.newaxis] > 0, pixelated_segment, final_result)
129
+
130
+ return final_result
131
+
132
+
133
+ def create_solid_color_mask_image(
134
+ image: np.ndarray,
135
+ masks: List,
136
+ color_hex: str = DEFAULT_COLOR
137
+ ):
138
+ final_result = image.copy()
139
+
140
+ def hex_to_bgr(hex_color: str):
141
+ hex_color = hex_color.lstrip('#')
142
+ rgb = tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))
143
+ return rgb[::-1]
144
+
145
+ color_bgr = hex_to_bgr(color_hex)
146
+
147
+ for info in masks:
148
+ rle = info['segmentation']
149
+ mask = decode_to_mask(rle)
150
+
151
+ solid_color_mask = np.full(image.shape, color_bgr, dtype=np.uint8)
152
+
153
+ final_result = np.where(mask[:, :, np.newaxis] > 0, solid_color_mask, final_result)
154
+
155
+ return final_result
156
+
157
+
158
  def insert_psd_layer(
159
  psd: PsdFile,
160
  image_data: np.ndarray,