use mask_color variable

This commit is contained in:
deeppomf 2018-02-27 11:36:20 -05:00
parent 092c214103
commit 4e4c2d74ef

View File

@ -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)