Spaces:
Configuration error
Configuration error
| import torch | |
| import numpy as np | |
| import os | |
| from lib.config import cfg | |
| import trimesh | |
| def update_loss_img(output, batch): | |
| mse = torch.mean((output['rgb_map'] - batch['rgb'])**2, dim=2)[0] | |
| mse = mse.detach().cpu().numpy().astype(np.float32) | |
| # load the loss img | |
| img_path = batch['meta']['img_path'][0] | |
| paths = img_path.split('/') | |
| paths[-1] = os.path.basename(img_path).replace('.jpg', '.npy') | |
| loss_img_path = os.path.join(paths[0], 'loss', *paths[1:]) | |
| if os.path.exists(loss_img_path): | |
| loss_img = np.load(loss_img_path) | |
| else: | |
| os.system("mkdir -p '{}'".format(os.path.dirname(loss_img_path))) | |
| H, W = int(cfg.H * cfg.ratio), int(cfg.W * cfg.ratio) | |
| loss_img = mse.mean() * np.ones([H, W]).astype(np.float32) | |
| coord = batch['img_coord'][0] | |
| coord = coord.detach().cpu().numpy() | |
| loss_img[coord[:, 0], coord[:, 1]] = mse | |
| np.save(loss_img_path, loss_img) | |
| def init_smpl(smpl): | |
| data_root = 'data/light_stage' | |
| smpl_dir = os.path.join(data_root, cfg.smpl, cfg.human) | |
| for i in range(cfg.ni): | |
| smpl_path = os.path.join(smpl_dir, '{}.ply'.format(i + 1)) | |
| ply = trimesh.load(smpl_path) | |
| xyz = np.array(ply.vertices).ravel() | |
| smpl.weight.data[i] = torch.FloatTensor(xyz) | |
| return smpl | |
| def pts_to_can_pts(pts, batch): | |
| """transform pts from the world coordinate to the smpl coordinate""" | |
| Th = batch['Th'] | |
| pts = pts - Th | |
| R = batch['R'] | |
| pts = torch.matmul(pts, batch['R']) | |
| return pts | |
| def pts_to_coords(pts, min_xyz): | |
| pts = pts.clone().detach() | |
| # convert xyz to the voxel coordinate dhw | |
| dhw = pts[..., [2, 1, 0]] | |
| min_dhw = min_xyz[:, [2, 1, 0]] | |
| dhw = dhw - min_dhw[:, None] | |
| dhw = dhw / torch.tensor(cfg.voxel_size).to(dhw) | |
| return dhw | |