import open3d as o3d import argparse import os import sys import numpy import numpy as np import torch import torch.utils.data from torch.utils.data import DataLoader from tqdm import tqdm # Only if the files are in example folder. BASE_DIR = os.path.dirname(os.path.abspath(__file__)) if BASE_DIR[-8:] == 'examples': sys.path.append(os.path.join(BASE_DIR, os.pardir)) os.chdir(os.path.join(BASE_DIR, os.pardir)) from learning3d.models import MaskNet2 from learning3d.data_utils import RegistrationData, ModelNet40Data def pc2open3d(data): if torch.is_tensor(data): data = data.detach().cpu().numpy() if len(data.shape) == 2: pc = o3d.geometry.PointCloud() pc.points = o3d.utility.Vector3dVector(data) return pc else: print("Error in the shape of data given to Open3D!, Shape is ", data.shape) def display_results(template, source, masked_template, masked_source): template = pc2open3d(template) source = pc2open3d(source) masked_template = pc2open3d(masked_template) masked_source = pc2open3d(masked_source) template.paint_uniform_color([1, 0, 0]) source.paint_uniform_color([0, 1, 0]) # masked_template.paint_uniform_color([0, 0, 1]) masked_template.paint_uniform_color([1, 0, 0]) masked_source.paint_uniform_color([0, 1, 0]) o3d.visualization.draw_geometries([template, source]) o3d.visualization.draw_geometries([masked_template, masked_source]) def evaluate_metrics(TP, FP, FN, TN, gt_mask): # TP, FP, FN, TN: True +ve, False +ve, False -ve, True -ve # gt_mask: Ground Truth mask [Nt, 1] accuracy = (TP + TN)/gt_mask.shape[1] misclassification_rate = (FN + FP)/gt_mask.shape[1] # Precision: (What portion of positive identifications are actually correct?) precision = TP / (TP + FP) # Recall: (What portion of actual positives are identified correctly?) recall = TP / (TP + FN) fscore = (2*precision*recall) / (precision + recall) return accuracy, precision, recall, fscore # Function used to evaluate the predicted mask with ground truth mask. def evaluate_mask(gt_mask, predicted_mask, predicted_mask_idx): # gt_mask: Ground Truth Mask [Nt, 1] # predicted_mask: Mask predicted by network [Nt, 1] # predicted_mask_idx: Point indices chosen by network [Ns, 1] if torch.is_tensor(gt_mask): gt_mask = gt_mask.detach().cpu().numpy() if torch.is_tensor(gt_mask): predicted_mask = predicted_mask.detach().cpu().numpy() if torch.is_tensor(predicted_mask_idx): predicted_mask_idx = predicted_mask_idx.detach().cpu().numpy() gt_mask, predicted_mask, predicted_mask_idx = gt_mask.reshape(1,-1), predicted_mask.reshape(1,-1), predicted_mask_idx.reshape(1,-1) gt_idx = np.where(gt_mask == 1)[1].reshape(1,-1) # Find indices of points which are actually in source. # TP + FP = number of source points. TP = np.intersect1d(predicted_mask_idx[0], gt_idx[0]).shape[0] # is inliner and predicted as inlier (True Positive) (Find common indices in predicted_mask_idx, gt_idx) FP = len([x for x in predicted_mask_idx[0] if x not in gt_idx]) # isn't inlier but predicted as inlier (False Positive) FN = FP # is inlier but predicted as outlier (False Negative) (due to binary classification) TN = gt_mask.shape[1] - gt_idx.shape[1] - FN # is outlier and predicted as outlier (True Negative) return evaluate_metrics(TP, FP, FN, TN, gt_mask) def test_one_epoch(args, model, test_loader): model.eval() test_loss = 0.0 pred = 0.0 count = 0 for i, data in enumerate(tqdm(test_loader)): template, source, igt, gt_template_mask, gt_source_mask = data template = template.to(args.device) source = source.to(args.device) igt = igt.to(args.device) # [source] = [igt]*[template] gt_template_mask = gt_template_mask.to(args.device) gt_source_mask = gt_source_mask.to(args.device) masked_template, masked_source, template_mask, source_mask = model(template, source) # TODO: Implement evaluation strategy. ''' Evaluate mask based on classification metrics. accuracy, precision, recall, fscore = evaluate_mask(gt_template_mask, template_mask, predicted_mask_idx = model.mask_idx) precision_list.append(precision) ''' # Different ways to visualize results. display_results(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], masked_template.detach().cpu().numpy()[0], masked_source.detach().cpu().numpy()[0]) def test(args, model, test_loader): test_one_epoch(args, model, test_loader) def options(): parser = argparse.ArgumentParser(description='MaskNet: A Fully-Convolutional Network For Inlier Estimation (Testing)') # settings for input data parser.add_argument('--num_points', default=1024, type=int, metavar='N', help='points in point-cloud (default: 1024)') parser.add_argument('--partial_source', default=True, type=bool, help='create partial source point cloud in dataset.') parser.add_argument('--partial_template', default=True, type=bool, help='create partial source point cloud in dataset.') parser.add_argument('--noise', default=False, type=bool, help='Add noise in source point clouds.') parser.add_argument('--outliers', default=False, type=bool, help='Add outliers to template point cloud.') # settings for on testing parser.add_argument('-j', '--workers', default=1, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument('-b', '--test_batch_size', default=1, type=int, metavar='N', help='test-mini-batch size (default: 1)') parser.add_argument('--pretrained', default='learning3d/pretrained/exp_masknet2/models/best_model_0.7.t7', type=str, metavar='PATH', help='path to pretrained model file (default: null (no-use))') parser.add_argument('--device', default='cuda:0', type=str, metavar='DEVICE', help='use CUDA if available') parser.add_argument('--unseen', default=False, type=bool, help='Use first 20 categories for training and last 20 for testing') args = parser.parse_args() return args def main(): args = options() torch.backends.cudnn.deterministic = True testset = RegistrationData('PointNetLK', ModelNet40Data(train=False, num_points=args.num_points), partial_template=args.partial_template, partial_source=args.partial_source, noise=args.noise, additional_params={'use_masknet': True, 'partial_point_cloud_method': 'planar_crop'}) test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers) if not torch.cuda.is_available(): args.device = 'cpu' args.device = torch.device(args.device) # Load Pretrained MaskNet. model = MaskNet2() if args.pretrained: assert os.path.isfile(args.pretrained) model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) model = model.to(args.device) test(args, model, test_loader) if __name__ == '__main__': main()