jhj0517 commited on
Commit
3810a85
1 Parent(s): 16bf670

Update decode

Browse files
Files changed (1) hide show
  1. modules/mask_utils.py +6 -2
modules/mask_utils.py CHANGED
@@ -7,8 +7,12 @@ 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]:
11
- return seg.astype(np.uint8) * 255
 
 
 
 
12
 
13
 
14
  def generate_random_color():
 
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
+
12
+ if isinstance(seg, np.ndarray) and seg.dtype == np.bool_:
13
+ return seg.astype(np.uint8) * 255
14
+ else:
15
+ return seg.astype(np.uint8)
16
 
17
 
18
  def generate_random_color():