jhj0517 commited on
Commit
dfc0c1a
1 Parent(s): 8d52a7d

Fix `create_mask_combined_images()`

Browse files
Files changed (1) hide show
  1. modules/mask_utils.py +23 -8
modules/mask_utils.py CHANGED
@@ -1,9 +1,8 @@
1
  import cv2
2
  import numpy as np
3
  from typing import Dict, List
4
- from pycocotools import mask as coco_mask
5
  from pytoshop import layers
6
- import pytoshop
7
  from pytoshop.enums import BlendMode
8
  from pytoshop.core import PsdFile
9
 
@@ -13,7 +12,11 @@ def decode_to_mask(seg: np.ndarray[np.bool_]) -> np.ndarray[np.uint8]:
13
 
14
 
15
  def generate_random_color():
16
- return np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)
 
 
 
 
17
 
18
 
19
  def create_base_layer(image: np.ndarray):
@@ -68,19 +71,31 @@ def create_mask_combined_images(
68
  masks: List
69
  ):
70
  final_result = np.zeros_like(image)
 
71
 
72
  for info in masks:
73
  rle = info['segmentation']
74
  mask = decode_to_mask(rle)
75
 
76
- color = generate_random_color()
 
 
 
 
 
77
  colored_mask = np.zeros_like(image)
78
- colored_mask[mask == 1] = color
 
 
 
 
 
79
 
80
- final_result = cv2.addWeighted(final_result, 1, colored_mask, 0.5, 0)
 
 
81
 
82
- combined_image = cv2.addWeighted(image, 1, final_result, 0.5, 0)
83
- return [combined_image, "masked"]
84
 
85
 
86
  def insert_psd_layer(
 
1
  import cv2
2
  import numpy as np
3
  from typing import Dict, List
4
+ import colorsys
5
  from pytoshop import layers
 
6
  from pytoshop.enums import BlendMode
7
  from pytoshop.core import PsdFile
8
 
 
12
 
13
 
14
  def generate_random_color():
15
+ h = np.random.randint(0, 360)
16
+ s = np.random.randint(70, 100) / 100
17
+ v = np.random.randint(70, 100) / 100
18
+ r, g, b = colorsys.hsv_to_rgb(h/360, s, v)
19
+ return int(r * 255), int(g * 255), int(b * 255)
20
 
21
 
22
  def create_base_layer(image: np.ndarray):
 
71
  masks: List
72
  ):
73
  final_result = np.zeros_like(image)
74
+ used_colors = set()
75
 
76
  for info in masks:
77
  rle = info['segmentation']
78
  mask = decode_to_mask(rle)
79
 
80
+ while True:
81
+ color = generate_random_color()
82
+ if color not in used_colors:
83
+ used_colors.add(color)
84
+ break
85
+
86
  colored_mask = np.zeros_like(image)
87
+ colored_mask[mask > 0] = color
88
+
89
+ blended = cv2.addWeighted(image, 0.3, colored_mask, 0.7, 0)
90
+ final_result = np.where(mask[:, :, np.newaxis] > 0, blended, final_result)
91
+
92
+ combined_image = np.where(final_result != 0, final_result, image)
93
 
94
+ hsv = cv2.cvtColor(combined_image, cv2.COLOR_BGR2HSV)
95
+ hsv[:, :, 1] = np.clip(hsv[:, :, 1] * 1.5, 0, 255)
96
+ enhanced = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
97
 
98
+ return [enhanced, "Masked"]
 
99
 
100
 
101
  def insert_psd_layer(