mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
#!/usr/bin/python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import os.path as osp
|
|
import os
|
|
import cv2
|
|
from transform import *
|
|
from PIL import Image
|
|
|
|
face_data = '/home/zll/data/CelebAMask-HQ/CelebA-HQ-img'
|
|
face_sep_mask = '/home/zll/data/CelebAMask-HQ/CelebAMask-HQ-mask-anno'
|
|
mask_path = '/home/zll/data/CelebAMask-HQ/mask'
|
|
counter = 0
|
|
total = 0
|
|
for i in range(15):
|
|
|
|
atts = ['skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r',
|
|
'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat']
|
|
|
|
for j in range(i * 2000, (i + 1) * 2000):
|
|
|
|
mask = np.zeros((512, 512))
|
|
|
|
for l, att in enumerate(atts, 1):
|
|
total += 1
|
|
file_name = ''.join([str(j).rjust(5, '0'), '_', att, '.png'])
|
|
path = osp.join(face_sep_mask, str(i), file_name)
|
|
|
|
if os.path.exists(path):
|
|
counter += 1
|
|
sep_mask = np.array(Image.open(path).convert('P'))
|
|
# print(np.unique(sep_mask))
|
|
|
|
mask[sep_mask == 225] = l
|
|
cv2.imwrite('{}/{}.png'.format(mask_path, j), mask)
|
|
print(j)
|
|
|
|
print(counter, total) |