| from PIL import Image | |
| import cv2 | |
| import os | |
| import numpy as np | |
| i = 0 | |
| threshold = 240 | |
| m = 0 | |
| for img in os.listdir("orig_masks"): | |
| mask = cv2.imread("orig_masks/" + img) | |
| for d in range(3): | |
| mask[:, :, d][mask[:, :, d] < threshold] = 0 | |
| mask[:, :, d][mask[:, :, d] > 0] = 1 | |
| mask = mask[:, :, 0] | mask[:, :, 1] | mask[:, :, 2] | |
| mask = (mask * 255).astype(np.uint8) | |
| #nonzero = [] | |
| #for i in mask[np.nonzero(mask)]: | |
| # nonzero.append(i) | |
| if np.all(mask == 0): | |
| print("The mask contains all zeros.") | |
| continue | |
| cv2.imwrite("masks/" + img, mask) | |
| img_split = img.split("_") | |
| os.system("cp orig_imgs/" + img_split[0] + "_" + img_split[2] + " imgs/") | |
| i += 1 | |
| print("Done ", str(i), str(m)) | |