Spaces:
Paused
Paused
File size: 1,857 Bytes
3de0e37 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
# -*- coding: utf-8 -*-
"""
# File name: script_get_sean_code.py
# Time : 2021/11/16 15:56
# Author: xyguoo@163.com
# Description:
"""
import os
import sys
sys.path.append('.')
from dataset_scripts.utils import merge_pickle_dir_to_dict
import cv2
import tqdm
from global_value_utils import GLOBAL_DATA_ROOT, DATASET_NAME
from hair_editor import HairEditor
import pickle
data_name = DATASET_NAME
root_dir = GLOBAL_DATA_ROOT
imgs_sub_dir = 'images_256'
target_dir = os.path.join(root_dir, 'hair_info_all_dataset/sean_code')
he = HairEditor(load_mask_model=False)
path_list = []
for d in data_name:
data_dir = os.path.join(root_dir, d, imgs_sub_dir)
path_list += [os.path.join(data_dir, pp) for pp in os.listdir(data_dir)]
path_list.sort()
# res_dict = {}
if not os.path.exists(target_dir):
os.makedirs(target_dir)
for img_path in tqdm.tqdm(path_list):
for dd in data_name:
if img_path.find(dd) != -1:
dataset_name = dd
break
else:
raise NotImplementedError
base_name = os.path.basename(img_path)
hair_path = os.path.join(root_dir, dataset_name, imgs_sub_dir, base_name)
hair_img = cv2.cvtColor(cv2.imread(hair_path), cv2.COLOR_BGR2RGB)
hair_parsing = cv2.imread(os.path.join(root_dir, dataset_name, 'label', base_name), cv2.IMREAD_GRAYSCALE)
# resize
hair_img = he.preprocess_img(hair_img)
hair_parsing = he.preprocess_mask(hair_parsing)
cur_code = he.get_code(hair_img, hair_parsing)
cur_code = cur_code.cpu().numpy()[0]
# res_dict['%s___%s' % (dataset_name, base_name)] = cur_code
target_file = os.path.join(target_dir, '%s___%s.pkl' % (dataset_name, base_name[:-4]))
with open(target_file, 'wb') as f:
pickle.dump(cur_code, f)
merge_pickle_dir_to_dict(target_dir, os.path.join(root_dir, 'sean_code_dict.pkl'))
|