mirror of
https://github.com/Deepshift/DeepCreamPy.git
synced 2025-01-05 21:55:57 +00:00
use mask_color variable
This commit is contained in:
parent
092c214103
commit
4e4c2d74ef
@ -74,7 +74,7 @@ def get_mask(x_batch):
|
||||
m = np.zeros((args.image_size, args.image_size, 1), dtype=np.uint8)
|
||||
for x in range(args.image_size):
|
||||
for y in range(args.image_size):
|
||||
if np.array_equal(raw[x][y], [0, 255, 0]):
|
||||
if np.array_equal(raw[x][y], mask_color):
|
||||
m[x, y] = 1
|
||||
mask.append(m)
|
||||
return np.array(mask)
|
||||
|
Loading…
Reference in New Issue
Block a user