-
Notifications
You must be signed in to change notification settings - Fork 887
/
evaluate.py
196 lines (172 loc) · 8.05 KB
/
evaluate.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
import argparse
import math
from datetime import datetime
import h5py
import numpy as np
import tensorflow as tf
import socket
import importlib
import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'models'))
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import tf_util
import part_dataset_all_normal
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
parser.add_argument('--model', default='pointnet2_part_seg', help='Model name [default: pointnet2_part_seg]')
parser.add_argument('--model_path', default='log/model.ckpt', help='model checkpoint file path [default: log/model.ckpt]')
parser.add_argument('--log_dir', default='log_eval', help='Log dir [default: log_eval]')
parser.add_argument('--num_point', type=int, default=2048, help='Point Number [default: 2048]')
parser.add_argument('--batch_size', type=int, default=32, help='Batch Size during training [default: 32]')
FLAGS = parser.parse_args()
VOTE_NUM = 12
EPOCH_CNT = 0
BATCH_SIZE = FLAGS.batch_size
NUM_POINT = FLAGS.num_point
GPU_INDEX = FLAGS.gpu
MODEL_PATH = FLAGS.model_path
MODEL = importlib.import_module(FLAGS.model) # import network module
MODEL_FILE = os.path.join(ROOT_DIR, 'models', FLAGS.model+'.py')
LOG_DIR = FLAGS.log_dir
if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR)
os.system('cp %s %s' % (MODEL_FILE, LOG_DIR)) # bkp of model def
os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(FLAGS)+'\n')
NUM_CLASSES = 50
# Shapenet official train/test split
DATA_PATH = os.path.join(ROOT_DIR, 'data', 'shapenetcore_partanno_segmentation_benchmark_v0_normal')
TEST_DATASET = part_dataset_all_normal.PartNormalDataset(root=DATA_PATH, npoints=NUM_POINT, classification=False, split='test')
def log_string(out_str):
LOG_FOUT.write(out_str+'\n')
LOG_FOUT.flush()
print(out_str)
def evaluate():
with tf.Graph().as_default():
with tf.device('/gpu:'+str(GPU_INDEX)):
pointclouds_pl, labels_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT)
is_training_pl = tf.placeholder(tf.bool, shape=())
print is_training_pl
print "--- Get model and loss"
pred, end_points = MODEL.get_model(pointclouds_pl, is_training_pl)
loss = MODEL.get_loss(pred, labels_pl)
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
# Restore variables from disk.
saver.restore(sess, MODEL_PATH)
ops = {'pointclouds_pl': pointclouds_pl,
'labels_pl': labels_pl,
'is_training_pl': is_training_pl,
'pred': pred,
'loss': loss}
eval_one_epoch(sess, ops)
def get_batch(dataset, idxs, start_idx, end_idx):
bsize = end_idx-start_idx
batch_data = np.zeros((bsize, NUM_POINT, 6))
batch_label = np.zeros((bsize, NUM_POINT), dtype=np.int32)
for i in range(bsize):
ps,normal,seg = dataset[idxs[i+start_idx]]
batch_data[i,:,0:3] = ps
batch_data[i,:,3:6] = normal
batch_label[i,:] = seg
return batch_data, batch_label
def eval_one_epoch(sess, ops):
""" ops: dict mapping from string to tf ops """
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
# Test on all data: last batch might be smaller than BATCH_SIZE
num_batches = (len(TEST_DATASET)+BATCH_SIZE-1)/BATCH_SIZE
total_correct = 0
total_seen = 0
loss_sum = 0
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
seg_classes = TEST_DATASET.seg_classes
shape_ious = {cat:[] for cat in seg_classes.keys()}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT))
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 6))
batch_label = np.zeros((BATCH_SIZE, NUM_POINT)).astype(np.int32)
for batch_idx in range(num_batches):
if batch_idx %20==0:
log_string('%03d/%03d'%(batch_idx, num_batches))
start_idx = batch_idx * BATCH_SIZE
end_idx = min(len(TEST_DATASET), (batch_idx+1) * BATCH_SIZE)
cur_batch_size = end_idx-start_idx
cur_batch_data, cur_batch_label = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
if cur_batch_size == BATCH_SIZE:
batch_data = cur_batch_data
batch_label = cur_batch_label
else:
batch_data[0:cur_batch_size] = cur_batch_data
batch_label[0:cur_batch_size] = cur_batch_label
# ---------------------------------------------------------------------
loss_val = 0
pred_val = np.zeros((BATCH_SIZE, NUM_POINT, NUM_CLASSES))
for _ in range(VOTE_NUM):
feed_dict = {ops['pointclouds_pl']: batch_data,
ops['labels_pl']: batch_label,
ops['is_training_pl']: is_training}
temp_loss_val, temp_pred_val = sess.run([ops['loss'], ops['pred']], feed_dict=feed_dict)
loss_val += temp_loss_val
pred_val += temp_pred_val
loss_val /= float(VOTE_NUM)
# ---------------------------------------------------------------------
# Select valid data
cur_pred_val = pred_val[0:cur_batch_size]
# Constrain pred to the groundtruth classes (selected by seg_classes[cat])
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32)
for i in range(cur_batch_size):
cat = seg_label_to_cat[cur_batch_label[i,0]]
logits = cur_pred_val_logits[i,:,:]
cur_pred_val[i,:] = np.argmax(logits[:,seg_classes[cat]], 1) + seg_classes[cat][0]
correct = np.sum(cur_pred_val == cur_batch_label)
total_correct += correct
total_seen += (cur_batch_size*NUM_POINT)
if cur_batch_size==BATCH_SIZE:
loss_sum += loss_val
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum(cur_batch_label==l)
total_correct_class[l] += (np.sum((cur_pred_val==l) & (cur_batch_label==l)))
for i in range(cur_batch_size):
segp = cur_pred_val[i,:]
segl = cur_batch_label[i,:]
cat = seg_label_to_cat[segl[0]]
part_ious = [0.0 for _ in range(len(seg_classes[cat]))]
for l in seg_classes[cat]:
if (np.sum(segl==l) == 0) and (np.sum(segp==l) == 0): # part is not present, no prediction as well
part_ious[l-seg_classes[cat][0]] = 1.0
else:
part_ious[l-seg_classes[cat][0]] = np.sum((segl==l) & (segp==l)) / float(np.sum((segl==l) | (segp==l)))
shape_ious[cat].append(np.mean(part_ious))
all_shape_ious = []
for cat in shape_ious.keys():
for iou in shape_ious[cat]:
all_shape_ious.append(iou)
shape_ious[cat] = np.mean(shape_ious[cat])
print len(all_shape_ious)
mean_shape_ious = np.mean(shape_ious.values())
log_string('eval mean loss: %f' % (loss_sum / float(len(TEST_DATASET)/BATCH_SIZE)))
log_string('eval accuracy: %f'% (total_correct / float(total_seen)))
log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float))))
for cat in sorted(shape_ious.keys()):
log_string('eval mIoU of %s:\t %f'%(cat, shape_ious[cat]))
log_string('eval mean mIoU: %f' % (mean_shape_ious))
log_string('eval mean mIoU (all shapes): %f' % (np.mean(all_shape_ious)))
if __name__ == "__main__":
log_string('pid: %s'%(str(os.getpid())))
evaluate()
LOG_FOUT.close()