|
- import torch
- import torch.utils.data
- import torch.nn.parallel
- import os
- import numpy as np
- import scipy.io as sio
- import argparse
- import fnmatch
- import Datasets
- import models
- from pyOMT_raw import pyOMT_raw, train_omt
- from torch.autograd.variable import Variable
- from im2mesh.utils.io import export_pointcloud
-
- from torchvision.utils import save_image
- '''for evaluation part packages'''
- from eval import Eval
- from Datasets.plyfile.plyfile import PlyData
- from models.emd.emd import earth_mover_distance
- from models.lossess import ChamferLoss
-
- # for PRD score
- # import hashlib
- # import prd_score as prd
-
-
- np.random.seed(5)
- torch.manual_seed(5)
-
- dataset_names = sorted(name for name in Datasets.__all__)
- model_names = sorted(name for name in models.__all__)
-
- parser = argparse.ArgumentParser(description= 'Point Cloud Disctribution Transformation SCOT', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- #parser.add_argument('config',default='/home/rayna/datasets/shape_net_core_uniform_samples_2048/', type=str, help='Path to config file.')
-
- #arguments for Saving Models and datasets
- #parser.add_argument('--GFV_path',default='./results/GFV/shapenet/08-19-08:33/encoder_pointnet/', help='Path GFV to Checkpoints')
- #parser.add_argument('--save_path',default='./results/GFV', help='Path to Data Set')
- parser.add_argument('--save',default= True,help= 'Save Models or not ?')
- parser.add_argument('--result', default='./results/', help='output result directory') # root directory of training and generating results')
- parser.add_argument('--pretrained',default='./ckpts/AE_ckpts/shapenet/08-19-08_33/ae_pointnet,Adam,400epochs,b24,lr0.001/model_best.pth.tar',help= 'Use Pretrained Model for testing or resuming training') ## TODO
- #Path to train dataset # TODO
- parser.add_argument('-d', '--data', metavar='DIR', default='/data/rayna/datasets/shape_net_core_uniform_samples_2048/four/', help='Path to Complete Point Cloud Data Set')
- #Path to test dataset # TODO
- parser.add_argument('-dw', '--datatest', default='/data/rayna/datasets/shape_net_core_uniform_samples_2048/four/', help='Path to Complete Point Cloud Data Set')
- # data processing
- parser.add_argument('-n', '--dataName', metavar='Data Set Name', default='shapenet', choices= dataset_names)
- parser.add_argument('-ad', '--adddata', metavar='aDIR', default='', help='Additional path to dataset')
- parser.add_argument('-s','--split_value',default = [0.85,0.3333], help='Ratio of train and validation data split')
-
- # Arguments for Torch Data Loader
- parser.add_argument('-b','--batch_size', type=int, default=1, help='input batch size') #
- parser.add_argument('-w','--workers',type=int, default=8, help='Set the number of workers')
-
- # Arguments for Model Settings
- #parser.add_argument('-m','--model',default='ae_pointnet',help='Chose Your Model Here',choices=['ae_pointnet','ae_rsnet','vae_pointnet']) # TODO
- parser.add_argument('-nt','--net_name',default='auto_encoder',help='Chose The name of your network',choices=['auto_encoder','shape_completion'])
- #parser.add_argument('-md','--model_decoder',default='decoder_sonet',help='Chose Your Decoder Model Here',choices=['decoder_sonet']) # TODO
-
- # Optimizer Settings
- parser.add_argument('-op','--optim',default= 'Adam',help='Specify the Optimizer to use')
- parser.add_argument('--lr',default=0.001,help='Learning Rate for the optimizer') #
- parser.add_argument('--momentum',default=0.9,help='Momentum for the adam optimizer')
- parser.add_argument('--beta',default=0.999,help='beta for the adam optimizer')
- parser.add_argument('--milestones',default=[60,120,180,500,800],help='For learning rate scheduler, will decay learning rate by gamma after each milestone')
- parser.add_argument('--gamma',default=0.5,help='gamma for the learning rate scheduler')
- parser.add_argument('--bias_decay',default=0,help='bias decay')
- # Training Settings
- parser.add_argument('--epochs',default=400,help='Number of epochs to run')
- parser.add_argument('--start_epoch',default=0,help='Starting Epoch')
-
- # Model Parameters
- parser.add_argument('--output_fc_pc_num', type=int, default=256, help='# of fc decoder output points')
- parser.add_argument('--output_conv_pc_num', type=int, default=4096, help='# of conv decoder output points')
- parser.add_argument('--feature_num', type=int, default=1024, help='length of encoded feature')
- parser.add_argument('--activation', type=str, default='relu', help='activation function: relu, elu')
- parser.add_argument('--normalization', type=str, default='batch', help='normalization function: batch, instance')
- ##OT_part
- parser.add_argument('--action', nargs='+',type=str, default='', help=' mode of OT: extract_feature, train_OT'
- ', generate, decode_feature, decode_test')
- parser.add_argument('-a','--angle', type=float, default=0.7, help='angle threshold to generate features') #
- parser.add_argument('-dis','--dissim', type=float, default=0, help='dissimarlty') #
- # GPU Settings
- parser.add_argument('--gpu_id', type=int, default=0, help='gpu ids: e.g. 0, 1. -1 is no GPU')
-
- # Parameter Initialize
- args = parser.parse_args()
- args.device = torch.device("cuda:%d" % (args.gpu_id) if torch.cuda.is_available() else "cpu") # for selecting device for chamfer loss
- #cuda.select_device
- torch.cuda.set_device(args.gpu_id)
- print('Using GPU : ',torch.cuda.current_device())
- os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
-
- if args.action == '':
- actions = ['extract_feature', 'train_OT', 'generate', 'decode_feature', 'decode_test']
- else:
- actions = args.action
-
- def main():
- '''------------------------------------- Generating args----------------------------------------- ------------'''
- max_gen_samples = 10000 # max number of generated samples. Used to avoid out of memory error.
- angle_threshold = args.angle # angle threshold of OT generator ranging from [0,1]. See paper for details.
- rec_gen_distance = args.dissim # dis-similarity between reconstructed samples and generated samples, ranging from [0,1] with smaller meaning more similar
- args.model = args.pretrained.split('/')[5].split(',')[0]
- if args.data.split('/')[5] == 'four':
- numOfClassess = 4
- elif args.data.split('/')[5] == 'eight':
- numOfClassess = 8
- else:
- numOfClassess = 1
- nameOfClass = args.data.split('/')[5]
-
- labels = ['02691156','02933112','03001627','04379243',
- '02828884','02958343','03211117','04256520']
- #labels_one = ['airplane', 'cabinet', 'chair', 'table',
- # 'bench', 'car', 'display', 'sofa']
- """------------------------------------- Creating output directories----------------------------------------- """
- '''Get the corresponding model file name'''
- import re
- time_stamp = re.split('\/', args.pretrained)[-3]
- '''Create result directories'''
- result_root_path = args.result + time_stamp
- print('==> results dir:', result_root_path)
- if not os.path.exists(result_root_path):
- os.mkdir(result_root_path)
-
- '''1. extract_feature output dir: GFV'''
- save_path = '{}'.format(args.model) # encoder_pointnet
- save_path = os.path.join(args.dataName,save_path)
- save_path = os.path.join('GFV', save_path)
- save_path = os.path.join(result_root_path, save_path)
- print('==> GFV dir:{}', save_path)
- if not os.path.exists(save_path):
- os.makedirs(save_path)
- feature_save_path = os.path.join(save_path, 'GFV.pt')
- '''2. train_OT output dir: '''
- ot_path = os.path.join(result_root_path,'OT')
- '''3. generate output dir: h '''
- outP_path = os.path.join(ot_path, str(angle_threshold) + '_' + str(rec_gen_distance))
- h_path = os.path.join(ot_path, 'h.pt')
- print('==> generate h path{} and outP_path{}', h_path, outP_path)
- gen_feature_path = os.path.join(outP_path, 'output_P_gen')
- if not os.path.exists(ot_path):
- os.makedirs(ot_path)
- '''4. decode_feature output dir : for evaluate trainset'''
- gen_pc_pair_path = os.path.join(ot_path, str(angle_threshold)+'_'+str(rec_gen_distance), 'gen_pcs_pairs/')
- gen_pc_pair_drop_path = os.path.join(ot_path, str(angle_threshold)+'_'+str(rec_gen_distance), 'gen_pcs_pairs_drop/')
- print('==> gen_pc_pair_path:', gen_pc_pair_path)
- if not os.path.exists(gen_pc_pair_path):
- os.makedirs(gen_pc_pair_path)
- eval_train_path = os.path.join(gen_pc_pair_path[:-14], 'eval_train.csv')
- eval_train_drop_path = os.path.join(gen_pc_pair_path[:-14], 'eval_train_drop.csv')
- '''5. decode_test output dir: for evaluate testset'''
- gen_pc_path = os.path.join(ot_path, str(angle_threshold) + '_' + str(rec_gen_distance), 'gen_pcs/')
- print('==> gen_pc_path:', gen_pc_path)
- if not os.path.exists(gen_pc_path):
- os.makedirs(gen_pc_path)
- root_eval = gen_pc_path[:-8]
- eval_test_path = os.path.join(root_eval, 'eval_test.csv')
-
- """------------------------------------- Data Loader---------------------------------------------------------- """
- [train_dataset, valid_dataset, test_dataset] = Datasets.__dict__[args.dataName](input_root=args.data,
- target_root=None,
- split=args.split_value,
- net_name=args.net_name,
- input_transforms=None,
- target_transforms=None,
- co_transforms=None,
- give_name=True)
- train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size,
- num_workers=args.workers,
- shuffle=False,
- pin_memory=True)
-
- test_loader = torch.utils.data.DataLoader(test_dataset,
- batch_size=args.batch_size,
- num_workers=args.workers,
- shuffle=False,
- pin_memory=True)
- print('==> Prepare: Data load complete')
- print('train_dataset len:',len(train_loader))
-
- """------------------------------------- Model Loader---------------------------------------------------------- """
- network_data = torch.load(args.pretrained)
- if (args.model == 'ae_pointnet' or args.model == 'vae_pointnet'):
- model = models.__dict__[args.model](args, num_points=2048, global_feat=True, data=network_data).cuda()
- else:
- model = models.__dict__[args.model](network_data).cuda()
- model.eval()
- print('==> Prepare: Model load complete')
- eval = Eval()
-
- """------------------------------------- Loss Loader---------------------------------------------------------- """
- chamfer = ChamferLoss(args)
- """------------------------------------- Run Part-------------------------------------------------------------- """
- print('==> Run actions:',actions)
- for action in actions:
- print('action', action)
- if action == 'extract_feature': # notice here batch_size need to be 1
- '''1. extract train and test dataset feature, namely GFV.pt and GFV_test.pt
- 2. eval chamfer and emd distance : train_errors.csv and test_errors.csv
- 3. evaluation TSNE and UMAP of train and test latent code.
- '''
- print('-----------------------------------1. extract_feature---------------------------------------------')
- from models.VAE_pointnet import kl_divergence, reparametrize
- emd = earth_mover_distance
- c_losses = 0
- e_losses = 0
- epath = os.path.join(save_path, 'train_errors.csv')
- test_error = os.path.join(save_path, 'test_errors.csv')
- args.batch_size = 1
- epoch = 0
- j = 1
- if args.model =='ae_pointnet':
- features_train = torch.empty([len(train_loader), 128], dtype=torch.float, requires_grad=False)
- features_test = torch.empty([len(test_loader), 128], dtype=torch.float, requires_grad=False)
- print('feature_train.size:', features_train.size())
- else:
- print('model not exist:', args.model)
- class_labels_train = torch.empty(len(train_loader), dtype=torch.uint8, requires_grad=False)
- class_labels_test = torch.empty(len(test_loader), dtype=torch.uint8, requires_grad=False)
- header = ['chamfer', 'emd']
-
- ''' extract features of train dataset'''
- epoch_size = len(train_loader)
- print('in 1 stage:train_loader', epoch_size)
- if os.path.exists(epath):
- os.remove(epath)
- with open(epath, 'a') as f:
- for i, (input, input_name) in enumerate(train_loader):
- save_path_old = save_path
- input_name = input_name[0]
- save_path = os.path.join(save_path, input_name[:8])
- class_labels_train[i] = labels.index(input_name[:8])
- root_name = os.path.basename(input_name)
- if not os.path.exists(save_path):
- os.makedirs(save_path)
- save_file = os.path.join(save_path, root_name)
- with torch.no_grad():
- input = input.cuda()
- input_var = Variable(input, requires_grad=True)
- if args.model == 'ae_pointnet':
- trans_input = torch.transpose(input_var, 1, 2)
- x = torch.squeeze(input_var, dim=1)
- x = torch.transpose(x, 1, 2)
- [encoder_out, _] = model.encoder(x, )
- features_train[i] = encoder_out
- decoder_out = model.decoder(encoder_out)
- c_loss = chamfer(trans_input, decoder_out)
- e_loss = emd(trans_input, decoder_out)
- e_loss = torch.sqrt(e_loss).mean()
- #print(decoder_out.shape)
- np.save(save_file, encoder_out.cpu())
- e_losses = e_losses + e_loss
-
- c_losses = c_losses + c_loss
-
- # measured elapsed time
- # from im2mesh.utils.io import export_pointcloud
- # y_rec = decoder_out.cpu()
- # y_rec = y_rec.transpose(1, 2)
- # y_rec = np.squeeze(y_rec)
- # #print(y_rec.shape)
- # out_file = save_file + '.ply'
- # export_pointcloud(y_rec.numpy(), out_file, False)
-
- if i % 8000 == 0:
- print(
- '1. Train part Epoch: [{0}][{1}/{2}]\t c_Loss: {3}, e_loss:{3}'.format(epoch, i, epoch_size, c_loss, e_loss))
- j += 1
-
- save_path = save_path_old
- errors = [c_loss.item(), e_loss.item()]
- if i == 0: ## first time to write items of each
- np.savetxt(f, [header], delimiter=',', fmt='%s')
- np.savetxt(f, [errors], delimiter=' ', fmt='%.6f')
- np.savetxt(f, ['average'], delimiter=',', fmt='%s')
- np.savetxt(f, [c_losses/len(train_loader), e_losses/len(train_loader)], delimiter=',', fmt='%.6f')
- f.close()
-
- ''' extract features of test dataset'''
- epoch_size = len(test_loader)
- if os.path.exists(test_error):
- os.remove(test_error)
- with open(test_error, 'a') as f:
- for i, (input, input_name) in enumerate(test_loader):
- save_path_old = save_path
- input_name = input_name[0]
- save_path = os.path.join(save_path, 'test', input_name[:8])
- class_labels_test[i] = labels.index(input_name[:8])
- root_name = os.path.basename(input_name)
- if not os.path.exists(save_path):
- os.makedirs(save_path)
- save_file = os.path.join(save_path, root_name)
- with torch.no_grad():
- input = input.cuda()
- input_var = Variable(input, requires_grad=True)
-
- if args.model == 'ae_pointnet':
- trans_input = torch.transpose(input_var, 1, 2)
- x = torch.squeeze(input_var, dim=1)
- x = torch.transpose(x, 1, 2)
- [encoder_out, _] = model.encoder(x, )
- features_train[i] = encoder_out
- decoder_out = model.decoder(encoder_out)
- c_loss = chamfer(trans_input, decoder_out)
- e_loss = emd(trans_input, decoder_out)
- e_loss = torch.sqrt(e_loss).mean()
-
- np.save(save_file, encoder_out.cpu())
-
- e_losses = e_losses + e_loss
- c_losses = c_losses + c_loss
- if i % 8000 == 0:
- print(
- '1. Test part:Epoch: [{0}][{1}/{2}]\t c_Loss: {3}, e_loss:{3}'.format(epoch, i, epoch_size, c_loss, e_loss))
- j += 1
-
- save_path = save_path_old
- errors = [c_loss.item(), e_loss.item()]
- if i == 0: ## first time to write items of each
- np.savetxt(f, [header], delimiter=',', fmt='%s')
- np.savetxt(f, [errors], delimiter=' ', fmt='%.6f')
- np.savetxt(f, ['average'], delimiter=',', fmt='%s')
- np.savetxt(f, [c_losses/len(test_loader), e_losses/len(test_loader)], delimiter=',', fmt='%.6f')
- f.close()
- print('1. Average emd Loss :{}'.format(torch.mean(e_losses)))
- '''save features of train and test dataset'''
- feature_train_save_path = os.path.join(save_path, 'GFV.pt')
- torch.save(features_train, feature_train_save_path)
- #torch.save(features_train, feature_train_save_path, _use_new_zipfile_serialization=False)
- feature_test_save_path = os.path.join(save_path, 'GFV_test.pt')
- torch.save(features_test, feature_test_save_path)
- #torch.save(features_test, feature_test_save_path, _use_new_zipfile_serialization=False)
- print('1. NO of GFV features:', str(features_train.size()) + ' Save to:' + feature_train_save_path)
- print('1. NO of GFV features of test:', str(features_test.size()) + ' Save to:' + feature_test_save_path)
-
- '''save TSNE and UMAP of train and test dataset in latent space'''
- if numOfClassess >1:
- eval.Vis_latent_code(numOfClassess,feature_train_save_path, save_path, class_labels_train, flag='train')
- eval.Vis_latent_code(numOfClassess, feature_test_save_path, save_path,class_labels_test, flag='test')
- else: # deal with only one class
- eval.Vis_latent_code_one(numOfClassess, feature_train_save_path, save_path, nameOfClass, flag='train')
- eval.Vis_latent_code_one(numOfClassess, feature_test_save_path, save_path, nameOfClass, flag='test')
- print('1. TSNE and UMAP of latent code complete of train and test dataset')
- if action == 'train_OT':
- '''
- input: features GFV.pt
- output: OT ckpts
- '''
- print('-----------------------------------2. Train OT part ----------------------------------------------')
- print('2. GFV path: ' + feature_save_path)
- print('2. OT model path: ' + h_path)
- print('2. Gen_feature path: ' + gen_feature_path)
- compute_ot(args, feature_save_path, h_path, gen_feature_path, mode=action, thresh=angle_threshold ,
- dissim=rec_gen_distance, max_gen_samples=max_gen_samples, time=time_stamp)
- if action == 'generate':
- '''
- input: trained_OT
- output: h.pt
- '''
- print('-----------------------------------3. generate_OT map---------------------------------------------')
- ot_model_load_path = h_path
- if not os.path.exists(ot_path):
- for file in os.listdir('.h/'):
- if fnmatch.fnmatch(file, '*.pt'):
- ot_model_load_path = os.path.join('.h/',file)
- print('3. Successfully loaded OT model' + ot_model_load_path)
- print('3. Generating features with OT solver...')
- compute_ot(args, feature_save_path, ot_model_load_path, gen_feature_path, mode=action, thresh=angle_threshold,
- dissim=rec_gen_distance, max_gen_samples=max_gen_samples)
- torch.cuda.empty_cache()
-
- if action == 'decode_feature':
- '''
- input: h.pt
- output: generated point cloud and its evaluations in eval*.csv
- '''
- print('-----------------------------------4. decode features---------------------------------------------')
- ## evaluation metrics:
- ## Begin: Added by ZRN 20201222
- header=['kl1_g','kl2_g','js_g','cham_g','emd_g','f_g', 'sink_g',
- 'kl1_r','kl2_r','js_r','cham_r','emd_r','f_r','sink_r']
- std = np.sqrt(0.10 ** 2)
- th_list = [std, 2 * std, 3 * std]
- sum_mode_g = 0
- sum_mode_r = 0
- sum_quality_g = 0
- sum_quality_r = 0
- sum_kl_g = []
- sum_kl_r = []
- sum_kl2_g = []
- sum_kl2_r = []
- sum_emd_g = []
- sum_emd_r = []
- sum_cham_g = []
- sum_cham_r = []
- sum_f_g = []
- sum_f_r = []
- sum_js_g = []
- sum_js_r = []
- sum_Ndiv_g = []
- sum_Ndiv_r = []
- sum_sink_g = []
- sum_sink_r = []
-
- feature_dict = sio.loadmat(gen_feature_path) # here is generated new feature samples
- features = feature_dict['features']
- ids = feature_dict['ids']
- #print(ids[0,1])
- num_feature = features.shape[0] # 1000
- #print('feature_num:', num_feature)
- num_ids = ids.size
- z = torch.from_numpy(features).cuda() # 1000x128
- #print(z.size())
- #print(feature_save_path)
- features_rec = torch.load(feature_save_path)
- #print(features_rec.size())
- z_rec = features_rec.cuda()
- #print(z_rec.size())
-
- ###################################################################################################
- ####################PRD part#######################################################################
- # print('==> PRD part:computing PRD')
- # real_embeddings = z_rec
- # eval_embeddings = z[0:17737]
- # prd_data = []
- # num_clusters = 20
- # num_angles = 1001
- # num_runs = 10
- # prd_data.append(prd.compute_prd_from_embedding(
- # eval_data=eval_embeddings.cpu().numpy(),
- # ref_data=real_embeddings.cpu().numpy(),
- # num_clusters=num_clusters,
- # num_angles=num_angles,
- # num_runs=num_runs))
- # print('==> PRD part:plotting results')
- # f_beta_data = [prd.prd_to_max_f_beta_pair(precision, recall, beta=8)
- # for precision, recall in prd_data]
- # print('==> PRD part: F_8 F_1/8 model')
-
- # prd.plot(prd_data, out_path=args.result)
-
- ###################################################################################################
- with torch.no_grad():
-
- if args.model == 'ae_pointnet':
- y = model.decoder(z) # y_size [1000,3,2048]
- y_re = model.decoder(z_rec)
- else:
- y, _, _ = model.decoder(z)
- y_re, _, _ = model.decoder(z_rec)
-
- '''Load GT data method 1'''
- lines=[]
- for i, (input, input_name) in enumerate(train_loader):
- '''input_name is a list with one element, so input_name[0] means to exact the element'''
- name_temp = str(input_name[0]) + '.ply'
- lines.append(name_temp)
-
- '''record those droped features'''
- numOfTrain = len(lines)
- drop_ids = []
- #print('ids[0,:]', ids[0,:])
- for i in range(numOfTrain):
- if i not in ids[0,:]:
- drop_ids.append(i)
- #print('drop_ids:', drop_ids, 'len:', len(drop_ids))
-
- # =====================generate reconstructed-generated point cloud pairs===========
- if os.path.exists(eval_train_path):
- #print('file:',eval_train_path,'exist, remove it')
- os.remove(eval_train_path)
- with open(eval_train_path, 'a') as f:
- for i in range(num_ids):
-
- source_name = lines[ids[0, i]]
- #print('source name:', source_name)
- source_path = os.path.join(args.data, source_name) #/n
- #print(source_path)
- classname = source_name.split('/') # 04379246/xxxxxxx.ply
- class_path = os.path.join(gen_pc_pair_path, ''.join(classname[0]))
- #print(class_path)
- #print(class_path)
- if not os.path.exists(class_path):
- os.makedirs(class_path)
- #root_name = os.path.basename(source_name)[:-4] # .ply
- root_name = classname[1][:-4] # .ply
- org_path = os.path.join(class_path, root_name + '_ori.ply')
- rec_path = os.path.join(class_path, root_name + '_rec.ply')
- gen_path = os.path.join(class_path, root_name + '_gen.ply')
-
- #export_pointcloud(input[j].cpu().numpy(), org_path, False) # here input mush change
- from shutil import copyfile
- copyfile(source_path, org_path)
-
- ply_data = PlyData.read(source_path)
- points = ply_data['vertex']
- points = np.vstack([points['x'], points['y'], points['z']]).T
- y_org = torch.from_numpy(points).cuda()
-
- #y_rec = y[i + num_ids, :, :]
- y_rec = y_re[ids[0,i], :, :]
- y_rec = y_rec.transpose(0,1)
- #print(y_rec.size())
- export_pointcloud(y_rec.cpu().numpy(), rec_path, False)
- y_gen = y[i, :, :]
- y_gen = y_gen.transpose(0,1)
- export_pointcloud(y_gen.cpu().numpy(), gen_path, False)
-
-
-
- ###################################eval part############################################
- '''KL diverserce'''
- kl1_g = eval.KL(y_org, y_gen)
- kl2_g = eval.KL(y_gen, y_org)
- kl1_r = eval.KL(y_org, y_rec)
- kl2_r = eval.KL(y_rec, y_org)
- #print('gen_kl1:, k2:', kl1_g, kl2_g)
- #print('rec_kl1:, k2:', kl1_r, kl2_r)
-
- '''JS distance'''
- js_g = eval.JS_divergence(y_org, y_gen)
- js_r = eval.JS_divergence(y_org, y_rec)
- #print('g:gen_JS:', js_g)
- #print('r:gen_JS:', js_r)
-
- '''Chamfer distance'''
- '''here need to change data to [1, 3, 2048]'''
- # print('y_org,size', y_org.size()) # 2048, 3
- trans_org = torch.transpose(y_org, 0, 1)
- trans_org = torch.unsqueeze(trans_org, 0)
-
- #trans_gen = torch.squeeze(y_gen, dim=1)
- trans_gen = torch.transpose(y_gen, 0,1)
- trans_gen = torch.unsqueeze(trans_gen, 0)
-
- #trans_rec = torch.squeeze(y_rec, dim=1)
- trans_rec = torch.transpose(y_rec, 0,1)
- trans_rec = torch.unsqueeze(trans_rec, 0)
- c_g = chamfer(trans_org, trans_gen)
- c_r = chamfer(trans_org, trans_rec)
-
- '''EMD distance'''
- ##here the input need to be put intp cuda
- emd_g = eval.emd(y_org, y_gen, transpose=False)
- emd_r = eval.emd(y_org, y_rec, transpose=False)
-
- '''F-score'''
- f_g = eval.f_score_coverage(y_org.cpu().numpy(), y_gen.cpu().numpy())
- f_r = eval.f_score_coverage(y_org.cpu().numpy(), y_rec.cpu().numpy())
-
- #
- # ## NDIV loss
- # ndiv_g = eval.NDiv_loss(z, y_gen)
- # print('gen_NDiv loss:', ndiv_g)
- # ndiv_r = eval.NDiv_loss(z_rec, y_rec)
- # print('rec_NDiv loss:', ndiv_r)
-
- '''Sinkhorn loss'''
- '''when blur =0 equals to Wass distance'''
- sinkhorn_g = eval.Sinkhorn(y_gen, y_org)
- sinkhorn_r = eval.Sinkhorn(y_rec, y_org)
- #print('gen_SinkHorn:', sinkhorn_g)
- #print('rec_SinkHorn:', sinkhorn_r)
- #print('Decoding {}/{}...'.format(i, num_ids))
-
-
- sum_kl_g.append(kl1_g.cpu().numpy())
- sum_kl_r.append(kl1_r.cpu().numpy())
- sum_kl2_g.append(kl2_g.cpu().numpy())
- sum_kl2_r.append(kl2_r.cpu().numpy())
- sum_js_g.append(js_g.cpu().numpy())
- sum_js_r.append(js_r.cpu().numpy())
- # sum_Ndiv_g = sum_Ndiv_g + ndiv_g
- # sum_Ndiv_r = sum_Ndiv_r + ndiv_r
- sum_emd_g.append(emd_g.cpu().numpy())
- sum_emd_r.append(emd_r.cpu().numpy())
- sum_cham_g.append(c_g.cpu().numpy())
- sum_cham_r.append(c_r.cpu().numpy())
- sum_f_g.append(f_g)
- sum_f_r.append(f_r)
- sum_sink_g.append(sinkhorn_g.cpu().numpy())
- sum_sink_r.append(sinkhorn_r.cpu().numpy())
- #print('mean kl', np.mean(sum_kl_g), np.mean(sum_kl2_g), np.mean(sum_js_g))
- #print('mean tensor', torch.mean(sum_emd_g).cpu().numpy(), torch.mean(sum_Ndiv_g).cpu().numpy(),torch.mean(sum_sink_g).cpu().numpy())
-
- eval_all = [kl1_g.cpu().numpy(), kl2_g.cpu().numpy(), js_g.cpu().numpy(), c_g.cpu().numpy(),
- emd_g.cpu().numpy()[0], f_g, sinkhorn_g.cpu().numpy(),
- kl1_r.cpu().numpy(), kl2_r.cpu().numpy(), js_r.cpu().numpy(), c_r.cpu().numpy(),
- emd_r.cpu().numpy()[0], f_r,sinkhorn_r.cpu().numpy()]
-
- if i == 0: ## first time to write items of each
- np.savetxt(f, [ids[0,:]], delimiter=',', fmt='%u')
- np.savetxt(f, [header], delimiter=',', fmt='%s')
- np.savetxt(f, [eval_all], delimiter=',', fmt='%.8f')
- print('4. Evaluations total average:')
- '''notice here sum_kl_g should be a list, not a number'''
- eval_sum = [np.mean(sum_kl_g), np.mean(sum_kl2_g), np.mean(sum_js_g),
- np.mean(sum_cham_g),np.mean(sum_emd_g), np.mean(sum_f_g), np.mean(sum_sink_g),
- np.mean(sum_kl_r), np.mean(sum_kl2_r), np.mean(sum_js_r),
- np.mean(sum_cham_r),np.mean(sum_emd_r), np.mean(sum_f_r), np.mean(sum_sink_r)]
-
- np.savetxt(f, ['average'], delimiter=',', fmt='%s')
- np.savetxt(f, [eval_sum], delimiter=',', fmt='%.8f')
- f.close()
-
- sum_kl_r = []
- sum_kl2_r = []
- sum_emd_r = []
- sum_cham_r = []
- sum_f_r = []
- sum_js_r = []
- sum_sink_r = []
- header = ['kl1_r', 'kl2_r', 'js_r', 'cham_r', 'emd_r', 'f_r','sink_r']
- if os.path.exists(eval_train_drop_path):
- #print('file:',eval_train_path,'exist, remove it')
- os.remove(eval_train_drop_path)
- with open(eval_train_drop_path, 'a') as f:
- for i in range(len(drop_ids)):
- source_name = lines[drop_ids[i]]
- #print(source_name)
- #print(ids[0,i])
- #print(source_name)
- source_path = os.path.join(args.data, source_name) #/n
- #print(source_path)
- classname = source_name.split('/') # 04379246/xxxxxxx.ply
- class_path = os.path.join(gen_pc_pair_drop_path, ''.join(classname[0]))
- #print(class_path)
- #print(class_path)
- if not os.path.exists(class_path):
- os.makedirs(class_path)
- #root_name = os.path.basename(source_name)[:-4] # .ply
- root_name = classname[1][:-4] # .ply
- org_path = os.path.join(class_path, root_name + '_ori.ply')
- rec_path = os.path.join(class_path, root_name + '_rec.ply')
-
- #export_pointcloud(input[j].cpu().numpy(), org_path, False) # here input mush change
- from shutil import copyfile
- copyfile(source_path, org_path)
-
- ply_data = PlyData.read(source_path)
- points = ply_data['vertex']
- points = np.vstack([points['x'], points['y'], points['z']]).T
- y_org = torch.from_numpy(points).cuda()
-
- y_rec = y_re[drop_ids[i], :, :]
- y_rec = y_rec.transpose(0,1)
- export_pointcloud(y_rec.cpu().numpy(), rec_path, False)
- ###################################eval part############################################
- #print('4. Evaluation on dropped ones: ', org_path)
-
- '''KL diverserce'''
- kl1_r = eval.KL(y_org, y_rec)
- kl2_r = eval.KL(y_rec, y_org)
-
- '''JS distance'''
- js_r = eval.JS_divergence(y_org, y_rec)
-
- '''Chamfer distance'''
- trans_org = torch.transpose(y_org, 0, 1)
- trans_org = torch.unsqueeze(trans_org, 0)
-
- trans_rec = torch.transpose(y_rec, 0,1)
- trans_rec = torch.unsqueeze(trans_rec, 0)
- c_r = chamfer(trans_org, trans_rec)
-
- '''EMD distance'''
- emd_r = eval.emd(y_org, y_rec, transpose=False)
-
- '''F-score'''
- f_r = eval.f_score_coverage(y_org.cpu().numpy(), y_rec.cpu().numpy())
-
- '''Sinkhorn loss'''
- sinkhorn_r = eval.Sinkhorn(y_rec, y_org)
-
- sum_kl_r.append(kl1_r.cpu().numpy())
- sum_kl2_r.append(kl2_r.cpu().numpy())
- sum_js_r.append(js_r.cpu().numpy())
- sum_emd_r.append(emd_r.cpu().numpy())
- sum_f_r.append(f_r)
- sum_cham_r.append(c_r.cpu().numpy())
- sum_sink_r.append(sinkhorn_r.cpu().numpy())
-
- eval_all = [kl1_r.cpu().numpy(), kl2_r.cpu().numpy(), js_r.cpu().numpy(), c_r.cpu().numpy(),
- emd_r.cpu().numpy()[0], f_r, sinkhorn_r.cpu().numpy()]
- if i == 0: ## first time to write items of each
- np.savetxt(f, [drop_ids], delimiter=',', fmt='%u')
- np.savetxt(f, [header], delimiter=',', fmt='%s')
- np.savetxt(f, [eval_all], delimiter=',', fmt='%.8f')
- print('4. Evaluations total average:')
- '''notice here sum_kl_g should be a list, not a number'''
- eval_sum = [np.mean(sum_kl_r), np.mean(sum_kl2_r), np.mean(sum_js_r),
- np.mean(sum_cham_r),np.mean(sum_emd_r), np.mean(f_r), np.mean(sum_sink_r)]
- np.savetxt(f, ['average'], delimiter=',', fmt='%s')
- np.savetxt(f, [eval_sum], delimiter=',', fmt='%.8f')
-
- f.close()
- if action== 'decode_test':
- '''
- input: features
- output: random generate point cloud as test sets.
- '''
- print('-----------------------------------5. decode generate test---------------------------------------------')
- '''generate random point cloud from OT'''
- feature_dict = sio.loadmat(gen_feature_path) # here is generated new feature samples
- features = feature_dict['features']
- max_gen_samples = len(test_loader)
- y_all = torch.empty([max_gen_samples, 2048, 3]) # what's meaning? 64x64 3 channels ?
- num_bat_y = features.shape[0] // args.batch_size
- features = features[:num_bat_y * args.batch_size, :]
- numOfgen = min(num_bat_y, max_gen_samples)
- z_features = torch.empty([numOfgen, 128], dtype=torch.float, requires_grad=False)
- #priint(num_bat_y) # 41
- #print(len(features)) # 984 = 41*24
- count = 0
- if os.path.exists(eval_test_path):
- #print('file:',eval_train_path,'exist, remove it')
- os.remove(eval_test_path)
- with open(eval_test_path, 'a') as f:
- for i in range(numOfgen):
- z = torch.from_numpy(features[i * args.batch_size: (i + 1) * args.batch_size, :]).cuda()
- z = z.view(args.batch_size, -1) # [1,128]
- z_features[i] = z
- #print(z)
- if args.model =='ae_pointnet':
- y = model.decoder(z) # z: 24x128
- else:
- y, _,_ = model.decoder(z)
-
- y = torch.transpose(y,1,2)
- print('5. Decoding {}/{}...'.format(i * args.batch_size, features.shape[0]))
- for ii in range(args.batch_size):
- export_pointcloud(y[ii,:,:].cpu().detach().numpy(), os.path.join(gen_pc_path,'gen_ply_{0:03d}.ply'.format(count)), False)
- if count < 64:
- y_all[count] = y[ii, :,:].cpu()
- count += 1
- '''here need to change, to save evaluation metrics'''
- # np.savetxt(f, [eval_sum], delimiter=',', fmt='%.6f')
- f.close()
- name_g = 'z_generated.png'
- print('root eval:', root_eval)
- feature_gen_path = os.path.join(root_eval, 'GFV_gen.pt') # new generated samples
- torch.save(z_features, feature_gen_path)
- eval.Vis_latent_code_z(numOfClassess,z_features,root_eval,name_g)
-
- print('Decoding complete.')
-
- def gen_P(p_s, numX, output_P_gen, thresh=0.7, topk=5, dissim=0.75, max_gen_samples=None): # input : map p_s; numToGenerate:numX; threshold:dissim
- '''construct topk*numOfGenerated tensors'''
- I_all = -torch.ones([topk, numX], dtype=torch.long) # set in comput_ot :20
- #print('numX:', numX) # 50000
- #print('I_all size:',I_all.size())
- num_bat_x = numX // p_s.bat_size_n
- bat_size_x = min(numX, p_s.bat_size_n)
- for ii in range(max(num_bat_x, 1)):
- p_s.pre_cal(ii)
- p_s.cal_measure()
- _, I = torch.topk(p_s.d_U, topk, dim=0) # I==indices of topk of ps.d_U
- for k in range(topk):
- I_all[k, ii * bat_size_x:(ii + 1) * bat_size_x].copy_(I[k, 0:bat_size_x])
- #print('p_s.d_U:',p_s.d_U)
- #print('U_size:', p_s.d_U.size()) #absolute each trained feature have one h
- I_all_2 = -torch.ones([2, (topk - 1) * numX], dtype=torch.long) #
- for ii in range(topk - 1):
- I_all_2[0, ii * numX:(ii + 1) * numX] = I_all[0, :] # topest
- I_all_2[1, ii * numX:(ii + 1) * numX] = I_all[ii + 1, :] #
- I_all = I_all_2
- #print('I_all:', I_all[0,:])
-
- if torch.sum(I_all < 0) > 0:
- print('3. Error: numX is not a multiple of bat_size_n')
-
- '''compute angles'''
- P = p_s.h_P
- nm = torch.cat([P, -torch.ones(p_s.num_P, 1)], dim=1) # concatnate with D1.
- nm /= torch.norm(nm, dim=1).view(-1, 1) #p=1 normalization
- cs = torch.sum(nm[I_all[0, :], :] * nm[I_all[1, :], :], 1) # element-wise multiplication, return to a list
- cs = torch.min(torch.ones([cs.shape[0]]), cs)
- theta = torch.acos(cs)
-
- ''' To modify thresh value based on the data not specified'''
- max_theta = torch.max(theta)
- min_theta = torch.min(theta)
- thresh = min_theta + thresh*(max_theta-min_theta)
-
- '''filter out generated samples with theta larger than threshold'''
-
- I_gen = I_all[:, theta <= thresh]
-
- I_gen, _ = torch.sort(I_gen, dim=0)
- _, uni_gen_id = np.unique(I_gen[0, :].numpy(), return_index=True)
- np.random.shuffle(uni_gen_id)
- I_gen = I_gen[:, torch.from_numpy(uni_gen_id)]
- # pdb.set_trace()
-
- numGen = I_gen.shape[1]
- if max_gen_samples is not None:
- numGen = min(numGen, max_gen_samples)
- I_gen = I_gen[:, :numGen]
- print('3. OT successfully generated {} samples'.format(numGen))
-
- '''generate new features'''
- if dissim == 0:
- rand_w = torch.rand([numGen,1])
- rand_w = torch.div(rand_w, torch.sum(rand_w))
- #print('sum of rand_w', torch.sum(rand_w))
- #rand_w = torch.div((rand_w - torch.min(rand_w)),(torch.max(rand_w) - torch.min(rand_w)))
- else:
- rand_w = dissim * torch.ones([numGen, 1])
-
- P_gen = (torch.mul(P[I_gen[0, :], :], 1-rand_w) + torch.mul(P[I_gen[1, :], :], rand_w)).numpy()
- P_gen2 = P[I_gen[0, :], :]
- P_gen = np.concatenate((P_gen, P_gen2))
- id_gen = I_gen[0, :].squeeze().numpy().astype(int)
- output_P_gen = output_P_gen + '.mat'
- sio.savemat(output_P_gen, {'features': P_gen, 'ids': id_gen})
-
- def compute_ot(args, input_P, output_h, output_P_gen, mode='train_OT', thresh=0.7, topk=20, dissim=0.75, max_gen_samples=None, time='20201120'):
- '''args for omt'''
- TRAIN = False
- GENERATE = False
- print('mode:', mode)
- if mode=='train_OT':
- TRAIN = True
- elif mode=='generate':
- GENERATE = True
- else:
- print('2. unrecogonized OT computation action: ' , mode)
-
- h_P = torch.load(input_P)
- num_P = h_P.shape[0]
- dim_y = h_P.shape[1]
- maxIter = 20000
- lr = 5e-1
- bat_size_P = num_P
- bat_size_n = args.batch_size
- init_num_bat_n = 20
- if not TRAIN:
- maxIter = 0
- '''args for generation'''
- num_gen_x = 50000 #a multiple of bat_size_n
-
- #crop h_P to fit bat_size_P
- h_P = h_P[0:num_P//bat_size_P*bat_size_P,:]
- num_P = h_P.shape[0]
-
- p_s = pyOMT_raw(h_P, num_P, dim_y, maxIter, lr, bat_size_P, bat_size_n)
- '''train omt'''
- if TRAIN: # or change to save several steps
- #train_omt(p_s, init_num_bat_n,time,thresh=thresh,dissim=dissim)
- train_omt(p_s, init_num_bat_n, time)
- torch.save(p_s.d_h, output_h)
- '''record result'''
- output_h_csv = output_h[:-2]+'csv'
- np.savetxt(output_h_csv, p_s.d_h.cpu().numpy(), delimiter=',')
- else:
- p_s.set_h(torch.load(output_h))
-
- if GENERATE:
- '''generate new samples'''
- gen_P(p_s, num_gen_x, output_P_gen, thresh=thresh, topk=topk, dissim=dissim, max_gen_samples=max_gen_samples)
-
- if __name__=='__main__':
- main()
|