# -*- 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 logging import warnings warnings.filterwarnings("ignore") logging.getLogger("lightning").setLevel(logging.ERROR) logging.getLogger("trimesh").setLevel(logging.ERROR) import os import numpy as np import torch import torchvision import trimesh from pytorch3d.ops import SubdivideMeshes from huggingface_hub import hf_hub_download from termcolor import colored from tqdm import tqdm from apps.IFGeo import IFGeo from apps.Normal import Normal from lib.common.BNI import BNI from lib.common.BNI_utils import save_normal_tensor from lib.common.config import cfg from lib.common.imutils import blend_rgb_norm from lib.common.local_affine import register from lib.common.render import query_color, Render from lib.common.train_util import Format, init_loss from lib.common.voxelize import VoxelGrid from lib.dataset.mesh_util import * from lib.dataset.TestDataset import TestDataset from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis torch.backends.cudnn.benchmark = True def generate_video(vis_tensor_path): in_tensor = torch.load(vis_tensor_path) render = Render(size=512, device=torch.device("cuda:0")) # visualize the final results in self-rotation mode verts_lst = in_tensor["body_verts"] + in_tensor["BNI_verts"] faces_lst = in_tensor["body_faces"] + in_tensor["BNI_faces"] # self-rotated video tmp_path = vis_tensor_path.replace("_in_tensor.pt", "_tmp.mp4") out_path = vis_tensor_path.replace("_in_tensor.pt", ".mp4") render.load_meshes(verts_lst, faces_lst) render.get_rendered_video_multi(in_tensor, tmp_path) os.system(f"ffmpeg -y -loglevel quiet -stats -i {tmp_path} -vcodec libx264 {out_path}") return out_path import sys class Logger: def __init__(self, filename): self.terminal = sys.stdout self.log = open(filename, "w") def write(self, message): self.terminal.write(message) self.log.write(message) def flush(self): self.terminal.flush() self.log.flush() def isatty(self): return False def generate_model(in_path, fitting_step=50): sys.stdout = Logger("./output.log") out_dir = "./results" # cfg read and merge cfg.merge_from_file("./configs/econ.yaml") cfg.merge_from_file("./lib/pymafx/configs/pymafx_config.yaml") device = torch.device(f"cuda:0") # setting for testing on in-the-wild images cfg_show_list = [ "test_gpus", [0], "mcube_res", 512, "clean_mesh", True, "test_mode", True, "batch_size", 1 ] cfg.merge_from_list(cfg_show_list) cfg.freeze() # load normal model normal_net = Normal.load_from_checkpoint( cfg=cfg, checkpoint_path=hf_hub_download( repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.normal_path ), map_location=device, strict=False ) normal_net = normal_net.to(device) normal_net.netG.eval() print( colored( f"Resume Normal Estimator from : {cfg.normal_path} ", "green" ) ) # SMPLX object SMPLX_object = SMPLX() dataset_param = { "image_path": in_path, "use_seg": True, # w/ or w/o segmentation "hps_type": cfg.bni.hps_type, # pymafx/pixie "vol_res": cfg.vol_res, "single": True, } if cfg.bni.use_ifnet: # load IFGeo model ifnet = IFGeo.load_from_checkpoint( cfg=cfg, checkpoint_path=hf_hub_download( repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.ifnet_path ), map_location=device, strict=False ) ifnet = ifnet.to(device) ifnet.netG.eval() print(colored(f"Resume IF-Net+ from : {cfg.ifnet_path} ", "green")) print(colored(f"Complete with : IF-Nets+ (Implicit) ", "green")) else: print(colored(f"Complete with : SMPL-X (Explicit) ", "green")) dataset = TestDataset(dataset_param, device) print(colored(f"Dataset Size: {len(dataset)}", "green")) data = dataset[0] losses = init_loss() print(f"Subject name: {data['name']}") # final results rendered as image (PNG) # 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) # 4. Blend the cropped image with predicted cloth normal (xxx_crop.png) os.makedirs(osp.join(out_dir, cfg.name, "png"), exist_ok=True) # final reconstruction meshes (OBJ) # 1. SMPL mesh (xxx_smpl_xx.obj) # 2. SMPL params (xxx_smpl.npy) # 3. d-BiNI surfaces (xxx_BNI.obj) # 4. seperate face/hand mesh (xxx_hand/face.obj) # 5. full shape impainted by IF-Nets+ after remeshing (xxx_IF.obj) # 6. sideded or occluded parts (xxx_side.obj) # 7. final reconstructed clothed human (xxx_full.obj) os.makedirs(osp.join(out_dir, cfg.name, "obj"), exist_ok=True) in_tensor = { "smpl_faces": data["smpl_faces"], "image": data["img_icon"].to(device), "mask": data["img_mask"].to(device) } # The optimizer and variables optimed_pose = data["body_pose"].requires_grad_(True) optimed_trans = data["trans"].requires_grad_(True) optimed_betas = data["betas"].requires_grad_(True) optimed_orient = data["global_orient"].requires_grad_(True) optimizer_smpl = torch.optim.Adam([optimed_pose, optimed_trans, optimed_betas, optimed_orient], lr=1e-2, amsgrad=True) scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer_smpl, mode="min", factor=0.5, verbose=0, min_lr=1e-5, patience=5, ) # [result_loop_1, result_loop_2, ...] per_data_lst = [] N_body, N_pose = optimed_pose.shape[:2] smpl_path = f"{out_dir}/{cfg.name}/obj/{data['name']}_smpl_00.obj" # remove this line if you change the loop_smpl and obtain different SMPL-X fits if osp.exists(smpl_path): smpl_verts_lst = [] smpl_faces_lst = [] for idx in range(N_body): smpl_obj = f"{out_dir}/{cfg.name}/obj/{data['name']}_smpl_{idx:02d}.obj" smpl_mesh = trimesh.load(smpl_obj) smpl_verts = torch.tensor(smpl_mesh.vertices).to(device).float() smpl_faces = torch.tensor(smpl_mesh.faces).to(device).long() smpl_verts_lst.append(smpl_verts) smpl_faces_lst.append(smpl_faces) batch_smpl_verts = torch.stack(smpl_verts_lst) batch_smpl_faces = torch.stack(smpl_faces_lst) # render optimized mesh as normal [-1,1] in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal( batch_smpl_verts, batch_smpl_faces ) with torch.no_grad(): in_tensor["normal_F"], in_tensor["normal_B"] = normal_net.netG(in_tensor) in_tensor["smpl_verts"] = batch_smpl_verts * torch.tensor([1., -1., 1.]).to(device) in_tensor["smpl_faces"] = batch_smpl_faces[:, :, [0, 2, 1]] else: # smpl optimization loop_smpl = tqdm(range(fitting_step)) for i in loop_smpl: per_loop_lst = [] optimizer_smpl.zero_grad() N_body, N_pose = optimed_pose.shape[:2] # 6d_rot to rot_mat optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1, 6)).view(N_body, 1, 3, 3) optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1, 6)).view(N_body, N_pose, 3, 3) smpl_verts, smpl_landmarks, smpl_joints = 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"] smpl_joints = (smpl_joints + optimed_trans) * data["scale"] * torch.tensor([ 1.0, 1.0, -1.0 ]).to(device) # landmark errors smpl_joints_3d = ( smpl_joints[:, dataset.smpl_data.smpl_joint_ids_45_pixie, :] + 1.0 ) * 0.5 in_tensor["smpl_joint"] = smpl_joints[:, dataset.smpl_data.smpl_joint_ids_24_pixie, :] ghum_lmks = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], :2].to(device) ghum_conf = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], -1].to(device) smpl_lmks = smpl_joints_3d[:, SMPLX_object.ghum_smpl_pairs[:, 1], :2].to(device) # render optimized mesh as normal [-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_image(type="mask") with torch.no_grad(): in_tensor["normal_F"], in_tensor["normal_B"] = normal_net.netG(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"]) # silhouette loss smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1) gt_arr = in_tensor["mask"].repeat(1, 1, 2) diff_S = torch.abs(smpl_arr - gt_arr) losses["silhouette"]["value"] = diff_S.mean() # large cloth_overlap --> big difference between body and cloth mask # for loose clothing, reply more on landmarks instead of silhouette+normal loss cloth_overlap = diff_S.sum(dim=[1, 2]) / gt_arr.sum(dim=[1, 2]) cloth_overlap_flag = cloth_overlap > cfg.cloth_overlap_thres losses["joint"]["weight"] = [50.0 if flag else 5.0 for flag in cloth_overlap_flag] # small body_overlap --> large occlusion or out-of-frame # for highly occluded body, reply only on high-confidence landmarks, no silhouette+normal loss # BUG: PyTorch3D silhouette renderer generates dilated mask bg_value = in_tensor["T_normal_F"][0, 0, 0, 0].to(device) smpl_arr_fake = torch.cat([ in_tensor["T_normal_F"][:, 0].ne(bg_value).float(), in_tensor["T_normal_B"][:, 0].ne(bg_value).float() ], dim=-1) body_overlap = (gt_arr * smpl_arr_fake.gt(0.0) ).sum(dim=[1, 2]) / smpl_arr_fake.gt(0.0).sum(dim=[1, 2]) body_overlap_mask = (gt_arr * smpl_arr_fake).unsqueeze(1) body_overlap_flag = body_overlap < cfg.body_overlap_thres losses["normal"]["value"] = ( diff_F_smpl * body_overlap_mask[..., :512] + diff_B_smpl * body_overlap_mask[..., 512:] ).mean() / 2.0 losses["silhouette"]["weight"] = [0 if flag else 1.0 for flag in body_overlap_flag] occluded_idx = torch.where(body_overlap_flag)[0] ghum_conf[occluded_idx] *= ghum_conf[occluded_idx] > 0.95 losses["joint"]["value"] = (torch.norm(ghum_lmks - smpl_lmks, dim=2) * ghum_conf).mean(dim=1) # Weighted sum of the losses smpl_loss = 0.0 pbar_desc = "Body Fitting -- " for k in ["normal", "silhouette", "joint"]: per_loop_loss = (losses[k]["value"] * torch.tensor(losses[k]["weight"]).to(device)).mean() pbar_desc += f"{k}: {per_loop_loss:.3f} | " smpl_loss += per_loop_loss pbar_desc += f"Total: {smpl_loss:.3f}" loose_str = ''.join([str(j) for j in cloth_overlap_flag.int().tolist()]) occlude_str = ''.join([str(j) for j in body_overlap_flag.int().tolist()]) pbar_desc += colored(f"| loose:{loose_str}, occluded:{occlude_str}", "yellow") loop_smpl.set_description(pbar_desc) print(pbar_desc) # save intermediate results if (i == fitting_step - 1): per_loop_lst.extend([ in_tensor["image"], in_tensor["T_normal_F"], in_tensor["normal_F"], diff_S[:, :, :512].unsqueeze(1).repeat(1, 3, 1, 1), ]) per_loop_lst.extend([ in_tensor["image"], in_tensor["T_normal_B"], in_tensor["normal_B"], diff_S[:, :, 512:].unsqueeze(1).repeat(1, 3, 1, 1), ]) per_data_lst.append( get_optim_grid_image(per_loop_lst, None, nrow=N_body * 2, type="smpl") ) 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) in_tensor["smpl_faces"] = in_tensor["smpl_faces"][:, :, [0, 2, 1]] per_data_lst[-1].save(osp.join(out_dir, cfg.name, f"png/{data['name']}_smpl.png")) img_crop_path = osp.join(out_dir, cfg.name, "png", f"{data['name']}_crop.png") torchvision.utils.save_image( torch.cat([ data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5, (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5 ], dim=3), img_crop_path ) rgb_norm_F = blend_rgb_norm(in_tensor["normal_F"], data) rgb_norm_B = blend_rgb_norm(in_tensor["normal_B"], data) img_overlap_path = osp.join(out_dir, cfg.name, f"png/{data['name']}_overlap.png") torchvision.utils.save_image( torch.cat([data["img_raw"], rgb_norm_F, rgb_norm_B], dim=-1) / 255., img_overlap_path ) smpl_obj_lst = [] for idx in range(N_body): smpl_obj = trimesh.Trimesh( in_tensor["smpl_verts"].detach().cpu()[idx] * torch.tensor([1.0, -1.0, 1.0]), in_tensor["smpl_faces"].detach().cpu()[0][:, [0, 2, 1]], process=False, maintains_order=True, ) smpl_obj_path = f"{out_dir}/{cfg.name}/obj/{data['name']}_smpl_{idx:02d}.obj" if not osp.exists(smpl_obj_path): smpl_obj.export(smpl_obj_path) smpl_obj.export(smpl_obj_path.replace(".obj", ".glb")) smpl_info = { "betas": optimed_betas[idx].detach().cpu().unsqueeze(0), "body_pose": rotation_matrix_to_angle_axis(optimed_pose_mat[idx].detach()).cpu().unsqueeze(0), "global_orient": rotation_matrix_to_angle_axis(optimed_orient_mat[idx].detach()).cpu().unsqueeze(0), "transl": optimed_trans[idx].detach().cpu(), "expression": data["exp"][idx].cpu().unsqueeze(0), "jaw_pose": rotation_matrix_to_angle_axis(data["jaw_pose"][idx]).cpu().unsqueeze(0), "left_hand_pose": rotation_matrix_to_angle_axis(data["left_hand_pose"][idx]).cpu().unsqueeze(0), "right_hand_pose": rotation_matrix_to_angle_axis(data["right_hand_pose"][idx]).cpu().unsqueeze(0), "scale": data["scale"][idx].cpu(), } np.save( smpl_obj_path.replace(".obj", ".npy"), smpl_info, allow_pickle=True, ) smpl_obj_lst.append(smpl_obj) del optimizer_smpl del optimed_betas del optimed_orient del optimed_pose del optimed_trans torch.cuda.empty_cache() # ------------------------------------------------------------------------------------------------------------------ # clothing refinement per_data_lst = [] batch_smpl_verts = in_tensor["smpl_verts"].detach() * torch.tensor([1.0, -1.0, 1.0], device=device) batch_smpl_faces = in_tensor["smpl_faces"].detach()[:, :, [0, 2, 1]] in_tensor["depth_F"], in_tensor["depth_B"] = dataset.render_depth( batch_smpl_verts, batch_smpl_faces ) per_loop_lst = [] in_tensor["BNI_verts"] = [] in_tensor["BNI_faces"] = [] in_tensor["body_verts"] = [] in_tensor["body_faces"] = [] for idx in range(N_body): final_path = f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_full.obj" side_mesh = smpl_obj_lst[idx].copy() face_mesh = smpl_obj_lst[idx].copy() hand_mesh = smpl_obj_lst[idx].copy() smplx_mesh = smpl_obj_lst[idx].copy() # save normals, depths and masks BNI_dict = save_normal_tensor( in_tensor, idx, osp.join(out_dir, cfg.name, f"BNI/{data['name']}_{idx}"), cfg.bni.thickness, ) # BNI process BNI_object = BNI( dir_path=osp.join(out_dir, cfg.name, "BNI"), name=data["name"], BNI_dict=BNI_dict, cfg=cfg.bni, device=device ) BNI_object.extract_surface(False) in_tensor["body_verts"].append(torch.tensor(smpl_obj_lst[idx].vertices).float()) in_tensor["body_faces"].append(torch.tensor(smpl_obj_lst[idx].faces).long()) # requires shape completion when low overlap # replace SMPL by completed mesh as side_mesh if cfg.bni.use_ifnet: side_mesh_path = f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_IF.obj" side_mesh = apply_face_mask(side_mesh, ~SMPLX_object.smplx_eyeball_fid_mask) # mesh completion via IF-net in_tensor.update( dataset.depth_to_voxel({ "depth_F": BNI_object.F_depth.unsqueeze(0), "depth_B": BNI_object.B_depth.unsqueeze(0) }) ) occupancies = VoxelGrid.from_mesh(side_mesh, cfg.vol_res, loc=[ 0, ] * 3, scale=2.0).data.transpose(2, 1, 0) occupancies = np.flip(occupancies, axis=1) in_tensor["body_voxels"] = torch.tensor(occupancies.copy() ).float().unsqueeze(0).to(device) with torch.no_grad(): sdf = ifnet.reconEngine(netG=ifnet.netG, batch=in_tensor) verts_IF, faces_IF = ifnet.reconEngine.export_mesh(sdf) if ifnet.clean_mesh_flag: verts_IF, faces_IF = clean_mesh(verts_IF, faces_IF) side_mesh = trimesh.Trimesh(verts_IF, faces_IF) side_mesh = remesh_laplacian(side_mesh, side_mesh_path) else: side_mesh = apply_vertex_mask( side_mesh, ( SMPLX_object.front_flame_vertex_mask + SMPLX_object.smplx_mano_vertex_mask + SMPLX_object.eyeball_vertex_mask ).eq(0).float(), ) #register side_mesh to BNI surfaces side_mesh = Meshes( verts=[torch.tensor(side_mesh.vertices).float()], faces=[torch.tensor(side_mesh.faces).long()], ).to(device) sm = SubdivideMeshes(side_mesh) side_mesh = register(BNI_object.F_B_trimesh, sm(side_mesh), device) side_verts = torch.tensor(side_mesh.vertices).float().to(device) side_faces = torch.tensor(side_mesh.faces).long().to(device) # Possion Fusion between SMPLX and BNI # 1. keep the faces invisible to front+back cameras # 2. keep the front-FLAME+MANO faces # 3. remove eyeball faces # export intermediate meshes BNI_object.F_B_trimesh.export(f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj") full_lst = [] if "face" in cfg.bni.use_smpl: # only face face_mesh = apply_vertex_mask(face_mesh, SMPLX_object.front_flame_vertex_mask) face_mesh.vertices = face_mesh.vertices - np.array([0, 0, cfg.bni.thickness]) # remove face neighbor triangles BNI_object.F_B_trimesh = part_removal( BNI_object.F_B_trimesh, face_mesh, cfg.bni.face_thres, device, smplx_mesh, region="face" ) side_mesh = part_removal( side_mesh, face_mesh, cfg.bni.face_thres, device, smplx_mesh, region="face" ) face_mesh.export(f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_face.obj") full_lst += [face_mesh] if "hand" in cfg.bni.use_smpl and (True in data['hands_visibility'][idx]): hand_mask = torch.zeros(SMPLX_object.smplx_verts.shape[0], ) if data['hands_visibility'][idx][0]: hand_mask.index_fill_( 0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["left_hand"]), 1.0 ) if data['hands_visibility'][idx][1]: hand_mask.index_fill_( 0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["right_hand"]), 1.0 ) # only hands hand_mesh = apply_vertex_mask(hand_mesh, hand_mask) # remove hand neighbor triangles BNI_object.F_B_trimesh = part_removal( BNI_object.F_B_trimesh, hand_mesh, cfg.bni.hand_thres, device, smplx_mesh, region="hand" ) side_mesh = part_removal( side_mesh, hand_mesh, cfg.bni.hand_thres, device, smplx_mesh, region="hand" ) hand_mesh.export(f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_hand.obj") full_lst += [hand_mesh] full_lst += [BNI_object.F_B_trimesh] # initial side_mesh could be SMPLX or IF-net side_mesh = part_removal( side_mesh, sum(full_lst), 2e-2, device, smplx_mesh, region="", clean=False ) full_lst += [side_mesh] # # export intermediate meshes BNI_object.F_B_trimesh.export(f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj") side_mesh.export(f"{out_dir}/{cfg.name}/obj/{data['name']}_{idx}_side.obj") final_mesh = poisson( sum(full_lst), final_path, cfg.bni.poisson_depth, ) print( colored(f"Poisson completion to : {final_path} ", "yellow") ) dataset.render.load_meshes(final_mesh.vertices, final_mesh.faces) rotate_recon_lst = dataset.render.get_image(cam_type="four") per_loop_lst.extend([in_tensor['image'][idx:idx + 1]] + rotate_recon_lst) if cfg.bni.texture_src == 'image': # coloring the final mesh (front: RGB pixels, back: normal colors) final_colors = query_color( torch.tensor(final_mesh.vertices).float(), torch.tensor(final_mesh.faces).long(), in_tensor["image"][idx:idx + 1], device=device, ) final_mesh.visual.vertex_colors = final_colors final_mesh.export(final_path) final_mesh.export(final_path.replace(".obj", ".glb")) elif cfg.bni.texture_src == 'SD': # !TODO: add texture from Stable Diffusion pass if len(per_loop_lst) > 0: per_data_lst.append(get_optim_grid_image(per_loop_lst, None, nrow=5, type="cloth")) per_data_lst[-1].save(osp.join(out_dir, cfg.name, f"png/{data['name']}_cloth.png")) # for video rendering in_tensor["BNI_verts"].append(torch.tensor(final_mesh.vertices).float()) in_tensor["BNI_faces"].append(torch.tensor(final_mesh.faces).long()) os.makedirs(osp.join(out_dir, cfg.name, "vid"), exist_ok=True) in_tensor["uncrop_param"] = data["uncrop_param"] in_tensor["img_raw"] = data["img_raw"] torch.save(in_tensor, osp.join(out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt")) smpl_glb_path = smpl_obj_path.replace(".obj", ".glb") # smpl_npy_path = smpl_obj_path.replace(".obj", ".npy") # refine_obj_path = final_path refine_glb_path = final_path.replace(".obj", ".glb") overlap_path = img_overlap_path vis_tensor_path = osp.join(out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt") # clean all the variables for element in dir(): if 'path' not in element: del locals()[element] import gc gc.collect() torch.cuda.empty_cache() return [smpl_glb_path, refine_glb_path, overlap_path, vis_tensor_path]