Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
| # holder of all proprietary rights on this computer program. | |
| # You can only use this computer program if you have closed | |
| # a license agreement with MPG or you get the right to use the computer | |
| # program from someone who is authorized to grant you that right. | |
| # Any use of the computer program without a valid license is prohibited and | |
| # liable to prosecution. | |
| # | |
| # Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
| # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
| # for Intelligent Systems. All rights reserved. | |
| # | |
| # Contact: ps-license@tuebingen.mpg.de | |
| import os | |
| import gc | |
| import logging | |
| from lib.common.config import cfg | |
| from lib.dataset.mesh_util import ( | |
| load_checkpoint, | |
| update_mesh_shape_prior_losses, | |
| blend_rgb_norm, | |
| unwrap, | |
| remesh, | |
| tensor2variable, | |
| rot6d_to_rotmat | |
| ) | |
| from lib.dataset.TestDataset import TestDataset | |
| from lib.common.render import query_color | |
| from lib.net.local_affine import LocalAffine | |
| from pytorch3d.structures import Meshes | |
| from apps.ICON import ICON | |
| from termcolor import colored | |
| import numpy as np | |
| from PIL import Image | |
| import trimesh | |
| import numpy as np | |
| from tqdm import tqdm | |
| import torch | |
| torch.backends.cudnn.benchmark = True | |
| logging.getLogger("trimesh").setLevel(logging.ERROR) | |
| def generate_model(in_path, model_type): | |
| torch.cuda.empty_cache() | |
| if model_type == 'ICON': | |
| model_type = 'icon-filter' | |
| else: | |
| model_type = model_type.lower() | |
| config_dict = {'loop_smpl': 100, | |
| 'loop_cloth': 200, | |
| 'patience': 5, | |
| 'out_dir': './results', | |
| 'hps_type': 'pymaf', | |
| 'config': f"./configs/{model_type}.yaml"} | |
| # cfg read and merge | |
| cfg.merge_from_file(config_dict['config']) | |
| cfg.merge_from_file("./lib/pymaf/configs/pymaf_config.yaml") | |
| os.makedirs(config_dict['out_dir'], exist_ok=True) | |
| cfg_show_list = [ | |
| "test_gpus", | |
| [0], | |
| "mcube_res", | |
| 256, | |
| "clean_mesh", | |
| True, | |
| ] | |
| cfg.merge_from_list(cfg_show_list) | |
| cfg.freeze() | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "0" | |
| device = torch.device(f"cuda:0") | |
| # load model and dataloader | |
| model = ICON(cfg) | |
| model = load_checkpoint(model, cfg) | |
| dataset_param = { | |
| 'image_path': in_path, | |
| 'seg_dir': None, | |
| 'has_det': True, # w/ or w/o detection | |
| 'hps_type': 'pymaf' # pymaf/pare/pixie | |
| } | |
| if config_dict['hps_type'] == "pixie" and "pamir" in config_dict['config']: | |
| print(colored("PIXIE isn't compatible with PaMIR, thus switch to PyMAF", "red")) | |
| dataset_param["hps_type"] = "pymaf" | |
| dataset = TestDataset(dataset_param, device) | |
| print(colored(f"Dataset Size: {len(dataset)}", "green")) | |
| pbar = tqdm(dataset) | |
| for data in pbar: | |
| pbar.set_description(f"{data['name']}") | |
| in_tensor = {"smpl_faces": data["smpl_faces"], "image": data["image"]} | |
| # The optimizer and variables | |
| optimed_pose = torch.tensor( | |
| data["body_pose"], device=device, requires_grad=True | |
| ) # [1,23,3,3] | |
| optimed_trans = torch.tensor( | |
| data["trans"], device=device, requires_grad=True | |
| ) # [3] | |
| optimed_betas = torch.tensor( | |
| data["betas"], device=device, requires_grad=True | |
| ) # [1,10] | |
| optimed_orient = torch.tensor( | |
| data["global_orient"], device=device, requires_grad=True | |
| ) # [1,1,3,3] | |
| optimizer_smpl = torch.optim.Adam( | |
| [optimed_pose, optimed_trans, optimed_betas, optimed_orient], | |
| lr=1e-3, | |
| amsgrad=True, | |
| ) | |
| scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau( | |
| optimizer_smpl, | |
| mode="min", | |
| factor=0.5, | |
| verbose=0, | |
| min_lr=1e-5, | |
| patience=config_dict['patience'], | |
| ) | |
| losses = { | |
| # Cloth: Normal_recon - Normal_pred | |
| "cloth": {"weight": 1e1, "value": 0.0}, | |
| # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2) | |
| "stiffness": {"weight": 1e5, "value": 0.0}, | |
| # Cloth: det(R) = 1 | |
| "rigid": {"weight": 1e5, "value": 0.0}, | |
| # Cloth: edge length | |
| "edge": {"weight": 0, "value": 0.0}, | |
| # Cloth: normal consistency | |
| "nc": {"weight": 0, "value": 0.0}, | |
| # Cloth: laplacian smoonth | |
| "laplacian": {"weight": 1e2, "value": 0.0}, | |
| # Body: Normal_pred - Normal_smpl | |
| "normal": {"weight": 1e0, "value": 0.0}, | |
| # Body: Silhouette_pred - Silhouette_smpl | |
| "silhouette": {"weight": 1e0, "value": 0.0}, | |
| } | |
| # smpl optimization | |
| loop_smpl = tqdm(range(config_dict['loop_smpl'])) | |
| for _ in loop_smpl: | |
| 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) | |
| if dataset_param["hps_type"] != "pixie": | |
| smpl_out = dataset.smpl_model( | |
| betas=optimed_betas, | |
| body_pose=optimed_pose_mat, | |
| global_orient=optimed_orient_mat, | |
| pose2rot=False, | |
| ) | |
| smpl_verts = ((smpl_out.vertices) + | |
| optimed_trans) * data["scale"] | |
| else: | |
| smpl_verts, _, _ = dataset.smpl_model( | |
| shape_params=optimed_betas, | |
| expression_params=tensor2variable(data["exp"], device), | |
| body_pose=optimed_pose_mat, | |
| global_pose=optimed_orient_mat, | |
| jaw_pose=tensor2variable(data["jaw_pose"], device), | |
| left_hand_pose=tensor2variable( | |
| data["left_hand_pose"], device), | |
| right_hand_pose=tensor2variable( | |
| data["right_hand_pose"], device), | |
| ) | |
| smpl_verts = (smpl_verts + optimed_trans) * data["scale"] | |
| # render optimized mesh (normal, T_normal, image [-1,1]) | |
| in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal( | |
| smpl_verts * | |
| torch.tensor([1.0, -1.0, -1.0] | |
| ).to(device), in_tensor["smpl_faces"] | |
| ) | |
| T_mask_F, T_mask_B = dataset.render.get_silhouette_image() | |
| with torch.no_grad(): | |
| in_tensor["normal_F"], in_tensor["normal_B"] = model.netG.normal_filter( | |
| in_tensor | |
| ) | |
| diff_F_smpl = torch.abs( | |
| in_tensor["T_normal_F"] - in_tensor["normal_F"]) | |
| diff_B_smpl = torch.abs( | |
| in_tensor["T_normal_B"] - in_tensor["normal_B"]) | |
| losses["normal"]["value"] = (diff_F_smpl + diff_B_smpl).mean() | |
| # silhouette loss | |
| smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0] | |
| gt_arr = torch.cat( | |
| [in_tensor["normal_F"][0], in_tensor["normal_B"][0]], dim=2 | |
| ).permute(1, 2, 0) | |
| gt_arr = ((gt_arr + 1.0) * 0.5).to(device) | |
| bg_color = ( | |
| torch.Tensor([0.5, 0.5, 0.5]).unsqueeze( | |
| 0).unsqueeze(0).to(device) | |
| ) | |
| gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float() | |
| diff_S = torch.abs(smpl_arr - gt_arr) | |
| losses["silhouette"]["value"] = diff_S.mean() | |
| # Weighted sum of the losses | |
| smpl_loss = 0.0 | |
| pbar_desc = "Body Fitting --- " | |
| for k in ["normal", "silhouette"]: | |
| pbar_desc += f"{k}: {losses[k]['value'] * losses[k]['weight']:.3f} | " | |
| smpl_loss += losses[k]["value"] * losses[k]["weight"] | |
| pbar_desc += f"Total: {smpl_loss:.3f}" | |
| loop_smpl.set_description(pbar_desc) | |
| smpl_loss.backward() | |
| optimizer_smpl.step() | |
| scheduler_smpl.step(smpl_loss) | |
| in_tensor["smpl_verts"] = smpl_verts * \ | |
| torch.tensor([1.0, 1.0, -1.0]).to(device) | |
| # visualize the optimization process | |
| # 1. SMPL Fitting | |
| # 2. Clothes Refinement | |
| os.makedirs(os.path.join(config_dict['out_dir'], cfg.name, | |
| "refinement"), exist_ok=True) | |
| # visualize the final results in self-rotation mode | |
| os.makedirs(os.path.join(config_dict['out_dir'], | |
| cfg.name, "vid"), exist_ok=True) | |
| # final results rendered as image | |
| # 1. Render the final fitted SMPL (xxx_smpl.png) | |
| # 2. Render the final reconstructed clothed human (xxx_cloth.png) | |
| # 3. Blend the original image with predicted cloth normal (xxx_overlap.png) | |
| os.makedirs(os.path.join(config_dict['out_dir'], | |
| cfg.name, "png"), exist_ok=True) | |
| # final reconstruction meshes | |
| # 1. SMPL mesh (xxx_smpl.obj) | |
| # 2. SMPL params (xxx_smpl.npy) | |
| # 3. clohted mesh (xxx_recon.obj) | |
| # 4. remeshed clothed mesh (xxx_remesh.obj) | |
| # 5. refined clothed mesh (xxx_refine.obj) | |
| os.makedirs(os.path.join(config_dict['out_dir'], | |
| cfg.name, "obj"), exist_ok=True) | |
| norm_pred_F = ( | |
| ((in_tensor["normal_F"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0) | |
| .detach() | |
| .cpu() | |
| .numpy() | |
| .astype(np.uint8) | |
| ) | |
| norm_pred_B = ( | |
| ((in_tensor["normal_B"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0) | |
| .detach() | |
| .cpu() | |
| .numpy() | |
| .astype(np.uint8) | |
| ) | |
| norm_orig_F = unwrap(norm_pred_F, data) | |
| norm_orig_B = unwrap(norm_pred_B, data) | |
| mask_orig = unwrap( | |
| np.repeat( | |
| data["mask"].permute(1, 2, 0).detach().cpu().numpy(), 3, axis=2 | |
| ).astype(np.uint8), | |
| data, | |
| ) | |
| rgb_norm_F = blend_rgb_norm(data["ori_image"], norm_orig_F, mask_orig) | |
| rgb_norm_B = blend_rgb_norm(data["ori_image"], norm_orig_B, mask_orig) | |
| Image.fromarray( | |
| np.concatenate( | |
| [data["ori_image"].astype(np.uint8), rgb_norm_F, rgb_norm_B], axis=1) | |
| ).save(os.path.join(config_dict['out_dir'], cfg.name, f"png/{data['name']}_overlap.png")) | |
| smpl_obj = trimesh.Trimesh( | |
| in_tensor["smpl_verts"].detach().cpu()[0] * | |
| torch.tensor([1.0, -1.0, 1.0]), | |
| in_tensor['smpl_faces'].detach().cpu()[0], | |
| process=False, | |
| maintains_order=True | |
| ) | |
| smpl_obj.visual.vertex_colors = (smpl_obj.vertex_normals+1.0)*255.0*0.5 | |
| smpl_obj.export( | |
| f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.obj") | |
| smpl_obj.export( | |
| f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb") | |
| smpl_info = {'betas': optimed_betas, | |
| 'pose': optimed_pose_mat, | |
| 'orient': optimed_orient_mat, | |
| 'trans': optimed_trans} | |
| np.save( | |
| f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.npy", smpl_info, allow_pickle=True) | |
| # ------------------------------------------------------------------------------------------------------------------ | |
| # cloth optimization | |
| # cloth recon | |
| in_tensor.update( | |
| dataset.compute_vis_cmap( | |
| in_tensor["smpl_verts"][0], in_tensor["smpl_faces"][0] | |
| ) | |
| ) | |
| if cfg.net.prior_type == "pamir": | |
| in_tensor.update( | |
| dataset.compute_voxel_verts( | |
| optimed_pose, | |
| optimed_orient, | |
| optimed_betas, | |
| optimed_trans, | |
| data["scale"], | |
| ) | |
| ) | |
| with torch.no_grad(): | |
| verts_pr, faces_pr, _ = model.test_single(in_tensor) | |
| recon_obj = trimesh.Trimesh( | |
| verts_pr, faces_pr, process=False, maintains_order=True | |
| ) | |
| recon_obj.visual.vertex_colors = ( | |
| recon_obj.vertex_normals+1.0)*255.0*0.5 | |
| recon_obj.export( | |
| os.path.join(config_dict['out_dir'], cfg.name, | |
| f"obj/{data['name']}_recon.obj") | |
| ) | |
| # Isotropic Explicit Remeshing for better geometry topology | |
| verts_refine, faces_refine = remesh(os.path.join(config_dict['out_dir'], cfg.name, | |
| f"obj/{data['name']}_recon.obj"), 0.5, device) | |
| # define local_affine deform verts | |
| mesh_pr = Meshes(verts_refine, faces_refine).to(device) | |
| local_affine_model = LocalAffine( | |
| mesh_pr.verts_padded().shape[1], mesh_pr.verts_padded().shape[0], mesh_pr.edges_packed()).to(device) | |
| optimizer_cloth = torch.optim.Adam( | |
| [{'params': local_affine_model.parameters()}], lr=1e-4, amsgrad=True) | |
| scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau( | |
| optimizer_cloth, | |
| mode="min", | |
| factor=0.1, | |
| verbose=0, | |
| min_lr=1e-5, | |
| patience=config_dict['patience'], | |
| ) | |
| final = None | |
| if config_dict['loop_cloth'] > 0: | |
| loop_cloth = tqdm(range(config_dict['loop_cloth'])) | |
| for _ in loop_cloth: | |
| optimizer_cloth.zero_grad() | |
| deformed_verts, stiffness, rigid = local_affine_model( | |
| verts_refine.to(device), return_stiff=True) | |
| mesh_pr = mesh_pr.update_padded(deformed_verts) | |
| # losses for laplacian, edge, normal consistency | |
| update_mesh_shape_prior_losses(mesh_pr, losses) | |
| in_tensor["P_normal_F"], in_tensor["P_normal_B"] = dataset.render_normal( | |
| mesh_pr.verts_padded(), mesh_pr.faces_padded()) | |
| diff_F_cloth = torch.abs( | |
| in_tensor["P_normal_F"] - in_tensor["normal_F"]) | |
| diff_B_cloth = torch.abs( | |
| in_tensor["P_normal_B"] - in_tensor["normal_B"]) | |
| losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean() | |
| losses["stiffness"]["value"] = torch.mean(stiffness) | |
| losses["rigid"]["value"] = torch.mean(rigid) | |
| # Weighted sum of the losses | |
| cloth_loss = torch.tensor(0.0, requires_grad=True).to(device) | |
| pbar_desc = "Cloth Refinement --- " | |
| for k in losses.keys(): | |
| if k not in ["normal", "silhouette"] and losses[k]["weight"] > 0.0: | |
| cloth_loss = cloth_loss + \ | |
| losses[k]["value"] * losses[k]["weight"] | |
| pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.5f} | " | |
| pbar_desc += f"Total: {cloth_loss:.5f}" | |
| loop_cloth.set_description(pbar_desc) | |
| # update params | |
| cloth_loss.backward() | |
| optimizer_cloth.step() | |
| scheduler_cloth.step(cloth_loss) | |
| final = trimesh.Trimesh( | |
| mesh_pr.verts_packed().detach().squeeze(0).cpu(), | |
| mesh_pr.faces_packed().detach().squeeze(0).cpu(), | |
| process=False, maintains_order=True | |
| ) | |
| # only with front texture | |
| tex_colors = query_color( | |
| mesh_pr.verts_packed().detach().squeeze(0).cpu(), | |
| mesh_pr.faces_packed().detach().squeeze(0).cpu(), | |
| in_tensor["image"], | |
| device=device, | |
| ) | |
| # full normal textures | |
| norm_colors = (mesh_pr.verts_normals_padded().squeeze( | |
| 0).detach().cpu() + 1.0) * 0.5 * 255.0 | |
| final.visual.vertex_colors = tex_colors | |
| final.export( | |
| f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.obj") | |
| final.visual.vertex_colors = norm_colors | |
| final.export( | |
| f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.glb") | |
| # always export visualized video regardless of the cloth refinment | |
| verts_lst = [smpl_obj.vertices, final.vertices] | |
| faces_lst = [smpl_obj.faces, final.faces] | |
| # self-rotated video | |
| dataset.render.load_meshes( | |
| verts_lst, faces_lst) | |
| dataset.render.get_rendered_video( | |
| [data["ori_image"], rgb_norm_F, rgb_norm_B], | |
| os.path.join(config_dict['out_dir'], cfg.name, | |
| f"vid/{data['name']}_cloth.mp4"), | |
| ) | |
| smpl_obj_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.obj" | |
| smpl_glb_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.glb" | |
| smpl_npy_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_smpl.npy" | |
| refine_obj_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.obj" | |
| refine_glb_path = f"{config_dict['out_dir']}/{cfg.name}/obj/{data['name']}_refine.glb" | |
| video_path = os.path.join( | |
| config_dict['out_dir'], cfg.name, f"vid/{data['name']}_cloth.mp4") | |
| overlap_path = os.path.join( | |
| config_dict['out_dir'], cfg.name, f"png/{data['name']}_overlap.png") | |
| # clean all the variables | |
| for element in dir(): | |
| if 'path' not in element: | |
| del locals()[element] | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| return [smpl_glb_path, smpl_obj_path,smpl_npy_path, | |
| refine_glb_path, refine_obj_path, | |
| video_path, video_path, overlap_path] | |