# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. import random import numpy as np import ray import torch import os from tqdm import tqdm from load_data.interface import LoadData def read_all_data(folder_list, load_data, add_model_str=True, add_ori_name=False): all_data = [] for f in folder_list: if add_model_str: result = load_data.run(os.path.join(f, 'model', 'mesh')) elif add_ori_name: result = load_data.run(os.path.join(f, f.split('/')[-1], 'mesh')) else: result = load_data.run(os.path.join(f, 'mesh')) all_data.append(result) q8_table = all_data[0][0] align_10 = all_data[0][1] dest_ArtCoeff = [r[2][np.newaxis, :] for r in all_data] dest_FdCoeff_q8 = [r[3][np.newaxis, :] for r in all_data] dest_CirCoeff_q8 = [r[4][np.newaxis, :] for r in all_data] dest_EccCoeff_q8 = [r[5][np.newaxis, :] for r in all_data] SRC_ANGLE = 10 ANGLE = 10 CAMNUM = 10 ART_COEF = 35 FD_COEF = 10 n_shape = len(all_data) dest_ArtCoeff = torch.from_numpy(np.ascontiguousarray(np.concatenate(dest_ArtCoeff, axis=0))).int().cuda().reshape(n_shape, SRC_ANGLE, CAMNUM, ART_COEF) dest_FdCoeff_q8 = torch.from_numpy(np.ascontiguousarray(np.concatenate(dest_FdCoeff_q8, axis=0))).int().cuda().reshape(n_shape, ANGLE, CAMNUM, FD_COEF) dest_CirCoeff_q8 = torch.from_numpy(np.ascontiguousarray(np.concatenate(dest_CirCoeff_q8, axis=0))).int().cuda().reshape(n_shape, ANGLE, CAMNUM) dest_EccCoeff_q8 = torch.from_numpy(np.ascontiguousarray(np.concatenate(dest_EccCoeff_q8, axis=0))).int().cuda().reshape(n_shape, ANGLE, CAMNUM) q8_table = torch.from_numpy(np.ascontiguousarray(q8_table)).int().cuda().reshape(256, 256) align_10 = torch.from_numpy(np.ascontiguousarray(align_10)).int().cuda().reshape(60, 20) ## return q8_table.contiguous(), align_10.contiguous(), dest_ArtCoeff.contiguous(), \ dest_FdCoeff_q8.contiguous(), dest_CirCoeff_q8.contiguous(), dest_EccCoeff_q8.contiguous() def compute_lfd_all(src_folder_list, tgt_folder_list, log): load_data = LoadData() add_ori_name = False add_model_str = False src_folder_list.sort() tgt_folder_list.sort() q8_table, align_10, src_ArtCoeff, src_FdCoeff_q8, src_CirCoeff_q8, src_EccCoeff_q8 = read_all_data(src_folder_list, load_data, add_model_str=False) q8_table, align_10, tgt_ArtCoeff, tgt_FdCoeff_q8, tgt_CirCoeff_q8, tgt_EccCoeff_q8 = read_all_data(tgt_folder_list, load_data, add_model_str=add_model_str, add_ori_name=add_ori_name) ### from lfd_all_compute.lfd import LFD lfd = LFD() lfd_matrix = lfd.forward( q8_table, align_10, src_ArtCoeff, src_FdCoeff_q8, src_CirCoeff_q8, src_EccCoeff_q8, tgt_ArtCoeff, tgt_FdCoeff_q8, tgt_CirCoeff_q8, tgt_EccCoeff_q8, log) # print(lfd_matrix) # print(lfd_matrix.shape) mmd = lfd_matrix.float().min(dim=0)[0].mean() mmd_swp = lfd_matrix.float().min(dim=1)[0].mean() # print(mmd) # print(mmd_swp) return lfd_matrix.data.cpu().numpy() if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument("--save_name", type=str, required=True, help="path to the save resules shapenet dataset") parser.add_argument("--dataset_path", type=str, required=True, help="path to the preprocessed shapenet dataset") parser.add_argument("--gen_path", type=str, required=True, help="path to the generated models") parser.add_argument("--num_workers", type=int, default=1, help="number of workers to run in parallel") parser.add_argument("--list", type=str, default=None, help="list file in the training set") args = parser.parse_args() save_path = '/'.join(args.save_name.split('/')[:-1]) os.makedirs(save_path, exist_ok=True) num_workers = args.num_workers listfile = args.list ray.init( num_cpus=os.cpu_count(), num_gpus=num_workers, ) print(f"dataset_path: {args.dataset_path}") print(f"gen_path: {args.gen_path}") assert os.path.exists(args.dataset_path) and os.path.exists(args.gen_path) tgt_folder_list = sorted(os.listdir(args.dataset_path)) if listfile is not None: valid_folders = [item.strip() for item in open(listfile, 'r').readlines()] tgt_folder_list = sorted(list(set(valid_folders) & set(tgt_folder_list))) tgt_folder_list = [os.path.join(args.dataset_path, f) for f in tgt_folder_list] else: tgt_folder_list = [os.path.join(args.dataset_path, f) for f in tgt_folder_list] src_folder_list = os.listdir(args.gen_path) random.shuffle(src_folder_list) src_folder_list = sorted(src_folder_list[:3000]) src_folder_list = [os.path.join(args.gen_path, f) for f in src_folder_list] compute_lfd_all_remote = ray.remote(num_gpus=1, num_cpus=os.cpu_count() // num_workers)(compute_lfd_all) print("Check data") print(f"len of src_folder_list: {len(src_folder_list)}") print(f"len of tgt_folder_list: {len(tgt_folder_list)}") # print(src_folder_list[0]) # print(tgt_folder_list[0]) results = [] for i in range(num_workers): i_start = i * len(src_folder_list) // num_workers i_end = (i + 1) * len(src_folder_list) // num_workers # print(i, i_start, i_end) results.append(compute_lfd_all_remote.remote( src_folder_list[i_start:i_end], tgt_folder_list, i==0)) lfd_matrix = ray.get(results) lfd_matrix = np.concatenate(lfd_matrix, axis=0) import pickle save_name = args.save_name nearest_name = [tgt_folder_list[idx].split("/")[-1] for idx in lfd_matrix.argmin(axis=1)] src_folder_list = [src_folder_list[idx].split("/")[-1] for idx in range(len(src_folder_list))] pickle.dump([src_folder_list, nearest_name, lfd_matrix], open(save_name, 'wb')) print(f"pkl is saved to {save_name}")