def __getitem__(self, index):
item_info = self.train_set_info[index]
img = Image.open(
os.path.join(self.data_root, 'image',
item_info['image'])).convert('RGB')
mask = Image.open(
os.path.join(self.data_root, 'mask', item_info['mask']))
if self.transform is not None:
img = self.transform(img)
mask = self.transform(mask)
vns_dict = load_pickle_file(
os.path.join(self.data_root, 'verts_normals',
item_info['model'] + '.pkl'))
verts = vns_dict['verts']
normals = vns_dict['normals']
camera_dict = load_pickle_file(
os.path.join(self.data_root, 'cameras',
item_info['model'] + '.pkl'))
K = camera_dict['K']
output = {
'image': img,
'mask': mask,
'vertices': torch.from_numpy(verts).float(),
'normals': torch.from_numpy(normals).float(),
'translation':
torch.tensor(item_info['pose']['translation']).float(),
'rotation': torch.tensor(item_info['pose']['rotation']).float(),
'K': torch.from_numpy(K).float()
}
return output
def _read_vids_info(self, file_path):
vids_info = []
with open(file_path, 'r') as reader:
lines = []
for line in reader:
line = line.rstrip()
lines.append(line)
total = len(lines)
for i, line in enumerate(lines):
images_path = glob.glob(
os.path.join(self._vids_dir, line, '*.jpg'))
images_path.sort()
smpl_data = load_pickle_file(
os.path.join(self._smpls_dir, line, 'pose_shape.pkl'))
kps_data = load_pickle_file(
os.path.join(self._smpls_dir, line, 'kps.pkl'))
cams = smpl_data['cams']
kps = (kps_data['kps'] + 1) / 2.0 * 1024
assert len(images_path) == len(cams), '{} != {}'.format(
len(images_path), len(cams))
info = {
'images': images_path,
'cams': cams,
'thetas': smpl_data['pose'],
'betas': smpl_data['shape'],
'j2ds': kps,
'length': len(images_path)
}
vids_info.append(info)
self._dataset_size += info['length'] // self._intervals
# self._dataset_size += info['length']
self._num_videos += 1
print('loading video = {}, {} / {}'.format(line, i, total))
if self._opt.debug:
if i > 1:
break
return vids_info
def _read_pair_list(self, im_dir, pair_pkl_path):
pair_list = load_pickle_file(pair_pkl_path)
new_pair_list = []
for i, pairs in enumerate(pair_list):
src_path = os.path.join(im_dir, pairs[0])
dst_path = os.path.join(im_dir, pairs[1])
new_pair_list.append((src_path, dst_path))
return new_pair_list
def __init__(self, pkl_path, rotate=False):
"""
pkl_path is the path to a SMPL model
"""
super(SMPL, self).__init__()
self.rotate = rotate
# -- Load SMPL params --
dd = load_pickle_file(pkl_path)
# define faces
# self.register_buffer('faces', torch.from_numpy(undo_chumpy(dd['f']).astype(np.int32)).type(dtype=torch.int32))
self.faces = torch.from_numpy(dd['f'].astype(
np.int32)).type(dtype=torch.int32)
# Mean template vertices
self.register_buffer('v_template', torch.FloatTensor(dd['v_template']))
# Size of mesh [Number of vertices, 3], (6890, 3)
self.size = [self.v_template.shape[0], 3]
self.num_betas = dd['shapedirs'].shape[-1]
# Shape blend shape basis (shapedirs): (6980, 3, 10)
# reshaped to (6980*3, 10), transposed to (10, 6980*3)
self.register_buffer(
'shapedirs',
torch.FloatTensor(
np.reshape(dd['shapedirs'], [-1, self.num_betas]).T))
# Regressor for joint locations given shape -> (24, 6890)
# Transpose to shape (6890, 24)
self.register_buffer(
'J_regressor',
torch.FloatTensor(np.asarray(dd['J_regressor'].T.todense())))
# Pose blend shape basis: (6890, 3, 207)
num_pose_basis = dd['posedirs'].shape[-1]
# Pose blend pose basis is reshaped to (6890*3, 207)
# posedirs is transposed into (207, 6890*3)
self.register_buffer(
'posedirs',
torch.FloatTensor(
np.reshape(dd['posedirs'], [-1, num_pose_basis]).T))
# indices of parents for each joints
self.parents = np.array(dd['kintree_table'][0].astype(np.int32))
# LBS weights (6890, 24)
self.register_buffer('weights', torch.FloatTensor(dd['weights']))
# This returns 19 keypoints: 6890 x 19
joint_regressor = torch.FloatTensor(
np.asarray(dd['cocoplus_regressor'].T.todense()))
self.register_buffer('joint_regressor', joint_regressor)
def __init__(self, imgs_path, pose_shape_pkl_path, image_size=256):
imgs_fn_list = os.listdir(imgs_path)
imgs_fn_list.sort()
self.imgs_path_list = [
os.path.join(imgs_path, img_fn) for img_fn in imgs_fn_list
]
self.pose_shape_pkl = load_pickle_file(pose_shape_pkl_path)
self.Resize = transforms.Resize(image_size)
self.ToTensor = transforms.ToTensor()
if len(self.imgs_path_list) != self.pose_shape_pkl['pose'].shape[0]:
print('images: ', len(self.imgs_path_list))
print('smpls: ', self.pose_shape_pkl['pose'].shape[0])
def load_senet_model(self, pretrain_model):
# saved_data = torch.load(pretrain_model, encoding='latin1')
from utils.util import load_pickle_file
saved_data = load_pickle_file(pretrain_model)
save_weights_dict = dict()
for key, val in saved_data.items():
if key.startswith('fc'):
continue
save_weights_dict[key] = torch.from_numpy(val)
self.net.load_state_dict(save_weights_dict)
print('load face model from {}'.format(pretrain_model))
def load_sample(self, im_pairs, pkl_path):
# 1. load images
imgs = self.load_images(im_pairs)
# 2.load pickle data
pkl_data = load_pickle_file(pkl_path)
src_fim = pkl_data['from_face_index_map'][:, :,
0] # (img_size, img_size)
dst_fim = pkl_data['to_face_index_map'][:, :,
0] # (img_size, img_size)
T = pkl_data['T'] # (img_size, img_size, 2)
fims = np.stack([src_fim, dst_fim], axis=0)
fims_enc = self.map_fn[fims] # (2, h, w, c)
fims_enc = np.transpose(fims_enc, axes=(0, 3, 1, 2)) # (2, c, h, w)
sample = {
'images': torch.tensor(imgs).float(),
'src_fim': torch.tensor(src_fim).float(),
'tsf_fim': torch.tensor(dst_fim).float(),
'fims': torch.tensor(fims_enc).float(),
'T': torch.tensor(T).float(),
'j2d': torch.tensor(pkl_data['j2d']).float()
}
if 'warp' in pkl_data:
if len(pkl_data['warp'].shape) == 4:
sample['warp'] = torch.tensor(pkl_data['warp'][0],
dtype=torch.float32)
else:
sample['warp'] = torch.tensor(pkl_data['warp'],
dtype=torch.float32)
elif 'warp_R' in pkl_data:
sample['warp'] = torch.tensor(pkl_data['warp_R'][0],
dtype=torch.float32)
elif 'warp_T' in pkl_data:
sample['warp'] = torch.tensor(pkl_data['warp_T'][0],
dtype=torch.float32)
if 'T_cycle' in pkl_data:
sample['T_cycle'] = torch.tensor(pkl_data['T_cycle']).float()
if 'T_cycle_vis' in pkl_data:
sample['T_cycle_vis'] = torch.tensor(
pkl_data['T_cycle_vis']).float()
return sample
def load_sample(self, im_pairs, pkl_path):
# 1. load images
imgs = self.load_images(im_pairs)
# 2.load pickle data
pkl_data = load_pickle_file(pkl_path)
src_fim = pkl_data['from_face_index_map'][:, :,
0] # (img_size, img_size)
dst_fim = pkl_data['to_face_index_map'][:, :,
0] # (img_size, img_size)
T = pkl_data['T'] # (img_size, img_size, 2)
fims = np.stack([src_fim, dst_fim], axis=0)
fims_enc = self.map_fn[fims] # (2, h, w, c)
fims_enc = np.transpose(fims_enc, axes=(0, 3, 1, 2)) # (2, c, h, w)
heads_mask = self.head_fn[fims[1:]] # (1, h, w, 1)
heads_mask = np.transpose(heads_mask,
axes=(0, 3, 1, 2)) # (1, 1, h, w)
head_bbox, _ = cal_mask_bbox(heads_mask, factor=1.05)
body_bbox, _ = cal_mask_bbox(1 - fims_enc[1:, -1:], factor=1.2)
# print(head_bbox.shape, valid_bbox.shape)
sample = {
'images': torch.tensor(imgs).float(),
'fims': torch.tensor(fims_enc).float(),
'T': torch.tensor(T).float(),
'head_bbox': torch.tensor(head_bbox[0]).long(),
'body_bbox': torch.tensor(body_bbox[0]).long()
}
if 'warp' in pkl_data:
if len(pkl_data['warp'].shape) == 4:
sample['warp'] = torch.tensor(pkl_data['warp'][0],
dtype=torch.float32)
else:
sample['warp'] = torch.tensor(pkl_data['warp'],
dtype=torch.float32)
elif 'warp_R' in pkl_data:
sample['warp'] = torch.tensor(pkl_data['warp_R'][0],
dtype=torch.float32)
elif 'warp_T' in pkl_data:
sample['warp'] = torch.tensor(pkl_data['warp_T'][0],
dtype=torch.float32)
return sample
def _read_samples_info(self, im_dir, pkl_dir, pair_ids_filepath):
"""
Args:
im_dir:
pkl_dir:
pair_ids_filepath:
Returns:
"""
# 1. load image pair list
self.im_pair_list = load_pickle_file(pair_ids_filepath)
# 2. load pkl file paths
self.all_pkl_paths = sorted(glob.glob((os.path.join(pkl_dir, '*.pkl'))))
assert len(self.im_pair_list) == len(self.all_pkl_paths), '{} != {}'.format(
len(self.im_pair_list), len(self.all_pkl_paths)
)
self._dataset_size = len(self.im_pair_list)
def _read_samples_info(self, im_dir, pkl_dir, pair_ids_filepath):
"""
Args:
im_dir:
pkl_dir:
pair_ids_filepath:
Returns:
"""
# 1. load image pair list
im_pair_list = load_pickle_file(pair_ids_filepath)
# 2. load pkl file paths
all_pkl_paths = sorted(glob.glob((os.path.join(pkl_dir, '*.pkl'))))
# 3. filters the source image is not front
self.im_pair_list = []
self.all_pkl_paths = []
for pairs, pkl_path in zip(im_pair_list, all_pkl_paths):
src_path = os.path.join(im_dir, pairs[0])
dst_path = os.path.join(im_dir, pairs[1])
if 'side' in src_path or 'back' in src_path:
continue
src_path = os.path.join(im_dir, src_path)
dst_path = os.path.join(im_dir, dst_path)
self.im_pair_list.append((src_path, dst_path))
self.all_pkl_paths.append(pkl_path)
assert len(self.im_pair_list) == len(
self.all_pkl_paths), '{} != {}'.format(len(self.im_pair_list),
len(self.all_pkl_paths))
self._dataset_size = len(self.im_pair_list)
del im_pair_list
del all_pkl_paths
def __init__(self,
feat_dim,
hid_dim,
out_dim,
deformed=0.1,
adj_mat_pkl_path='assets/adj_mat_info.pkl'):
super(MeshDeformation, self).__init__()
self.deformed = deformed
adj_mat_info = load_pickle_file(adj_mat_pkl_path)
adj_mat = torch_sparse_tensor(adj_mat_info['indices'],
adj_mat_info['value'],
adj_mat_info['size'])
self.deformation = GBottleneck(2,
feat_dim,
hid_dim,
out_dim,
adj_mat,
activation='relu')
self.to_verts = nn.Tanh()
def __init__(self,
pkl_path,
hresMapping_pkl_path='assets/hresMapping.pkl',
rotate=False,
isHres=False):
"""
pkl_path is the path to a SMPL model
"""
super(SMPL, self).__init__()
self.rotate = rotate
self.isHres = isHres
# -- Load SMPL params --
dd = load_pickle_file(pkl_path)
# define faces
# self.register_buffer('faces', torch.from_numpy(undo_chumpy(dd['f']).astype(np.int32)).type(dtype=torch.int32))
self.faces = torch.from_numpy(dd['f'].astype(
np.int32)).type(dtype=torch.int32)
# Mean template vertices
self.register_buffer('v_template', torch.FloatTensor(dd['v_template']))
# Size of mesh [Number of vertices, 3], (6890, 3)
self.size = [self.v_template.shape[0], 3]
self.num_betas = dd['shapedirs'].shape[-1]
# Shape blend shape basis (shapedirs): (6980, 3, 10)
# reshaped to (6980*3, 10), transposed to (10, 6980*3)
self.register_buffer(
'shapedirs',
torch.FloatTensor(
np.reshape(dd['shapedirs'], [-1, self.num_betas]).T))
# Regressor for joint locations given shape -> (24, 6890)
# Transpose to shape (6890, 24)
self.register_buffer(
'J_regressor',
torch.FloatTensor(np.asarray(dd['J_regressor'].T.todense())))
# Pose blend shape basis: (6890, 3, 207)
num_pose_basis = dd['posedirs'].shape[-1]
# Pose blend pose basis is reshaped to (6890*3, 207)
# posedirs is transposed into (207, 6890*3)
self.register_buffer(
'posedirs',
torch.FloatTensor(
np.reshape(dd['posedirs'], [-1, num_pose_basis]).T))
# indices of parents for each joints
self.parents = np.array(dd['kintree_table'][0].astype(np.int32))
# LBS weights (6890, 24)
self.register_buffer('weights', torch.FloatTensor(dd['weights']))
# This returns 19 keypoints: 6890 x 19
joint_regressor = torch.FloatTensor(
np.asarray(dd['cocoplus_regressor'].T.todense()))
self.register_buffer('joint_regressor', joint_regressor)
mapping, nf = load_pickle_file(hresMapping_pkl_path)
self.faces_hres = torch.from_numpy(nf.astype(
np.int32)).type(dtype=torch.int32)
weights_hres = torch.FloatTensor(
np.hstack([
np.expand_dims(np.mean(mapping.dot(
np.repeat(np.expand_dims(dd['weights'][:, i], -1),
3)).reshape(-1, 3),
axis=1),
axis=-1) for i in range(24)
]))
self.register_buffer('weights_hres', weights_hres)
mapping = coo_matrix_to_torch_sparse_tensor(mapping.tocoo())
self.register_buffer('mapping', mapping)