# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import json import numpy as np import os import time from pathlib import Path from scipy.spatial import cKDTree import torch import torch.backends.cudnn as cudnn import torch.nn.functional as F torch.set_num_threads(8) import utils.misc as misc import skinning_models.models as models from utils.misc import NativeScalerWithGradNormCount as NativeScaler from utils.skin_data import SkinData from utils.util import save_skin_weights_to_rig, post_filter def get_args_parser(): parser = argparse.ArgumentParser('Autoencoder', add_help=False) parser.add_argument('--batch_size', default=64, type=int, help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus') parser.add_argument('--lr', default=1e-4, type=float, help='learning rate') parser.add_argument('--model', default='SkinningNetStacked', type=str, metavar='MODEL', help='Name of model to train') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=0, type=int) parser.add_argument('--num_workers', default=60, type=int) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem') parser.set_defaults(pin_mem=False) # distributed training parameters parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') parser.add_argument('--local_rank', default=-1, type=int) parser.add_argument('--dist_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') parser.add_argument('--pretrained_weights', default=None, type=str, help='dataset path') parser.add_argument('--depth', default=1, type=int, help='network depth in transformer') parser.add_argument('--max_joints', default=70, type=int, help='max joints') parser.add_argument('--use_TAJA', action='store_true', default=True, help='whether to use TAJA') parser.add_argument('--save_folder', default="outputs", type=str, help='save folder') parser.add_argument('--save_skin_npy', action='store_true', default=False, help='save skinning weights as npy files') # for evaluation parser.add_argument('--eval', action='store_true', help='Perform evaluation only') parser.add_argument('--eval_data_path', default=None, type=str, help='eval dataset path') parser.add_argument('--pose_test', action='store_true', default=False, help='evaluate on diverse pose test set') parser.add_argument('--modelres_test', action='store_true', default=False, help='evaluate on modelresources test set') parser.add_argument('--xl_test', action='store_true', default=False, help='evaluate on articulation-xl test set') parser.add_argument('--filter_thre', default=0.15, type=float, help='filter threshold') # for generation parser.add_argument('--generate', action='store_true', default=False, help='Perform inference') parser.add_argument('--input_skel_folder', default=None, type=str, help='input skeleton folder') parser.add_argument('--mesh_folder', default=None, type=str, help='input mesh folder') parser.add_argument('--post_filter', action='store_true', default=False, help='whether to do post filtering') return parser @torch.no_grad() def evaluate(data_loader, model, device, args): model.eval() prec_total = [] rec_total = [] l1_dist_total = [] infer_all_time = [] output_dir = args.save_folder os.makedirs(output_dir, exist_ok=True) eval_file = os.path.join(output_dir, 'evaluate_pose_test.txt' if args.pose_test else 'evaluate_modelres_test.txt' if args.modelres_test else 'evaluate_xl_test.txt' if args.xl_test else 'evaluate_default.txt') if args.modelres_test: args.filter_thre = 0.35 with open(eval_file, 'w') as f: def log_print(*args, **kwargs): print(*args, **kwargs) print(*args, **kwargs, file=f) for data_iter_step, (sample_points, pc_w_norm, skeleton, valid_joints_mask, dist_graph, vertices, file_name, edges, gt_skin) in enumerate(data_loader): sample_points = sample_points.to(device, non_blocking=True) pc_w_norm = pc_w_norm.to(device, non_blocking=True) skeleton = skeleton.to(device, non_blocking=True) valid_joints_mask = valid_joints_mask.to(device, non_blocking=True) dist_graph = dist_graph.to(device, non_blocking=True) edges = edges.to(device, non_blocking=True) start_time = time.time() with torch.cuda.amp.autocast(enabled=False): generate_skin = model( sample_points, skeleton, pc_w_norm, dist_graph, valid_joints_mask ) infer_time_pre_mesh = time.time() - start_time infer_all_time.append(infer_time_pre_mesh) generate_skin_np = generate_skin.cpu().numpy() # (batch_size, ...) gt_skin_np = gt_skin.cpu().numpy() valid_joints_mask_np = valid_joints_mask.cpu().numpy() # (batch_size, num_joints) batch_size = generate_skin_np.shape[0] for i in range(batch_size): tree = cKDTree(sample_points[i][:,:3].cpu().numpy()) _, indices = tree.query(vertices[i].cpu().numpy()) current_generate_skin = generate_skin_np[i][indices] # (n_vertex, n_joints) current_gt_skin = gt_skin_np[i] # (n_vertex, num_joints) current_valid_joints_mask = valid_joints_mask_np[i] # (num_joints,) valid_joint_indices = np.where(current_valid_joints_mask == 1)[0] if len(valid_joint_indices) == 0: continue generate_skin_masked = current_generate_skin[:, valid_joint_indices] gt_skin_masked = current_gt_skin[:, valid_joint_indices] if generate_skin_masked.size == 0: continue generate_skin_masked[generate_skin_masked < 1e-3] = 0.0 if args.post_filter: generate_skin_masked = post_filter(generate_skin_masked, edges[i].cpu().numpy(), num_ring=1) generate_skin_masked[generate_skin_masked < np.max(generate_skin_masked, axis=1, keepdims=True) * args.filter_thre] = 0.0 generate_skin_masked = generate_skin_masked / (generate_skin_masked.sum(axis=1, keepdims=True)+1e-10) valid_rows = np.abs(np.sum(gt_skin_masked, axis=1) - 1) < 1e-2 generate_skin_masked = generate_skin_masked[valid_rows] gt_skin_masked = gt_skin_masked[valid_rows] if args.save_skin_npy: test_folder = ('xl_test' if args.xl_test else 'pose_test' if args.pose_test else 'modelres_test' if args.modelres_test else 'default') os.makedirs(os.path.join(output_dir, test_folder), exist_ok=True) npy_path = os.path.join(output_dir, test_folder, f"{file_name[i]}_skin.npy") np.save(npy_path, generate_skin_masked) # metrics precision = np.sum(np.logical_and(generate_skin_masked > 0, gt_skin_masked > 0)) / (np.sum(generate_skin_masked > 0) + 1e-10) recall = np.sum(np.logical_and(generate_skin_masked > 0, gt_skin_masked > 0)) / (np.sum(gt_skin_masked > 0) + 1e-10) mean_l1_dist = np.sum(np.abs(generate_skin_masked - gt_skin_masked)) /len(generate_skin_masked) log_print('for', data_iter_step, ',', file_name[i], ': precision:', precision, 'recall:', recall, 'mean_l1_dist:', mean_l1_dist) prec_total.append(precision) rec_total.append(recall) l1_dist_total.append(mean_l1_dist) print("number of items: " + str(len(l1_dist_total))) final_precision = np.mean(prec_total) if prec_total else 0.0 final_recall = np.mean(rec_total) if rec_total else 0.0 final_avg_l1_dist = np.mean(l1_dist_total) if l1_dist_total else 0.0 avg_infer_time = np.mean(infer_all_time) log_print('final_precision: ', final_precision, 'final_recall: ', final_recall, 'final_avg_l1_dist: ', final_avg_l1_dist, 'avg_infer_time: ', avg_infer_time) @torch.no_grad() def generate(data_loader, model, device, args): model.eval() for data_iter_step, (sample_points, pc_w_norm, skeleton, valid_joints_mask, dist_graph, vertices, file_name, edges) in enumerate(data_loader): sample_points = sample_points.to(device, non_blocking=True) pc_w_norm = pc_w_norm.to(device, non_blocking=True) skeleton = skeleton.to(device, non_blocking=True) valid_joints_mask = valid_joints_mask.to(device, non_blocking=True) dist_graph = dist_graph.to(device, non_blocking=True) edges = edges.to(device, non_blocking=True) if skeleton[0].shape[0] > args.max_joints: continue with torch.cuda.amp.autocast(enabled=False): generate_skin = model( sample_points, skeleton, pc_w_norm, dist_graph, valid_mask=valid_joints_mask, ) generate_skin_np = generate_skin.cpu().numpy() # (batch_size, ...) valid_joints_mask_np = valid_joints_mask.cpu().numpy() # (batch_size, num_joints) batch_size = generate_skin_np.shape[0] for i in range(batch_size): tree = cKDTree(sample_points[i][:,:3].cpu().numpy()) _, indices = tree.query(vertices[i].cpu().numpy()) current_generate_skin = generate_skin_np[i][indices] # (n_vertex, n_joints) current_valid_joints_mask = valid_joints_mask_np[i] # (num_joints,) valid_joint_indices = np.where(current_valid_joints_mask == 1)[0] if len(valid_joint_indices) == 0: continue generate_skin_masked = current_generate_skin[:, valid_joint_indices] if generate_skin_masked.size == 0: continue if args.post_filter: generate_skin_masked = post_filter(generate_skin_masked, edges[i].cpu().numpy(), num_ring=1) generate_skin_masked[generate_skin_masked < np.max(generate_skin_masked, axis=1, keepdims=True) * 0.35] = 0.0 generate_skin_masked = generate_skin_masked / (generate_skin_masked.sum(axis=1, keepdims=True)) output_dir = args.save_folder if not os.path.exists(output_dir): os.makedirs(output_dir) pred_rig_path = os.path.join(args.input_skel_folder, f'{file_name[i]}.txt') print(file_name[i]) # save rig files with skinning weights os.makedirs(os.path.join(output_dir, 'generate'), exist_ok=True) output_path = os.path.join(output_dir, f'generate/{file_name[i]}_skin.txt') print(output_path) save_skin_weights_to_rig(pred_rig_path, generate_skin_masked, output_path) np.save(os.path.join(output_dir, f"generate/{file_name[i]}_skin.npy"), generate_skin_masked) def main(args): misc.init_distributed_mode(args) print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) device = torch.device(args.device) # fix the seed for reproducibility seed = args.seed + misc.get_rank() torch.manual_seed(seed) np.random.seed(seed) cudnn.benchmark = True if args.eval: dataset_val = SkinData(args, mode='eval', query_num=8192) elif args.generate: dataset_val = SkinData(args, mode='generate', query_num=8192) else: dataset_train = SkinData(args, mode='train', query_num=8192) num_tasks = misc.get_world_size() global_rank = misc.get_rank() if not args.eval and not args.generate: sampler_train = torch.utils.data.DistributedSampler( dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True ) print("Sampler_train = %s" % str(sampler_train)) else: sampler_val = torch.utils.data.SequentialSampler(dataset_val) if args.eval or args.generate: data_loader_val = torch.utils.data.DataLoader( dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False ) else: data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True, persistent_workers=True, prefetch_factor=4 ) model = models.__dict__[args.model](args) model.to(device) n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) # print("Model = %s" % str(model)) print('number of params (M): %.2f' % (n_parameters / 1.e6)) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=False) model_without_ddp = model.module optimizer = torch.optim.AdamW(model_without_ddp.parameters(), lr=args.lr) if args.pretrained_weights is not None: pkg = torch.load(args.pretrained_weights, map_location=torch.device("cpu")) model_without_ddp.load_state_dict(pkg["model"]) if args.generate: generate(data_loader_val, model_without_ddp, device, args) elif args.eval: if not any([args.xl_test, args.pose_test, args.modelres_test]): raise ValueError("Please specify a test type: --xl_test, --pose_test, or --modelres_test") test_type = ('Articulation-XL2.0 Test' if args.xl_test else 'Diverse-Pose Test' if args.pose_test else 'ModelsResource Test') print(f"Running evaluation: {test_type}") evaluate(data_loader_val, model_without_ddp, device, args) if __name__ == '__main__': args = get_args_parser() args = args.parse_args() main(args)