switchablenorms / CelebAMask-HQ

A large-scale face dataset for face parsing, recognition, generation and editing.
2.05k stars 343 forks source link

Coderefactoring for g_mask.py #16

Closed bdseal closed 4 years ago

bdseal commented 4 years ago

Coderefactoring for g_mask.py

`face_sep_mask = 'CelebAMask-HQ-mask-anno'

mask_path = 'CelebAMask-HQ/mask'

for i in range(15):

    # files = os.listdir(osp.join(face_sep_mask, str(i)))

    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)`
jianwang-scu commented 4 years ago

There was nothing wrong with the original version, except that the '/ ' had to be changed to '/ / ' in lines 18-21 and 23.

steven413d commented 4 years ago

Hi, thanks for providing a faster and terse version. I have made a new version of g_mask.py