forked from ywcmaike/PSGN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
loss_psgn.py
143 lines (108 loc) · 4.4 KB
/
loss_psgn.py
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 torch
from utils.libkdtree import KDTree
import numpy as np
import os
from plyfile import PlyElement, PlyData
import numpy as np
def compute_iou(occ1, occ2):
''' Computes the Intersection over Union (IoU) value for two sets of
occupancy values.
Args:
occ1 (tensor): first set of occupancy values
occ2 (tensor): second set of occupancy values
'''
occ1 = np.asarray(occ1)
occ2 = np.asarray(occ2)
# Put all data in second dimension
# Also works for 1-dimensional data
if occ1.ndim >= 2:
occ1 = occ1.reshape(occ1.shape[0], -1)
if occ2.ndim >= 2:
occ2 = occ2.reshape(occ2.shape[0], -1)
# Convert to boolean values
occ1 = (occ1 >= 0.5)
occ2 = (occ2 >= 0.5)
# Compute IOU
area_union = (occ1 | occ2).astype(np.float32).sum(axis=-1)
area_intersect = (occ1 & occ2).astype(np.float32).sum(axis=-1)
iou = (area_intersect / area_union)
return iou
def chamfer_distance(points1, points2, use_kdtree=True, give_id=False):
''' Returns the chamfer distance for the sets of points.
Args:
points1 (numpy array): first point set
points2 (numpy array): second point set
use_kdtree (bool): whether to use a kdtree
give_id (bool): whether to return the IDs of nearest points
'''
if use_kdtree:
return chamfer_distance_kdtree(points1, points2, give_id=give_id)
else:
return chamfer_distance_naive(points1, points2)
def chamfer_distance_naive(points1, points2):
''' Naive implementation of the Chamfer distance.
Args:
points1 (numpy array): first point set
points2 (numpy array): second point set
'''
assert(points1.size() == points2.size())
batch_size, T, _ = points1.size()
points1 = points1.view(batch_size, T, 1, 3)
points2 = points2.view(batch_size, 1, T, 3)
distances = (points1 - points2).pow(2).sum(-1)
chamfer1 = distances.min(dim=1)[0].mean(dim=1)
chamfer2 = distances.min(dim=2)[0].mean(dim=1)
chamfer = chamfer1 + chamfer2
return chamfer
def chamfer_distance_kdtree(points1, points2, give_id=False):
''' KD-tree based implementation of the Chamfer distance.
Args:
points1 (numpy array): first point set
points2 (numpy array): second point set
give_id (bool): whether to return the IDs of the nearest points
'''
# Points have size batch_size x T x 3
batch_size = points1.size(0)
# First convert points to numpy
points1_np = points1.detach().cpu().numpy()
points2_np = points2.detach().cpu().numpy()
# Get list of nearest neighbors indieces
idx_nn_12, _ = get_nearest_neighbors_indices_batch(points1_np, points2_np)
idx_nn_12 = torch.LongTensor(idx_nn_12).to(points1.device)
# Expands it as batch_size x 1 x 3
idx_nn_12_expand = idx_nn_12.view(batch_size, -1, 1).expand_as(points1)
# Get list of nearest neighbors indieces
idx_nn_21, _ = get_nearest_neighbors_indices_batch(points2_np, points1_np)
idx_nn_21 = torch.LongTensor(idx_nn_21).to(points1.device)
# Expands it as batch_size x T x 3
idx_nn_21_expand = idx_nn_21.view(batch_size, -1, 1).expand_as(points2)
# Compute nearest neighbors in points2 to points in points1
# points_12[i, j, k] = points2[i, idx_nn_12_expand[i, j, k], k]
points_12 = torch.gather(points2, dim=1, index=idx_nn_12_expand)
# Compute nearest neighbors in points1 to points in points2
# points_21[i, j, k] = points2[i, idx_nn_21_expand[i, j, k], k]
points_21 = torch.gather(points1, dim=1, index=idx_nn_21_expand)
# Compute chamfer distance
chamfer1 = (points1 - points_12).pow(2).sum(2).mean(1)
chamfer2 = (points2 - points_21).pow(2).sum(2).mean(1)
# Take sum
chamfer = chamfer1 + chamfer2
# If required, also return nearest neighbors
if give_id:
return chamfer1, chamfer2, idx_nn_12, idx_nn_21
return chamfer
def get_nearest_neighbors_indices_batch(points_src, points_tgt, k=1):
''' Returns the nearest neighbors for point sets batchwise.
Args:
points_src (numpy array): source points
points_tgt (numpy array): target points
k (int): number of nearest neighbors to return
'''
indices = []
distances = []
for (p1, p2) in zip(points_src, points_tgt):
kdtree = KDTree(p2)
dist, idx = kdtree.query(p1, k=k)
indices.append(idx)
distances.append(dist)
return indices, distances