def set_inputs(src_info, tsf_info):
s2t_init_preds, s2t_T, s2t_tsf_inputs = initialize(src_info, tsf_info)
t2s_init_preds, t2s_T, t2s_tsf_inputs = initialize(tsf_info, src_info)
s2t_j2d = torch.cat([src_info['j2d'], tsf_info['j2d']], dim=0)
t2s_j2d = torch.cat([tsf_info['j2d'], src_info['j2d']], dim=0)
j2ds = torch.stack([s2t_j2d, t2s_j2d], dim=0)
init_preds = torch.cat([s2t_init_preds, t2s_init_preds], dim=0)
images = torch.cat([src_info['img'], tsf_info['img']], dim=0)
T = torch.cat([s2t_T, t2s_T], dim=0)
T_cycle = torch.cat([t2s_T, s2t_T], dim=0)
tsf_inputs = torch.cat([s2t_tsf_inputs, t2s_tsf_inputs], dim=0)
src_fim = torch.cat([src_info['fim'], tsf_info['fim']], dim=0)
tsf_fim = torch.cat([tsf_info['fim'], src_info['fim']], dim=0)
s2t_inputs = src_info['src_inputs']
t2s_inputs = tsf_info['src_inputs']
src_inputs = torch.cat([s2t_inputs, t2s_inputs], dim=0)
src_mask = util.morph(src_inputs[:, -1:, ], ks=self._opt.ft_ks, mode='erode')
tsf_mask = util.morph(tsf_inputs[:, -1:, ], ks=self._opt.ft_ks, mode='erode')
pseudo_masks = torch.cat([src_mask, tsf_mask], dim=0)
return src_fim, tsf_fim, j2ds, T, T_cycle, src_inputs, tsf_inputs, images, init_preds, pseudo_masks
def forward(self, aug_img, src_img, ref_img, src_smpl, ref_smpl):
# get smpl information
src_info = self._hmr.get_details(src_smpl)
ref_info = self._hmr.get_details(ref_smpl)
# process source inputs
src_f2verts, src_fim, _ = self._render.render_fim_wim(
src_info['cam'], src_info['verts'])
src_f2verts = src_f2verts[:, :, :, 0:2]
src_f2verts[:, :, :, 1] *= -1
src_cond, _ = self._render.encode_fim(src_info['cam'],
src_info['verts'],
fim=src_fim,
transpose=True)
src_crop_mask = util.morph(src_cond[:, -1:, :, :], ks=3, mode='erode')
_, ref_fim, ref_wim = self._render.render_fim_wim(
ref_info['cam'], ref_info['verts'])
ref_cond, _ = self._render.encode_fim(ref_info['cam'],
ref_info['verts'],
fim=ref_fim,
transpose=True)
T = self._render.cal_bc_transform(src_f2verts, ref_fim, ref_wim)
syn_img = F.grid_sample(src_img, T)
# src input
input_G_src = torch.cat([src_img * (1 - src_crop_mask), src_cond],
dim=1)
# tsf input
input_G_tsf = torch.cat([syn_img, ref_cond], dim=1)
# bg input
src_bg_mask = util.morph(src_cond[:, -1:, :, :], ks=15, mode='erode')
input_G_aug_bg = torch.cat([aug_img * src_bg_mask, src_bg_mask], dim=1)
input_G_src_bg = torch.cat([src_img * src_bg_mask, src_bg_mask], dim=1)
if self._opt.bg_both:
ref_bg_mask = util.morph(ref_cond[:, -1:, :, :],
ks=25,
mode='erode')
input_G_tsf_bg = torch.cat([ref_img * ref_bg_mask, ref_bg_mask],
dim=1)
else:
input_G_tsf_bg = None
# masks
tsf_crop_mask = util.morph(ref_cond[:, -1:, :, :], ks=3, mode='erode')
head_bbox = self.cal_head_bbox(ref_info['j2d'])
body_bbox = self.cal_body_bbox(ref_info['j2d'])
return input_G_aug_bg, input_G_src_bg, input_G_tsf_bg, input_G_src, input_G_tsf, \
T, src_crop_mask, tsf_crop_mask, head_bbox, body_bbox
def forward(self, src_img, src_smpl, ref_smpl):
# get smpl information
src_info = self._hmr.get_details(src_smpl)
ref_info = self._hmr.get_details(ref_smpl)
# process source inputs
src_f2verts, src_fim, _ = self._render.render_fim_wim(
src_info['cam'], src_info['verts'])
src_f2verts = src_f2verts[:, :, :, 0:2]
src_f2verts[:, :, :, 1] *= -1
src_cond, _ = self._render.encode_fim(src_info['cam'],
src_info['verts'],
fim=src_fim,
transpose=True)
src_bg_mask = util.morph(src_cond[:, -1:, :, :], ks=15, mode='erode')
src_crop_mask = util.morph(src_cond[:, -1:, :, :], ks=3, mode='erode')
# bg input
input_G_bg = torch.cat([src_img * src_bg_mask, src_bg_mask], dim=1)
# src input
input_G_src = torch.cat([src_img * (1 - src_crop_mask), src_cond],
dim=1)
# process reference inputs
_, ref_fim, ref_wim = self._render.render_fim_wim(
ref_info['cam'], ref_info['verts'])
ref_cond, _ = self._render.encode_fim(ref_info['cam'],
ref_info['verts'],
fim=ref_fim,
transpose=True)
T = self._render.cal_bc_transform(src_f2verts, ref_fim, ref_wim)
syn_img = F.grid_sample(src_img, T)
input_G_tsf = torch.cat([syn_img, ref_cond], dim=1)
# masks
tsf_crop_mask = util.morph(ref_cond[:, -1:, :, :], ks=3, mode='erode')
bg_mask = torch.cat([src_crop_mask, tsf_crop_mask], dim=0)
conds = torch.cat([src_cond, ref_cond], dim=1)
return syn_img, input_G_bg, input_G_src, input_G_tsf, T, bg_mask, conds, ref_info[
'j2d']
def preprocess(self, sample):
"""
Args:
sample (dict): items contain
--images (torch.FloatTensor): (2, 3, h, w)
--fims (torch.FloatTensor): (2, 3, h, w)
--T (torch.FloatTensor): (h, w, 2)
--warp (torch.FloatTensor): (3, h, w)
--head_bbox (torch.FloatTensor): (4), hear 4 = [lt_x, lt_y, rt_x, rt_y]
Returns:
new_sample (dict): items contain
--src_inputs (torch.FloatTensor): (3+3, h, w)
--tsf_inputs (torch.FloatTensor): (3+3, h, w)
--T (torch.FloatTensor): (h, w, 2)
--head_bbox (torch.FloatTensor): (4), hear 4 = [lt_x, lt_y, rt_x, rt_y]
--images (torch.FloatTensor): (2, 3, h, w)
--pseudo_masks (torch.FloatTensor) : (2, 1, h, w)
--bg_inputs (torch.FloatTensor): (3+1, h, w) or (2, 3+1, h, w) if self.is_both is True
"""
images = sample['images']
fims = sample['fims']
# 1. process the bg inputs
src_fim = fims[0]
src_img = images[0]
src_mask = src_fim[None, -1:, :, :] # (1, h, w)
# 2. process the src inputs
src_crop_mask = morph(src_mask,
ks=self.ft_ks,
mode='erode',
kernel=self.ft_kernel)[0]
src_inputs = torch.cat([src_img * (1 - src_crop_mask), src_fim])
# 3. process the tsf inputs
tsf_fim = fims[1]
tsf_mask = tsf_fim[None, -1:, :, :] # (1, h, w), bg is 0, front is 1
tsf_crop_mask = morph(tsf_mask,
ks=self.ft_ks,
mode='erode',
kernel=self.ft_kernel)[0]
if 'warp' not in sample:
warp = F.grid_sample(src_img[None], sample['T'][None])[0]
else:
warp = sample['warp']
tsf_inputs = torch.cat([warp, tsf_fim], dim=0)
if self.use_src_bg:
src_bg_mask = morph(
src_mask, ks=self.bg_ks, mode='erode',
kernel=self.bg_kernel)[0] # bg is 0, front is 1
bg_inputs = torch.cat([src_img * src_bg_mask, src_bg_mask], dim=0)
else:
tsf_img = images[1]
tsf_bg_mask = morph(tsf_mask,
ks=self.bg_ks,
mode='erode',
kernel=self.bg_kernel)[0]
bg_inputs = torch.cat([tsf_img * tsf_bg_mask, tsf_bg_mask], dim=0)
# 4. concat pseudo mask
pseudo_masks = torch.stack([src_crop_mask, tsf_crop_mask], dim=0)
new_sample = {
'images': images,
'pseudo_masks': pseudo_masks,
'head_bbox': sample['head_bbox'],
'body_bbox': sample['body_bbox'],
'bg_inputs': bg_inputs,
'src_inputs': src_inputs,
'tsf_inputs': tsf_inputs,
'T': sample['T'],
}
return new_sample
def personalize(self, src_path, src_smpl=None, output_path='', visualizer=None):
ori_img = cv_utils.read_cv2_img(src_path)
# resize image and convert the color space from [0, 255] to [-1, 1]
img = cv_utils.transform_img(ori_img, self._opt.image_size, transpose=True) * 2 - 1.0
img = torch.tensor(img, dtype=torch.float32).cuda()[None, ...]
if src_smpl is None:
img_hmr = cv_utils.transform_img(ori_img, 224, transpose=True) * 2 - 1.0
img_hmr = torch.tensor(img_hmr, dtype=torch.float32).cuda()[None, ...]
src_smpl = self.hmr(img_hmr)
else:
src_smpl = torch.tensor(src_smpl, dtype=torch.float32).cuda()[None, ...]
# source process, {'theta', 'cam', 'pose', 'shape', 'verts', 'j2d', 'j3d'}
src_info = self.hmr.get_details(src_smpl)
src_f2verts, src_fim, src_wim = self.render.render_fim_wim(src_info['cam'], src_info['verts'])
# src_f2pts = src_f2verts[:, :, :, 0:2]
src_info['fim'] = src_fim
src_info['wim'] = src_wim
src_info['cond'], _ = self.render.encode_fim(src_info['cam'], src_info['verts'], fim=src_fim, transpose=True)
src_info['f2verts'] = src_f2verts
src_info['p2verts'] = src_f2verts[:, :, :, 0:2]
src_info['p2verts'][:, :, :, 1] *= -1
if self._opt.only_vis:
src_info['p2verts'] = self.render.get_vis_f2pts(src_info['p2verts'], src_fim)
src_info['part'], _ = self.render.encode_fim(src_info['cam'], src_info['verts'],
fim=src_fim, transpose=True, map_fn=self.part_fn)
# add image to source info
src_info['img'] = img
src_info['image'] = ori_img
# 2. process the src inputs
if self.detector is not None:
bbox, body_mask = self.detector.inference(img[0])
bg_mask = 1 - body_mask
else:
bg_mask = util.morph(src_info['cond'][:, -1:, :, :], ks=self._opt.bg_ks, mode='erode')
body_mask = 1 - bg_mask
if self._opt.bg_model != 'ORIGINAL':
src_info['bg'] = self.bgnet(img, masks=body_mask, only_x=True)
else:
incomp_img = img * bg_mask
bg_inputs = torch.cat([incomp_img, bg_mask], dim=1)
img_bg = self.bgnet(bg_inputs)
src_info['bg_inputs'] = bg_inputs
# src_info['bg'] = img_bg
src_info['bg'] = incomp_img + img_bg * body_mask
ft_mask = 1 - util.morph(src_info['cond'][:, -1:, :, :], ks=self._opt.ft_ks, mode='erode')
src_inputs = torch.cat([img * ft_mask, src_info['cond']], dim=1)
src_info['feats'] = self.generator.encode_src(src_inputs)
src_info['src_inputs'] = src_inputs
src_info = src_info
# if visualizer is not None:
# self.visualize(visualizer, src=img, bg=src_info['bg'])
if output_path:
cv_utils.save_cv2_img(src_info['image'], output_path, image_size=self._opt.image_size)
return src_info
