39 lines
920 B
Python
39 lines
920 B
Python
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()
|
|
|