File size: 755 Bytes
2cb73a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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))
|