from core.remesh import calc_vertex_normals
from core.opt import MeshOptimizer
from utils.func import make_sparse_camera, make_round_views
from utils.render import NormalsRenderer
import torch.optim as optim
from tqdm import tqdm
from utils.video_utils import write_video
from omegaconf import OmegaConf
import numpy as np
import os
from PIL import Image
import kornia
import torch
import torch.nn as nn
import trimesh
from icecream import ic
from utils.project_mesh import multiview_color_projection, get_cameras_list
from utils.mesh_utils import to_py3d_mesh, rot6d_to_rotmat, tensor2variable
from utils.project_mesh import project_color, get_cameras_list
from utils.smpl_util import SMPLX
from lib.dataset.mesh_util import apply_vertex_mask, part_removal, poisson, keep_largest
from scipy.spatial.transform import Rotation as R
from scipy.spatial import KDTree
import argparse
#### ------------------- config----------------------
bg_color = np.array([1,1,1])
class colorModel(nn.Module):
def __init__(self, renderer, v, f, c):
super().__init__()
self.renderer = renderer
self.v = v
self.f = f
self.colors = nn.Parameter(c, requires_grad=True)
self.bg_color = torch.from_numpy(bg_color).float().to(self.colors.device)
def forward(self, return_mask=False):
rgba = self.renderer.render(self.v, self.f, colors=self.colors)
if return_mask:
return rgba
else:
mask = rgba[..., 3:]
return rgba[..., :3] * mask + self.bg_color * (1 - mask)
def scale_mesh(vert):
min_bbox, max_bbox = vert.min(0)[0], vert.max(0)[0]
center = (min_bbox + max_bbox) / 2
offset = -center
vert = vert + offset
max_dist = torch.max(torch.sqrt(torch.sum(vert**2, dim=1)))
scale = 1.0 / max_dist
return scale, offset
def save_mesh(save_name, vertices, faces, color=None):
trimesh.Trimesh(
vertices.detach().cpu().numpy(),
faces.detach().cpu().numpy(),
vertex_colors=(color.detach().cpu().numpy() * 255).astype(np.uint8) if color is not None else None) \
.export(save_name)
class ReMesh:
def __init__(self, opt, econ_dataset):
self.opt = opt
self.device = torch.device(f"cuda:{opt.gpu_id}" if torch.cuda.is_available() else "cpu")
self.num_view = opt.num_view
self.out_path = opt.res_path
os.makedirs(self.out_path, exist_ok=True)
self.resolution = opt.resolution
self.views = ['front_face', 'front_right', 'right', 'back', 'left', 'front_left' ]
self.weights = torch.Tensor([1., 0.4, 0.8, 1.0, 0.8, 0.4]).view(6,1,1,1).to(self.device)
self.renderer = self.prepare_render()
# pose prediction
self.econ_dataset = econ_dataset
self.smplx_face = torch.Tensor(econ_dataset.faces.astype(np.int64)).long().to(self.device)
def prepare_render(self):
### ------------------- prepare camera and renderer----------------------
mv, proj = make_sparse_camera(self.opt.cam_path, self.opt.scale, views=[0,1,2,4,6,7], device=self.device)
renderer = NormalsRenderer(mv, proj, [self.resolution, self.resolution], device=self.device)
return renderer
def proj_texture(self, fused_images, vertices, faces):
mesh = to_py3d_mesh(vertices, faces)
mesh = mesh.to(self.device)
camera_focal = 1/2
cameras_list = get_cameras_list([0, 45, 90, 180, 270, 315], device=self.device, focal=camera_focal)
mesh = multiview_color_projection(mesh, fused_images, camera_focal=camera_focal, resolution=self.resolution, weights=self.weights.squeeze().cpu().numpy(),
device=self.device, complete_unseen=True, confidence_threshold=0.2, cameras_list=cameras_list)
return mesh
def get_invisible_idx(self, imgs, vertices, faces):
mesh = to_py3d_mesh(vertices, faces)
mesh = mesh.to(self.device)
camera_focal = 1/2
if self.num_view == 6:
cameras_list = get_cameras_list([0, 45, 90, 180, 270, 315], device=self.device, focal=camera_focal)
elif self.num_view == 4:
cameras_list = get_cameras_list([0, 90, 180, 270], device=self.device, focal=camera_focal)
valid_vert_id = []
vertices_colors = torch.zeros((vertices.shape[0], 3)).float().to(self.device)
valid_cnt = torch.zeros((vertices.shape[0])).to(self.device)
for cam, img, weight in zip(cameras_list, imgs, self.weights.squeeze()):
ret = project_color(mesh, cam, img, eps=0.01, resolution=self.resolution, device=self.device)
# print(ret['valid_colors'].shape)
valid_cnt[ret['valid_verts']] += weight
vertices_colors[ret['valid_verts']] += ret['valid_colors']*weight
valid_mask = valid_cnt > 1
invalid_mask = valid_cnt < 1
vertices_colors[valid_mask] /= valid_cnt[valid_mask][:, None]
# visibility
invisible_vert = valid_cnt < 1
invisible_vert_indices = torch.nonzero(invisible_vert).squeeze()
# vertices_colors[invalid_vert] = torch.tensor([1.0, 0.0, 0.0]).float().to("cuda")
return vertices_colors, invisible_vert_indices
def inpaint_missed_colors(self, all_vertices, all_colors, missing_indices):
all_vertices = all_vertices.detach().cpu().numpy()
all_colors = all_colors.detach().cpu().numpy()
missing_indices = missing_indices.detach().cpu().numpy()
non_missing_indices = np.setdiff1d(np.arange(len(all_vertices)), missing_indices)
kdtree = KDTree(all_vertices[non_missing_indices])
for missing_index in missing_indices:
missing_vertex = all_vertices[missing_index]
_, nearest_index = kdtree.query(missing_vertex.reshape(1, -1))
interpolated_color = all_colors[non_missing_indices[nearest_index]]
all_colors[missing_index] = interpolated_color
return torch.from_numpy(all_colors).to(self.device)
def load_training_data(self, case):
###------------------ load target images -------------------------------
kernal = torch.ones(3, 3)
erode_iters = 2
normals = []
masks = []
colors = []
for idx, view in enumerate(self.views):
# for idx in [0,2,3,4]:
normal = Image.open(f'{self.opt.mv_path}/{case}/normals_{view}_masked.png')
# normal = Image.open(f'{data_path}/{case}/normals/{idx:02d}_rgba.png')
normal = normal.convert('RGBA').resize((self.resolution, self.resolution), Image.BILINEAR)
normal = np.array(normal).astype(np.float32) / 255.
mask = normal[..., 3:] # alpha
mask_troch = torch.from_numpy(mask).unsqueeze(0)
for _ in range(erode_iters):
mask_torch = kornia.morphology.erosion(mask_troch, kernal)
mask_erode = mask_torch.squeeze(0).numpy()
masks.append(mask_erode)
normal = normal[..., :3] * mask_erode
normals.append(normal)
color = Image.open(f'{self.opt.mv_path}/{case}/color_{view}_masked.png')
color = color.convert('RGBA').resize((self.resolution, self.resolution), Image.BILINEAR)
color = np.array(color).astype(np.float32) / 255.
color_mask = color[..., 3:] # alpha
# color_dilate = color[..., :3] * color_mask + bg_color * (1 - color_mask)
color_dilate = color[..., :3] * mask_erode + bg_color * (1 - mask_erode)
colors.append(color_dilate)
masks = np.stack(masks, 0)
masks = torch.from_numpy(masks).to(self.device)
normals = np.stack(normals, 0)
target_normals = torch.from_numpy(normals).to(self.device)
colors = np.stack(colors, 0)
target_colors = torch.from_numpy(colors).to(self.device)
return masks, target_colors, target_normals
def preprocess(self, color_pils, normal_pils):
###------------------ load target images -------------------------------
kernal = torch.ones(3, 3)
erode_iters = 2
normals = []
masks = []
colors = []
for normal, color in zip(normal_pils, color_pils):
normal = normal.resize((self.resolution, self.resolution), Image.BILINEAR)
normal = np.array(normal).astype(np.float32) / 255.
mask = normal[..., 3:] # alpha
mask_troch = torch.from_numpy(mask).unsqueeze(0)
for _ in range(erode_iters):
mask_torch = kornia.morphology.erosion(mask_troch, kernal)
mask_erode = mask_torch.squeeze(0).numpy()
masks.append(mask_erode)
normal = normal[..., :3] * mask_erode
normals.append(normal)
color = color.resize((self.resolution, self.resolution), Image.BILINEAR)
color = np.array(color).astype(np.float32) / 255.
color_mask = color[..., 3:] # alpha
# color_dilate = color[..., :3] * color_mask + bg_color * (1 - color_mask)
color_dilate = color[..., :3] * mask_erode + bg_color * (1 - mask_erode)
colors.append(color_dilate)
masks = np.stack(masks, 0)
masks = torch.from_numpy(masks).to(self.device)
normals = np.stack(normals, 0)
target_normals = torch.from_numpy(normals).to(self.device)
colors = np.stack(colors, 0)
target_colors = torch.from_numpy(colors).to(self.device)
return masks, target_colors, target_normals
def optimize_case(self, case, pose, clr_img, nrm_img, opti_texture=True):
case_path = f'{self.out_path}/{case}'
os.makedirs(case_path, exist_ok=True)
if clr_img is not None:
masks, target_colors, target_normals = self.preprocess(clr_img, nrm_img)
else:
masks, target_colors, target_normals = self.load_training_data(case)
# rotation
rz = R.from_euler('z', 180, degrees=True).as_matrix()
ry = R.from_euler('y', 180, degrees=True).as_matrix()
rz = torch.from_numpy(rz).float().to(self.device)
ry = torch.from_numpy(ry).float().to(self.device)
scale, offset = None, None
global_orient = pose["global_orient"] # pymaf_res[idx]['smplx_params']['body_pose'][:, :1, :, :2].to(device).reshape(1, 1, -1) # data["global_orient"]
body_pose = pose["body_pose"] # pymaf_res[idx]['smplx_params']['body_pose'][:, 1:22, :, :2].to(device).reshape(1, 21, -1) # data["body_pose"]
left_hand_pose = pose["left_hand_pose"] # pymaf_res[idx]['smplx_params']['left_hand_pose'][:, :, :, :2].to(device).reshape(1, 15, -1)
right_hand_pose = pose["right_hand_pose"] # pymaf_res[idx]['smplx_params']['right_hand_pose'][:, :, :, :2].to(device).reshape(1, 15, -1)
beta = pose["betas"]
# The optimizer and variables
optimed_pose = torch.tensor(body_pose,
device=self.device,
requires_grad=True) # [1,23,3,3]
optimed_trans = torch.tensor(pose["trans"],
device=self.device,
requires_grad=True) # [3]
optimed_betas = torch.tensor(beta,
device=self.device,
requires_grad=True) # [1,200]
optimed_orient = torch.tensor(global_orient,
device=self.device,
requires_grad=True) # [1,1,3,3]
optimed_rhand = torch.tensor(right_hand_pose,
device=self.device,
requires_grad=True)
optimed_lhand = torch.tensor(left_hand_pose,
device=self.device,
requires_grad=True)
optimed_params = [
{'params': [optimed_lhand, optimed_rhand], 'lr': 1e-3},
{'params': [optimed_betas, optimed_trans, optimed_orient, optimed_pose], 'lr': 3e-3},
]
optimizer_smpl = torch.optim.Adam(
optimed_params,
amsgrad=True,
)
scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer_smpl,
mode="min",
factor=0.5,
verbose=0,
min_lr=1e-5,
patience=5,
)
smpl_steps = 100
for i in tqdm(range(smpl_steps)):
optimizer_smpl.zero_grad()
# 6d_rot to rot_mat
optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(
-1, 6)).unsqueeze(0)
optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(
-1, 6)).unsqueeze(0)
smpl_verts, smpl_landmarks, smpl_joints = self.econ_dataset.smpl_model(
shape_params=optimed_betas,
expression_params=tensor2variable(pose["exp"], self.device),
body_pose=optimed_pose_mat,
global_pose=optimed_orient_mat,
jaw_pose=tensor2variable(pose["jaw_pose"], self.device),
left_hand_pose=optimed_lhand,
right_hand_pose=optimed_rhand,
)
smpl_verts = smpl_verts + optimed_trans
v_smpl = torch.matmul(torch.matmul(smpl_verts.squeeze(0), rz.T), ry.T)
if scale is None:
scale, offset = scale_mesh(v_smpl.detach())
v_smpl = (v_smpl + offset) * scale * 2
# if i == 0:
# save_mesh(f'{case_path}/{case}_init_smpl.obj', v_smpl, self.smplx_face)
# exit()
normals = calc_vertex_normals(v_smpl, self.smplx_face)
nrm = self.renderer.render(v_smpl, self.smplx_face, normals=normals)
masks_ = nrm[..., 3:]
smpl_mask_loss = ((masks_ - masks) * self.weights).abs().mean()
smpl_nrm_loss = ((nrm[..., :3] - target_normals) * self.weights).abs().mean()
smpl_loss = smpl_mask_loss + smpl_nrm_loss
# smpl_loss = smpl_mask_loss
smpl_loss.backward()
optimizer_smpl.step()
scheduler_smpl.step(smpl_loss)
mesh_smpl = trimesh.Trimesh(vertices=v_smpl.detach().cpu().numpy(), faces=self.smplx_face.detach().cpu().numpy())
nrm_opt = MeshOptimizer(v_smpl.detach(), self.smplx_face.detach(), edge_len_lims=[0.01, 0.1])
vertices, faces = nrm_opt.vertices, nrm_opt.faces
# ###----------------------- optimization iterations-------------------------------------
for i in tqdm(range(self.opt.iters)):
nrm_opt.zero_grad()
normals = calc_vertex_normals(vertices,faces)
nrm = self.renderer.render(vertices,faces, normals=normals)
normals = nrm[..., :3]
# if i < 800:
loss = ((normals-target_normals) * self.weights).abs().mean()
# else:
# loss = ((normals-target_images) * masks).abs().mean()
alpha = nrm[..., 3:]
loss += ((alpha - masks) * self.weights).abs().mean()
loss.backward()
nrm_opt.step()
vertices,faces = nrm_opt.remesh()
if self.opt.debug and i % self.opt.snapshot_step == 0:
import imageio
os.makedirs(f'{case_path}/normals', exist_ok=True)
imageio.imwrite(f'{case_path}/normals/{i:02d}.png',(nrm.detach()[0,:,:,:3]*255).clamp(max=255).type(torch.uint8).cpu().numpy())
# mesh_remeshed = trimesh.Trimesh(vertices=vertices.detach().cpu().numpy(), faces=faces.detach().cpu().numpy())
# mesh_remeshed.export(f'{case_path}/{case}_remeshed_step{i}.obj')
torch.cuda.empty_cache()
mesh_remeshed = trimesh.Trimesh(vertices=vertices.detach().cpu().numpy(), faces=faces.detach().cpu().numpy())
mesh_remeshed.export(f'{case_path}/{case}_remeshed.obj')
# save_mesh(case, vertices, faces)
vertices = vertices.detach()
faces = faces.detach()
#### replace hand
smpl_data = SMPLX()
if self.opt.replace_hand and True in pose['hands_visibility'][0]:
hand_mask = torch.zeros(smpl_data.smplx_verts.shape[0], )
if pose['hands_visibility'][0][0]:
hand_mask.index_fill_(
0, torch.tensor(smpl_data.smplx_mano_vid_dict["left_hand"]), 1.0
)
if pose['hands_visibility'][0][1]:
hand_mask.index_fill_(
0, torch.tensor(smpl_data.smplx_mano_vid_dict["right_hand"]), 1.0
)
hand_mesh = apply_vertex_mask(mesh_smpl.copy(), hand_mask)
body_mesh = part_removal(
mesh_remeshed.copy(),
hand_mesh,
0.08,
self.device,
mesh_smpl.copy(),
region="hand"
)
final = poisson(sum([hand_mesh, body_mesh]), f'{case_path}/{case}_final.obj', 10, False)
else:
final = poisson(mesh_remeshed, f'{case_path}/{case}_final.obj', 10, False)
vertices = torch.from_numpy(final.vertices).float().to(self.device)
faces = torch.from_numpy(final.faces).long().to(self.device)
# Differing from paper, we use the texturing method in Unique3D
masked_color = []
for tmp in clr_img:
# tmp = Image.open(f'{self.opt.mv_path}/{case}/color_{view}_masked.png')
tmp = tmp.resize((self.resolution, self.resolution), Image.BILINEAR)
tmp = np.array(tmp).astype(np.float32) / 255.
masked_color.append(torch.from_numpy(tmp).permute(2, 0, 1).to(self.device))
meshes = self.proj_texture(masked_color, vertices, faces)
vertices = meshes.verts_packed().float()
faces = meshes.faces_packed().long()
colors = meshes.textures.verts_features_packed().float()
save_mesh(f'./{case_path}/result_clr_scale{self.opt.scale}_{case}.obj', vertices, faces, colors)
self.evaluate(vertices, colors, faces, save_path=f'{case_path}/result_clr_scale{self.opt.scale}_{case}.mp4', save_nrm=True)
def evaluate(self, target_vertices, target_colors, target_faces, save_path=None, save_nrm=False):
mv, proj = make_round_views(60, self.opt.scale, device=self.device)
renderer = NormalsRenderer(mv, proj, [512, 512], device=self.device)
target_images = renderer.render(target_vertices,target_faces, colors=target_colors)
target_images = target_images.detach().cpu().numpy()
target_images = target_images[..., :3] * target_images[..., 3:4] + bg_color * (1 - target_images[..., 3:4])
target_images = (target_images.clip(0, 1) * 255).astype(np.uint8)
if save_nrm:
target_normals = calc_vertex_normals(target_vertices, target_faces)
# target_normals[:, 2] *= -1
target_normals = renderer.render(target_vertices, target_faces, normals=target_normals)
target_normals = target_normals.detach().cpu().numpy()
target_normals = target_normals[..., :3] * target_normals[..., 3:4] + bg_color * (1 - target_normals[..., 3:4])
target_normals = (target_normals.clip(0, 1) * 255).astype(np.uint8)
frames = [np.concatenate([img, nrm], 1) for img, nrm in zip(target_images, target_normals)]
else:
frames = [img for img in target_images]
if save_path is not None:
write_video(frames, fps=25, save_path=save_path)
return frames
def run(self):
cases = sorted(os.listdir(self.opt.imgs_path))
for idx in range(len(cases)):
case = cases[idx].split('.')[0]
print(f'Processing {case}')
pose = self.econ_dataset.__getitem__(idx)
v, f, c = self.optimize_case(case, pose, None, None, opti_texture=True)
self.evaluate(v, c, f, save_path=f'{self.opt.res_path}/{case}/result_clr_scale{self.opt.scale}_{case}.mp4', save_nrm=True)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--config", help="path to the yaml configs file", default='config.yaml')
args, extras = parser.parse_known_args()
opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))
from econdataset import SMPLDataset
dataset_param = {'image_dir': opt.imgs_path, 'seg_dir': None, 'colab': False, 'has_det': True, 'hps_type': 'pixie'}
econdata = SMPLDataset(dataset_param, device='cuda')
EHuman = ReMesh(opt, econdata)
EHuman.run()