make black/blonde dataset for cyclegan

make_dataset.py

import os
from PIL import Image

# Attribute
f = open('./list_attr_celeba.txt', 'r')
text = f.readlines()

# Number of images
text[0] 

# Attributes
text[1]

attr2idx = dict()
for i, attr in enumerate(text[1].split()):
    attr2idx[attr] = i
    
attr2idx['Black_Hair'] # 8
attr2idx['Blond_Hair'] # 9


image_path = '/usr/share/CelebA/CelebA_nocrop/images/'
images = os.listdir(image_path)
img = Image.open(image_path + '000007.jpg')
img


splits = text[8].split()
img_file = splits[0]
attrs = splits[1:]

black_idx = attr2idx['Black_Hair']  # 8
blonde_idx = attr2idx['Blond_Hair']  # 9

black_path = './CelebA_hair/black/'
blonde_path = './CelebA_hair/blonde/'

if not os.path.exists(black_path):
    os.makedirs(black_path)
if not os.path.exists(blonde_path):
    os.makedirs(blonde_path)
   
def save_if_black_or_blonde(img_file, attrs):
    if attrs[black_idx] == '1':
        # Save black hair image
        img = Image.open(image_path + img_file)
        img = img.resize([256, 256], Image.ANTIALIAS)
        img.save(black_path + img_file)

    elif attrs[blonde_idx] == '1':
        # Save black hair image
        img = Image.open(image_path + img_file)
        img = img.resize([256, 256], Image.ANTIALIAS)
        img.save(blonde_path + img_file)

for i, face_image in enumerate(text[2:]):
    splits = face_image.split()
    img_file = splits[0]
    attrs = splits[1:]
    save_if_black_or_blonde(img_file, attrs)
    
    if i % 100 == 0:
        print (i)

        
# Train/Test split
black_path = './CelebA_hair/trainA/'
blonde_path = './CelebA_hair/trainB/'
test_black_path = './CelebA_hair/testA/'
test_blonde_path = './CelebA_hair/testB/'

if not os.path.exists(test_black_path):
    os.makedirs(test_black_path)
if not os.path.exists(test_blonde_path):
    os.makedirs(test_blonde_path)
    
black_images = os.listdir(black_path)
blonde_images = os.listdir(blonde_path)

for img in black_images[0:1000]:
    os.rename(black_path + img, test_black_path + img)
for img in blonde_images[0:1000]:
    os.rename(blonde_path + img, test_blonde_path + img)