-
Notifications
You must be signed in to change notification settings - Fork 15
/
eval_S3DIS.py
147 lines (123 loc) · 6.51 KB
/
eval_S3DIS.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
144
145
146
147
import torch
import torch.nn.functional as F
from datasets.S3DIS import S3DIStest, cfl_collate_fn_test
import numpy as np
import MinkowskiEngine as ME
from torch.utils.data import DataLoader
from sklearn.utils.linear_assignment_ import linear_assignment # pip install scikit-learn==0.22.2
from sklearn.cluster import KMeans
from models.fpn import Res16FPN18
from lib.utils import get_fixclassifier
import warnings
import argparse
import os
warnings.filterwarnings('ignore')
###
def parse_args():
parser = argparse.ArgumentParser(description='PyTorch Unsuper_3D_Seg')
parser.add_argument('--data_path', type=str, default='data/S3DIS/input',
help='pont cloud data path')
parser.add_argument('--sp_path', type=str, default='data/S3DIS/initial_superpoints',
help='initial superpoint path')
parser.add_argument('--save_path', type=str, default='trained_models/S3DIS/',
help='model savepath')
###
parser.add_argument('--bn_momentum', type=float, default=0.02, help='batchnorm parameters')
parser.add_argument('--conv1_kernel_size', type=int, default=5, help='kernel size of 1st conv layers')
####
parser.add_argument('--workers', type=int, default=10, help='how many workers for loading data')
parser.add_argument('--cluster_workers', type=int, default=4, help='how many workers for loading data in clustering')
parser.add_argument('--seed', type=int, default=2022, help='random seed')
parser.add_argument('--voxel_size', type=float, default=0.05, help='voxel size in SparseConv')
parser.add_argument('--input_dim', type=int, default=6, help='network input dimension')### 6 for XYZGB
parser.add_argument('--primitive_num', type=int, default=300, help='how many primitives used in training')
parser.add_argument('--semantic_class', type=int, default=12, help='ground truth semantic class')
parser.add_argument('--feats_dim', type=int, default=128, help='output feature dimension')
parser.add_argument('--ignore_label', type=int, default=-1, help='invalid label')
return parser.parse_args()
def eval_once(args, model, test_loader, classifier, use_sp=False):
all_preds, all_label = [], []
for data in test_loader:
with torch.no_grad():
coords, features, inverse_map, labels, index, region = data
in_field = ME.TensorField(features, coords, device=0)
feats = model(in_field)
feats = F.normalize(feats, dim=1)
region = region.squeeze()
#
if use_sp:
region_inds = torch.unique(region)
region_feats = []
for id in region_inds:
if id != -1:
valid_mask = id == region
region_feats.append(feats[valid_mask].mean(0, keepdim=True))
region_feats = torch.cat(region_feats, dim=0)
#
scores = F.linear(F.normalize(feats), F.normalize(classifier.weight))
preds = torch.argmax(scores, dim=1).cpu()
region_scores = F.linear(F.normalize(region_feats), F.normalize(classifier.weight))
region_no = 0
for id in region_inds:
if id != -1:
valid_mask = id == region
preds[valid_mask] = torch.argmax(region_scores, dim=1).cpu()[region_no]
region_no +=1
else:
scores = F.linear(F.normalize(feats), F.normalize(classifier.weight))
preds = torch.argmax(scores, dim=1).cpu()
preds = preds[inverse_map.long()]
all_preds.append(preds[labels!=args.ignore_label]), all_label.append(labels[[labels!=args.ignore_label]])
return all_preds, all_label
def eval(epoch, args, test_areas = ['Area_5']):
model = Res16FPN18(in_channels=args.input_dim, out_channels=args.primitive_num, conv1_kernel_size=args.conv1_kernel_size, config=args).cuda()
model.load_state_dict(torch.load(os.path.join(args.save_path, 'model_' + str(epoch) + '_checkpoint.pth')))
model.eval()
cls = torch.nn.Linear(args.feats_dim, args.primitive_num, bias=False).cuda()
cls.load_state_dict(torch.load(os.path.join(args.save_path, 'cls_' + str(epoch) + '_checkpoint.pth')))
cls.eval()
primitive_centers = cls.weight.data###[300, 128]
print('Merging Primitives')
cluster_pred = KMeans(n_clusters=args.semantic_class, n_init=10, random_state=0, n_jobs=10).fit_predict(primitive_centers.cpu().numpy())#.astype(np.float64))
'''Compute Class Centers'''
centroids = torch.zeros((args.semantic_class, args.feats_dim))
for cluster_idx in range(args.semantic_class):
indices = cluster_pred ==cluster_idx
cluster_avg = primitive_centers[indices].mean(0, keepdims=True)
centroids[cluster_idx] = cluster_avg
# #
centroids = F.normalize(centroids, dim=1)
classifier = get_fixclassifier(in_channel=args.feats_dim, centroids_num=args.semantic_class, centroids=centroids).cuda()
classifier.eval()
test_dataset = S3DIStest(args, areas=test_areas)
test_loader = DataLoader(test_dataset, batch_size=1, collate_fn=cfl_collate_fn_test(), num_workers=4, pin_memory=True)
preds, labels = eval_once(args, model, test_loader, classifier)
all_preds = torch.cat(preds).numpy()
all_labels = torch.cat(labels).numpy()
'''Unsupervised, Match pred to gt'''
sem_num = args.semantic_class
mask = (all_labels >= 0) & (all_labels < sem_num)
histogram = np.bincount(sem_num * all_labels[mask] + all_preds[mask], minlength=sem_num ** 2).reshape(sem_num, sem_num)
'''Hungarian Matching'''
m = linear_assignment(histogram.max() - histogram)
o_Acc = histogram[m[:, 0], m[:, 1]].sum() / histogram.sum()*100.
m_Acc = np.mean(histogram[m[:, 0], m[:, 1]] / histogram.sum(1))*100
hist_new = np.zeros((sem_num, sem_num))
for idx in range(sem_num):
hist_new[:, idx] = histogram[:, m[idx, 1]]
'''Final Metrics'''
tp = np.diag(hist_new)
fp = np.sum(hist_new, 0) - tp
fn = np.sum(hist_new, 1) - tp
IoUs = tp / (tp + fp + fn + 1e-8)
m_IoU = np.nanmean(IoUs)
s = '| mIoU {:5.2f} | '.format(100 * m_IoU)
for IoU in IoUs:
s += '{:5.2f} '.format(100 * IoU)
return o_Acc, m_Acc, s
if __name__ == '__main__':
args = parse_args()
for epoch in range(10, 1500):
if epoch % 1270 == 0:
o_Acc, m_Acc, s = eval(epoch, args)
print('Epoch: {:02d}, oAcc {:.2f} mAcc {:.2f} IoUs'.format(epoch, o_Acc, m_Acc), s)