import argparse
import os

import numpy as np

import cv2
import glob


parser = argparse.ArgumentParser()
parser.add_argument('input', type=str, default='', help='path to input directory')
parser.add_argument('output', type=str, default='', help='path to output directory')

args = parser.parse_args()


target = args.input
out_dir = args.output
targets = glob.glob(os.path.join(target+'*.png'))

for target in targets:
    img = cv2.imread(target)
    r = img[:,:,2] == 255
    g = img[:,:,1] == 255
    b = img[:,:,1] == 255
    mask1 = np.logical_and(r, g, b)
    mask2 = np.logical_and(r, np.logical_not(g),np.logical_not(b))
    out = np.zeros(np.shape(mask1))
    out[mask1] = 1
    out[mask2] = 2
    cv2.imwrite(os.path.join(out_dir,os.path.basename(target)), out)

# for x in glob.glob(os.pathout_dir+'/*.png'):
#     test = cv2.imread(x,0)
#     print(test.shape)
#     pylab.imshow(test)
#     pylab.show()