File size: 5,936 Bytes
8d5039c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import glob
import tqdm
import os
import torch
import numpy as np
import open3d as o3d
from scipy.optimize import linear_sum_assignment
import itertools
from model.pointnet_util import *
from model.model_utils import *
def writePoints(points, clsRoad):
with open(clsRoad, 'w+') as file1:
for i in range(len(points)):
point = points[i]
file1.write(str(point[0]))
file1.write(' ')
file1.write(str(point[1]))
file1.write(' ')
file1.write(str(point[2]))
file1.write('\n')
def writeEdges(edges, clsRoad):
with open(clsRoad, 'w+') as file1:
for i in range(len(edges)):
edge = edges[i]
file1.write(str(edge[0] + 1))
file1.write(' ')
file1.write(str(edge[1] + 1))
file1.write(' ')
file1.write('\n')
def assign_targets(points, gvs, radius):
idx = ball_center_query(radius, points, gvs).type(torch.int64)
batch_size = gvs.size()[0]
idx_add = torch.arange(batch_size).to(idx.device).unsqueeze(-1).repeat(1, idx.shape[-1]) * gvs.shape[1]
gvs = gvs.view(-1, 3)
idx_add += idx
target_points = gvs[idx_add.view(-1)].view(batch_size, -1, 3)
dis = target_points - points
dis[idx < 0] = 0
dis /= radius
label = torch.where(idx >= 0, torch.ones(idx.shape).to(idx.device), torch.zeros(idx.shape).to(idx.device))
return dis, label
def test_model(model, data_loader, logger):
dataloader_iter = iter(data_loader)
with tqdm.trange(0, len(data_loader), desc='test', dynamic_ncols=True) as tbar:
model.use_edge = True
statistics = {'tp_pts': 0, 'num_label_pts': 0, 'num_pred_pts': 0, 'pts_bias': np.zeros(3, float),
'tp_edges': 0, 'num_label_edges': 0, 'num_pred_edges': 0}
for cur_it in tbar:
batch = next(dataloader_iter)
load_data_to_gpu(batch)
with torch.no_grad():
batch = model(batch)
load_data_to_cpu(batch)
# print("test")
keys = batch['frame_id']
wf_vertices = batch['refined_keypoint']#.squeeze(0)
wf_edges = batch['pair_points']#.squeeze(0)
return keys, wf_vertices, wf_edges
# def test_model(model, data_loader, logger):
# dataloader_iter = iter(data_loader)
# with tqdm.trange(0, len(data_loader), desc='test', dynamic_ncols=True) as tbar:
# model.use_edge = True
# statistics = {'tp_pts': 0, 'num_label_pts': 0, 'num_pred_pts': 0, 'pts_bias': np.zeros(3, float),
# 'tp_edges': 0, 'num_label_edges': 0, 'num_pred_edges': 0}
# for cur_it in tbar:
# batch = next(dataloader_iter)
# load_data_to_gpu(batch)
# with torch.no_grad():
# batch = model(batch)
# load_data_to_cpu(batch)
# eval_process(batch, statistics)
# bias = statistics['pts_bias'] / statistics['tp_pts']
# logger.info('pts_recall: %f' % (statistics['tp_pts'] / statistics['num_label_pts']))
# logger.info('pts_precision: %f' % (statistics['tp_pts'] / statistics['num_pred_pts']))
# logger.info('pts_bias: %f, %f, %f' % (bias[0], bias[1], bias[2]))
# logger.info('edge_recall: %f' % (statistics['tp_edges'] / statistics['num_label_edges']))
# logger.info('edge_precision: %f' % (statistics['tp_edges'] / statistics['num_pred_edges']))
def eval_process(batch, statistics):
batch_size = batch['batch_size']
pts_pred, pts_refined, pts_label = batch['keypoint'], batch['refined_keypoint'], batch['vectors']
edge_pred, edge_label = batch['edge_score'], batch['edges']
mm_pts = batch['minMaxPt']
id = batch['frame_id']
idx = 0
for i in range(batch_size):
mm_pt = mm_pts[i]
minPt = mm_pt[0]
maxPt = mm_pt[1]
deltaPt = maxPt - minPt
p_pts = pts_refined[pts_pred[:, 0] == i]
l_pts = pts_label[i]
l_pts = l_pts[np.sum(l_pts, -1, keepdims=False) > -2e1]
vec_a = np.sum(p_pts ** 2, -1)
vec_b = np.sum(l_pts ** 2, -1)
dist_matrix = vec_a.reshape(-1, 1) + vec_b.reshape(1, -1) - 2 * np.matmul(p_pts, np.transpose(l_pts))
dist_matrix = np.sqrt(dist_matrix + 1e-6)
p_ind, l_ind = linear_sum_assignment(dist_matrix)
mask = dist_matrix[p_ind, l_ind] < 0.1 # 0.1
tp_ind, tl_ind = p_ind[mask], l_ind[mask]
#dis = np.abs(p_pts[tp_ind] - l_pts[tl_ind])
dis = np.abs( ((p_pts[tp_ind]*deltaPt) + minPt) - ((l_pts[tl_ind]*deltaPt) + minPt) )
statistics['tp_pts'] += tp_ind.shape[0]
statistics['num_label_pts'] += l_pts.shape[0]
statistics['num_pred_pts'] += p_pts.shape[0]
statistics['pts_bias'] += np.sum(dis, 0)
match_edge = list(itertools.combinations(l_ind, 2))
match_edge = np.array([tuple(sorted(e)) for e in match_edge])
score = edge_pred[idx:idx+len(match_edge)]
idx += len(match_edge)
l_edge = edge_label[i]
l_edge = l_edge[np.sum(l_edge, -1, keepdims=False) > 0]
l_edge = [tuple(e) for e in l_edge]
match_edge = match_edge[score > 0.5]
tp_edges = np.sum([tuple(e) in l_edge for e in match_edge])
statistics['tp_edges'] += tp_edges
statistics['num_label_edges'] += len(l_edge)
statistics['num_pred_edges'] += match_edge.shape[0]
def load_data_to_gpu(batch_dict):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for key, val in batch_dict.items():
if not isinstance(val, np.ndarray):
continue
batch_dict[key] = torch.from_numpy(val).float().to(device)
def load_data_to_cpu(batch_dict):
for key, val in batch_dict.items():
if not isinstance(val, torch.Tensor):
continue
batch_dict[key] = val.cpu().numpy() |