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