Skip to content

Instantly share code, notes, and snippets.

@heart4lor
Created Aug 10, 2019
Embed
What would you like to do?
selfie2anime processing
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment