File size: 5,610 Bytes
ef6a8c6 |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
'''
@Author: Jiapeng Zhou
@Email: [email protected]
@Desc: script to process all data of RECMV.
Before run this bash, what you only need to do is :
- run openpose to get keypoints.json file in the dir: (self.root, openpose)
- have a conda environment of videoavatar, named: "videoavatar"
Usage: python process_data_all.py --root ~/DeepDxxx/xinyu_a/ --gid 0 --gender female
'''
import pdb, os, os.path as osp, sys, argparse
from typing import Any
import numpy as np, cv2
class RECMVDataProcessor():
def __init__(self, args) -> None:
self.args = args
self.gender = args.gender
self.root = args.root
#out = osp.join(self.root, args.out)
#os.makedirs(out, exist_ok=True)
self.img_dir = osp.join(self.root, 'imgs')
assert osp.exists(self.img_dir) and len(os.listdir(self.img_dir)) > 0, 'img file does not exist'
self.gid = args.gid
def __len__(self):
return len(os.listdir(self.img_dir))
def __call__(self):
'''get normals, ATR-parsing, masks, camera, smpl_params, depth, featurelines'''
# get ATR_parsing
self.get_ATR_parsing()
# get masks
self.get_masks()
# get normal from pifuhd
self.get_normals()
'''cannot directly ca videoavatar
# get reconstructed_poses.hdf5 from videoavatar
self.get_poses()
# get camera.npz and smpl_rec.npz
self.get_camera_smpl_rec()
'''
# get depth
#self.get_depth()
# get featurelines
#self.get_featurelines()
def get_featurelines(self):
fl_dir = osp.join(self.root, 'featurelines')
if osp.exists(fl_dir) and len(os.listdir(fl_dir)) >= self.__len__():
print('featurelines files exists')
else:
print('featurelines file does not exist')
def get_depth(self):
depth_dir = osp.join(self.root, 'depth')
if osp.exists(depth_dir) and len(os.listdir(depth_dir)) >= self.__len__():
print('depth files exists')
else:
print('depth file does not exist')
def get_masks(self):
mask_dir = osp.join(self.root, 'masks')
if osp.exists(mask_dir) and len(os.listdir(mask_dir)) >= self.__len__():
print('mask files exists')
return
os.makedirs(mask_dir, exist_ok=True)
# mask是否需要按照单通道保存？
print('\n-----------------Getting Masks-------------------------')
os.system(f'cd ./RobustVideoMatting && python inference_itw_rotate.py \
--input_pth {self.img_dir} --output_pth {mask_dir} --device {self.gid}')
def get_normals(self):
normal_dir = osp.join(self.root, 'normals')
if osp.exists(normal_dir) and len(os.listdir(normal_dir)) >= self.__len__():
print('normal files exists')
return
print('\n-------------------Getting Normals----------------------')
os.makedirs(normal_dir, exist_ok=True)
os.system('cd ./pifuhd && python ./generate_normals.py --gid {} --imgpath {}'.format(self.gid, self.img_dir))
def get_ATR_parsing(self):
input_dir = osp.join(self.root, 'imgs')
output_dir = osp.join(self.root, 'parsing_SCH_ATR')
if osp.exists(output_dir) and len(os.listdir(output_dir)) >= self.__len__():
print('ATR parsing files exists')
return
print('\n--------------------Getting ATR_Parsing------------------')
os.makedirs(output_dir, exist_ok=True)
os.system(f'cd ./Self-Correction-Human-Parsing && python simple_extractor.py \
--dataset atr --model-restore exp-schp-201908301523-atr.pth --gpu {self.gid} \
--input-dir {self.img_dir} --output-dir {output_dir} --logits')
def get_poses(self):
openpose_p = osp.join(self.root, 'openpose')
assert osp.exists(openpose_p) and len(os.listdir(openpose_p)) >= self.__len__(), 'openpose keypoint folder doesnot exist or it is emtpy'
output_dir = osp.join(self.root, 'videoavatars')
if osp.exists(output_dir) and osp.exists(osp.join(output_dir, 'reconstructed_poses.hdf5')):
print('reconstructed_poses.hdf5 files exists')
return
print('\n------------------Getting Poses---------------------------')
os.makedirs(output_dir, exist_ok=True)
# Note: Here needs to conda actiavte the environment of videoavatar
os.system(f'cd ./videoavatars && bash get_reconstructed_poses.sh {self.root} {output_dir} {self.gender}')
def get_camera_smpl_rec(self):
if osp.exists(osp.join(self.root, 'camera.npz')) and osp.exists(osp.join(self.root, 'smpl_rec.npz')):
print('camera.npz and smpl_rec.npz files exists')
return
print('\n---------------------Getting camera and smpl_rec--------------------')
os.system(' python get_smpl_rec_camera.py --root {} --save_root {} --gender {}'.\
format(osp.join(self.root, 'videoavatars'), self.root, self.gender))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--gid', type=int, default=0)
parser.add_argument('--root', type=str, required=True)
parser.add_argument('--gender', type=str, choices=['male', 'female'], required=True)
args = parser.parse_args()
return args
def main():
args = parse_args()
processor = RECMVDataProcessor(args)
processor()
if __name__ == '__main__':
main()
|