selfie2anime processing

process.py

import os
from shutil import copyfile
import cv2


def selfie():
    base_dir = 'selfie'
    train_dir = 'trainA'
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    test_dir = 'testA'
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)
    info_file = open('selfie_dataset.txt')
    infos = [(info.split(' ')[0], info.split(' ')[1], info.split(' ')[2], info.split(' ')[3]) for info in info_file.readlines()] # filename, popularity, partial, female
    train_list = []
    for info in infos:
        if info[2] == '-1' and info[3] == '1':
            train_list.append((info[0], info[1]))
    # train_list = sorted(train_list, key = lambda x:float(x[1]), reverse=False)
    test_list = train_list[3400:3500]
    train_list = train_list[:3400]
    for file_name in train_list:
        file_name = file_name[0]+'.jpg'
        src = os.path.join(base_dir, file_name)
        dst = os.path.join(train_dir, file_name)
        copyfile(src, dst)
    for file_name in test_list:
        file_name = file_name[0]+'.jpg'
        src = os.path.join(base_dir, file_name)
        dst = os.path.join(test_dir, file_name)
        copyfile(src, dst)


def anime():
    base_dir = 'anime'
    train_dir = 'trainB'
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    test_dir = 'testB'
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)
    file_list = os.listdir(base_dir)
    infos = [(file_name, cv2.imread(os.path.join(base_dir, file_name)).size) for file_name in file_list]
    infos = sorted(infos, key=lambda x:x[1], reverse=True)[:3500]
    test_list = infos[3400:]
    train_list = infos[:3400]
    for file_name in train_list:
        file_name = file_name[0]
        src = os.path.join(base_dir, file_name)
        dst = os.path.join(train_dir, file_name)
        copyfile(src, dst)
    for file_name in test_list:
        file_name = file_name[0]
        src = os.path.join(base_dir, file_name)
        dst = os.path.join(test_dir, file_name)
        copyfile(src, dst)


if __name__ == "__main__":
    selfie()
    anime()