def eval_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = int(math.ceil((1.0*len(TEST_DATASET))/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)]
labelweights = np.zeros(NUM_CLASSES)
tp = np.zeros(NUM_CLASSES)
fp = np.zeros(NUM_CLASSES)
fn = np.zeros(NUM_CLASSES)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
if FLAGS.extra_dims:
aug_data = np.concatenate((provider.rotate_point_cloud_z(batch_data[:,:,0:3]),
batch_data[:,:,3:]),axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training}
summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'],
ops['loss'], ops['pred']], feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum((pred_val == batch_label) & (batch_label>0) & (batch_smpw>0)) # evaluate only all categories except unknown
total_correct += correct
total_seen += np.sum((batch_label>0) & (batch_smpw>0))
loss_sum += loss_val
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label==l) & (batch_smpw>0))
total_correct_class[l] += np.sum((pred_val==l) & (batch_label==l) & (batch_smpw>0))
tp[l] += ((pred_val==l) & (batch_label==l)).sum()
fp[l] += ((pred_val==l) & (batch_label!=l)).sum()
fn[l] += ((pred_val!=l) & (batch_label==l)).sum()
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point accuracy: %f'% (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(np.array(total_correct_class)/(np.array(total_seen_class,dtype=np.float)+1e-6))))
per_class_str = ' '
iou = np.divide(tp,tp+fp+fn)
for l in range(NUM_CLASSES):
per_class_str += 'class %d[%d] acc: %f, iou: %f; ' % (TEST_DATASET.decompress_label_map[l],l,total_correct_class[l]/float(total_seen_class[l]),iou[l])
log_string(per_class_str)
log_string('mIOU: {}'.format(iou.mean()))
return total_correct/float(total_seen)
def train_one_epoch(self, sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, self.train_data.__len__(
)) # train_idx ,idx is a list [1263, 396,1309 .... 41 398 495]
np.random.shuffle(train_idxs)
num_batches = int(math.ceil((1.0 * self.train_sz) / self.batch_sz))
self.log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * self.batch_sz
end_idx = (batch_idx + 1) * self.batch_sz
batch_data, batch_label, batch_smpw = self.get_batch_wdp(
self.train_data, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation and jitter
if FLAGS.extra_dims:
aug_data = np.concatenate((provider.rotate_point_cloud_z(
batch_data[:, :, 0:3]), batch_data[:, :, 3:]),
axis=2)
aug_data = np.concatenate((provider.jitter_point_cloud(
aug_data[:, :, 0:3]), aug_data[:, :, 3:]),
axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
aug_data = provider.jitter_point_cloud(aug_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val = sess.run(
[
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['pred']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
# pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (self.batch_sz * self.point_sz)
loss_sum += loss_val
if (batch_idx + 1) % 5 == 0:
self.log_string(' -- %03d / %03d --' %
(batch_idx + 1, num_batches))
self.log_string('mean loss: %f' % (loss_sum / 5))
self.log_string('accuracy: %f' %
(total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def train_one_epoch_my(sess, ops, train_writer):
is_training = True
# ops: dict mapping from string to tf ops """
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = int(math.ceil((1.0*len(TRAIN_DATASET))/BATCH_SIZE))
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(TRAIN_DATASET, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation
if FLAGS.extra_dims:
aug_data = np.concatenate((provider.rotate_point_cloud_z(batch_data[:,:,0:3]),
batch_data[:,:,3:]),axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']:batch_smpw,
ops['is_training_pl']: is_training,}
summary, step, _, loss_val,loss_batch, pred_val = sess.run([ops['merged'], ops['step'],
ops['train_op'], ops['loss'],ops['all_loss'], ops['pred']], feed_dict=feed_dict)
train_writer.add_summary(summary, step)
loss_batch_row = tf.reshape(loss_batch,[loss_batch.shape[0]*loss_batch.shape[1],])
top_k, index = tf.nn.top_k(loss_batch_row, int(loss_batch.shape[0]*loss_batch.shape[1]*0.05)) #提取前20%个loss最大索引
rows, cols = get_batch_index(index.eval(session=sess))
print("Alive1")
for i,j in zip(rows, cols):
# i = m.eval(session=sess)
# j = n.eval(session=sess)
bad_sample_data.append(batch_data[i,j,:])
bad_sample_label.append(batch_label[i,j])
bad_sample_smpw.append(batch_smpw[i,j])
print("Alive2")
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE*NUM_POINT)
loss_sum += loss_val
if (batch_idx+1)%10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx+1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def train_one_epoch(sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = int(math.ceil((1.0 * len(TRAIN_DATASET)) / BATCH_SIZE))
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(
TRAIN_DATASET, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation
if FLAGS.extra_dims:
aug_data = np.concatenate((provider.rotate_point_cloud_z(
batch_data[:, :, 0:3]), batch_data[:, :, 3:]),
axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val = sess.run([
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['pred']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
# print(np.unique(batch_label))
# print(np.unique(pred_val))
# input()
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
loss_sum += loss_val
if (batch_idx + 1) % 10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx + 1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def pre_processor(files, input_queue):
for file in files:
logging.info('Loading {}'.format(file))
pset = PointSet(file)
[block_data_total, block_label_total] = pset.cloud2blocks_finaltest()
test_data_psets = np.concatenate((block_data_total[0], block_data_total[1], block_data_total[2],
block_data_total[3], block_data_total[4], block_data_total[5],
block_data_total[6], block_data_total[7], block_data_total[8],
block_data_total[9]), axis=0) # concat all the blocks
sz=test_data_psets.__len__()
print(sz)
# psets = pset.split() # current splited
num_batches = int(math.ceil((1.0 * sz) / BATCH_SIZE))
print('num_batches: ', num_batches)
data = []
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
for k in range(FLAGS.n_angles):
batch_raw, batch_data = get_batch(test_data_psets, start_idx, end_idx)
if k == 0:
aug_data = batch_data
# aug_data = np.concatenate((provider.rotate_point_cloud_z(batch_data[:, :, 0:3]),
# batch_data[:, :, 3:]), axis=2)
# aug_data = np.concatenate((provider.jitter_point_cloud(aug_data[:, :, 0:3]),
# aug_data[:, :, 3:]), axis=2)
else:
ang = (1.0 * k) / (1.0 * FLAGS.n_angles) * 2 * np.pi
if FLAGS.extra_dims:
aug_data = np.concatenate((provider.rotate_point_cloud_z(batch_data[:, :, 0:3], angle=ang),
batch_data[:, :, 3:]), axis=2)
# aug_data = np.concatenate((provider.rotate_point_cloud_z(batch_data[:, :, 0:3]),
# batch_data[:, :, 3:]), axis=2)
# aug_data = np.concatenate((provider.jitter_point_cloud(aug_data[:, :, 0:3]),
# aug_data[:, :, 3:]), axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
aug_data = provider.jitter_point_cloud(aug_data)
data.append((batch_raw, aug_data))
logging.debug('Adding {} to queue'.format(file))
input_queue.put((pset, data))
logging.debug('Added {} to queue'.format(file))
logging.info('Pre-processing finished')
input_queue.put(None)
logging.debug('Pre-processing thread finished')
def train_one_epoch(sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
log_string('----')
# Shuffle train samples
train_idxs = np.arange(0, len(train_xyz))
np.random.shuffle(train_idxs)
num_batches = len(train_xyz) // BATCH_SIZE
# num_batches = len(trainSet) // BATCH_SIZE
# num_batches = 20
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw, batch_feats = get_batch_wdp(
'train', batch_idx)
if batch_idx % (num_batches / 2) == 0:
print('Current batch/total batch num: %d/%d' %
(batch_idx, num_batches))
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['feature_pl']: batch_feats,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val, lr_val = sess.run(
[
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['pred'], ops['learnrate']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
loss_sum += loss_val
log_string('learn rate: %f' % (lr_val))
log_string('mean loss: %f' % (loss_sum / float(num_batches)))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
mloss = loss_sum / float(num_batches)
macc = total_correct / float(total_seen)
return mloss, macc
def eval_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = len(TEST_DATASET) / 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)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----' % (EPOCH_CNT))
labelweights = np.zeros(21)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(
TEST_DATASET, test_idxs, start_idx, end_idx)
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training
}
summary, step, loss_val, pred_val = sess.run(
[ops['merged'], ops['step'], ops['loss'], ops['pred']],
feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum(
(pred_val == batch_label) &
(batch_smpw > 0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_smpw > 0))
loss_sum += loss_val
tmp, _ = np.histogram(batch_label, range(22))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l)
& (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l)
& (batch_smpw > 0))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class[0:]) /
(np.array(total_seen_class[0:], dtype=np.float) + 1e-6))))
EPOCH_CNT += 1
return total_correct / float(total_seen)
def train_one_epoch(sess, ops, train_writer, epoch=None):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = int(len(TRAIN_DATASET) / BATCH_SIZE)
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
total_iou_deno = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
#batch_data, batch_label, batch_smpw = get_batch_wdp(TRAIN_DATASET, train_idxs, start_idx, end_idx)
batch_data, batch_label, batch_smpw = get_batch(
TRAIN_DATASET, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation
batch_data[:, :, :3] = provider.rotate_point_cloud_z(
batch_data[:, :, :3])
#aug_data = provider.rotate_point_cloud(batch_data)
feed_dict = {
ops['pointclouds_pl']: batch_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val = sess.run([
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['pred']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
iou_deno = 0
for l in range(NUM_CLASSES):
iou_deno += np.sum((pred_val == l) | (batch_label == l))
total_iou_deno += iou_deno
loss_sum += loss_val
if (batch_idx + 1) % 10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx + 1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
log_string('total IoU: %f' %
(total_correct / float(total_iou_deno)))
total_correct = 0
total_seen = 0
loss_sum = 0
total_iou_deno = 0
def pre_processor(files, input_queue):
for file in files:
logging.info('Loading {}'.format(file))
pset = PointSet(file)
psets = pset.split()
num_batches = int(math.ceil((1.0 * len(psets)) / BATCH_SIZE))
data = []
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
for k in range(FLAGS.n_angles):
batch_raw, batch_data = get_batch(psets, start_idx, end_idx)
if k == 0:
aug_data = batch_data
else:
ang = (1.0 * k) / (1.0 * FLAGS.n_angles) * 2 * np.pi
if FLAGS.extra_dims:
aug_data = np.concatenate(
(provider.rotate_point_cloud_z(
batch_data[:, :, 0:3],
angle=ang), batch_data[:, :, 3:]),
axis=2)
else:
aug_data = provider.rotate_point_cloud_z(batch_data)
data.append((batch_raw, aug_data))
logging.debug('Adding {} to queue'.format(file))
input_queue.put((pset, data))
logging.debug('Added {} to queue'.format(file))
logging.info('Pre-processing finished')
input_queue.put(None)
logging.debug('Pre-processing thread finished')
def train_one_epoch(sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = len(TRAIN_DATASET) / BATCH_SIZE
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch_wdp(
TRAIN_DATASET, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops["pointclouds_pl"]: aug_data,
ops["labels_pl"]: batch_label,
ops["smpws_pl"]: batch_smpw,
ops["is_training_pl"]: is_training,
}
summary, step, _, loss_val, pred_val = sess.run(
[
ops["merged"], ops["step"], ops["train_op"], ops["loss"],
ops["pred"]
],
feed_dict=feed_dict,
)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += BATCH_SIZE * NUM_POINT
loss_sum += loss_val
if (batch_idx + 1) % 10 == 0:
log_string(" -- %03d / %03d --" % (batch_idx + 1, num_batches))
log_string("mean loss: %f" % (loss_sum / 10))
log_string("accuracy: %f" % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def visualization(sess, ops):
batch_data, batch_label, batch_smpw = get_batch(
TEST_DATASET, np.arange(0, len(TEST_DATASET)), 0, BATCH_SIZE)
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: False
}
summary, step, loss_val, pred_val = sess.run(
[ops['merged'], ops['step'], ops['loss'], ops['pred']],
feed_dict=feed_dict)
pred_val = np.argmax(pred_val, 2)
point_gt = [[] for i in range(NUM_CLASSES)]
point_pr = [[] for i in range(NUM_CLASSES)]
for i in range(batch_label[0].shape[0]):
point_gt[batch_label[0, i]].append(aug_data[0, i, :])
point_pr[pred_val[0, i]].append(aug_data[0, i, :])
drawpc(FLAGS.suffix, point_gt, point_pr)
def train_one_epoch(sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = len(TRAIN_DATASET)/BATCH_SIZE
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch_wdp(TRAIN_DATASET, train_idxs, start_idx, end_idx)
# Augment batched point clouds by rotation
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']:batch_smpw,
ops['is_training_pl']: is_training,}
summary, step, _, loss_val, pred_val = sess.run([ops['merged'], ops['step'],
ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE*NUM_POINT)
loss_sum += loss_val
if (batch_idx+1)%10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx+1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def training(self):
with tf.Graph().as_default():
self.build_graph()
# merge operator (for tensorboard)
merged = tf.summary.merge_all()
iter_in_epoch = self.train_sz // self.batch_sz
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
best_acc = 0.0
with tf.Session(config=config) as sess:
train_writer = tf.summary.FileWriter(TRAIN_LOG_PATH,
sess.graph)
evaluate_writer = tf.summary.FileWriter(
TEST_LOG_PATH, sess.graph)
sess.run(tf.global_variables_initializer())
epoch_sz = MAX_EPOCH
tic = time.time()
for epoch in range(epoch_sz):
ave_loss = 0
train_idxs = np.arange(0, self.train_data.__len__())
np.random.shuffle(train_idxs)
for _iter in range(iter_in_epoch):
start_idx = _iter * self.batch_sz
end_idx = (_iter + 1) * self.batch_sz
batch_data, batch_label, batch_smpw = self.get_batch_wdp(
self.train_data, train_idxs, start_idx, end_idx)
aug_data = provider.rotate_point_cloud_z(batch_data)
loss, _, summary, step = sess.run(
[self.loss, self.train_op, merged, self.batch],
feed_dict={
self.point_pl: aug_data,
self.label_pl: batch_label,
self.smpws_pl: batch_smpw,
self.is_train_pl: True
})
ave_loss += loss
train_writer.add_summary(summary, step)
ave_loss /= iter_in_epoch
#print("epoch %d , loss is %f take %.3f s" % (epoch + 1, ave_loss, time.time() - tic))
log_string("epoch %d , loss is %f take %.3f s" %
(epoch + 1, ave_loss, time.time() - tic))
tic = time.time()
if (epoch + 1) % 5 == 0:
acc = self.evaluate_one_epoch(sess, evaluate_writer,
step, epoch)
if acc > best_acc:
_path = saver.save(
sess,
os.path.join(
SAVE_PATH,
"best_seg_model_%d.ckpt" % (epoch + 1)))
log_string("epoch %d " % (epoch + 1))
#print("epoch %d, best saved in file: " % (epoch + 1), _path)
best_acc = acc
_path = saver.save(
sess, os.path.join(SAVE_PATH, 'train_base_seg_model.ckpt'))
print("Model saved in file: ", _path)
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('sem_seg')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/stanford_indoor3d/'
NUM_CLASSES = 13
NUM_POINT = args.npoint
BATCH_SIZE = args.batch_size
FEATURE_CHANNEL = 3 if args.with_rgb else 0
print("start loading training data ...")
TRAIN_DATASET = S3DISDataset(root, split='train', with_rgb=args.with_rgb, test_area=args.test_area, block_points=NUM_POINT)
print("start loading whole scene validation data ...")
TEST_DATASET_WHOLE_SCENE = S3DISDatasetWholeScene(root, split='test', with_rgb=args.with_rgb, test_area=args.test_area, block_points=NUM_POINT)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=BATCH_SIZE,shuffle=True, num_workers=4)
weights = TRAIN_DATASET.labelweights
weights = torch.Tensor(weights).cuda()
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE))
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(experiment_dir))
shutil.copy('models/pointnet_util.py', str(experiment_dir))
classifier = MODEL.get_model(NUM_CLASSES, with_rgb=args.with_rgb).cuda()
criterion = MODEL.get_loss().cuda()
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init)
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate
)
else:
optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
global_epoch = 0
best_iou = 0
for epoch in range(start_epoch,args.epoch):
log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch))
'''Adjust learning rate and BN momentum'''
lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
mean_correct = []
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(lambda x: bn_momentum_adjust(x,momentum))
'''learning one epoch'''
for i, data in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
points, target, weight = data
points = points.data.numpy()
points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3])
points[:,:, 0:3] = provider.rotate_point_cloud_z(points[:,:, 0:3])
points = torch.Tensor(points)
points, target = points.float().cuda(),target.long().cuda()
points = points.transpose(2, 1)
optimizer.zero_grad()
classifier = classifier.train()
seg_pred, trans_feat = classifier(points)
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
target = target.view(-1, 1)[:, 0]
pred_choice = seg_pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
mean_correct.append(correct.item() / (args.batch_size * args.npoint))
loss = criterion(seg_pred, target, trans_feat, weights)
loss.backward()
optimizer.step()
train_instance_acc = np.mean(mean_correct)
log_string('Train accuracy is: %.5f' % train_instance_acc)
# evaluate on whole scenes, for each block, only sample NUM_POINT points
if epoch % 10 ==0:
with torch.no_grad():
num_batches = len(TEST_DATASET_WHOLE_SCENE)
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)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
labelweights = np.zeros(NUM_CLASSES)
is_continue_batch = False
extra_batch_data = np.zeros((0, NUM_POINT, 3 + FEATURE_CHANNEL))
extra_batch_label = np.zeros((0, NUM_POINT))
extra_batch_smpw = np.zeros((0, NUM_POINT))
for batch_idx in tqdm(range(num_batches),total=num_batches):
if not is_continue_batch:
batch_data, batch_label, batch_smpw = TEST_DATASET_WHOLE_SCENE[batch_idx]
batch_data = np.concatenate((batch_data, extra_batch_data), axis=0)
batch_label = np.concatenate((batch_label, extra_batch_label), axis=0)
batch_smpw = np.concatenate((batch_smpw, extra_batch_smpw), axis=0)
else:
batch_data_tmp, batch_label_tmp, batch_smpw_tmp = TEST_DATASET_WHOLE_SCENE[batch_idx]
batch_data = np.concatenate((batch_data, batch_data_tmp), axis=0)
batch_label = np.concatenate((batch_label, batch_label_tmp), axis=0)
batch_smpw = np.concatenate((batch_smpw, batch_smpw_tmp), axis=0)
if batch_data.shape[0] < BATCH_SIZE:
is_continue_batch = True
continue
elif batch_data.shape[0] == BATCH_SIZE:
is_continue_batch = False
extra_batch_data = np.zeros((0, NUM_POINT, 3 + FEATURE_CHANNEL))
extra_batch_label = np.zeros((0, NUM_POINT))
extra_batch_smpw = np.zeros((0, NUM_POINT))
else:
is_continue_batch = False
extra_batch_data = batch_data[BATCH_SIZE:, :, :]
extra_batch_label = batch_label[BATCH_SIZE:, :]
extra_batch_smpw = batch_smpw[BATCH_SIZE:, :]
batch_data = batch_data[:BATCH_SIZE, :, :]
batch_label = batch_label[:BATCH_SIZE, :]
batch_smpw = batch_smpw[:BATCH_SIZE, :]
batch_label = torch.Tensor(batch_label)
batch_data = torch.Tensor(batch_data)
batch_data, batch_label = batch_data.float().cuda(), batch_label.long().cuda()
batch_data = batch_data.transpose(2, 1)
seg_pred, _ = classifier(batch_data)
seg_pred = seg_pred.contiguous()
batch_label = batch_label.cpu().data.numpy()
pred_val = seg_pred.cpu().data.max(2)[1].numpy()
correct = np.sum((pred_val == batch_label) & (batch_smpw > 0))
total_correct += correct
total_seen += np.sum(batch_smpw > 0)
tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l) & (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l) & (batch_label == l) & (batch_smpw > 0))
total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l)) & (batch_smpw > 0))
mIoU = np.mean(np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6))
log_string('eval whole scene mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % mIoU)
log_string('eval whole scene point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval whole scene point avg class acc: %f' % (
np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6))))
labelweights = labelweights.astype(np.float32) / np.sum(labelweights.astype(np.float32))
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])), labelweights[l],
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
if (mIoU >= best_iou):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'class_avg_iou': mIoU,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
global_epoch += 1
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('classification')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
'''DATA LOADING'''
log_string('Load dataset ...')
traindata_manager = Manager()
testdata_manager = Manager()
train_cache = traindata_manager.dict()
test_cache = testdata_manager.dict()
DATA_PATH = '/disk/users/sc468/no_backup/my_kitti'
TRAIN_DATASET = MyKittiDataLoader(root=DATA_PATH,
cache=train_cache,
npoint=args.num_point,
split='train',
normal_channel=args.normal)
TEST_DATASET = MyKittiDataLoader(root=DATA_PATH,
cache=test_cache,
npoint=args.num_point,
split='test',
normal_channel=args.normal)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
'''MODEL LOADING'''
num_class = 4
MODEL = importlib.import_module(args.model)
shutil.copy('./models/%s.py' % args.model, str(experiment_dir))
shutil.copy('./models/pointnet_util.py', str(experiment_dir))
classifier = MODEL.get_model(num_class, normal_channel=args.normal).cuda()
criterion = MODEL.get_loss().cuda()
try:
checkpoint = torch.load(
str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
else:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=0.01,
momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=20,
gamma=0.7)
global_epoch = 0
global_step = 0
best_instance_acc = 0.0
best_class_acc = 0.0
mean_correct = []
'''TRANING'''
logger.info('Start training...')
for epoch in range(start_epoch, args.epoch):
log_string('Epoch %d (%d/%s):' %
(global_epoch + 1, epoch + 1, args.epoch))
print("Train cache size {}".format(len(train_cache)))
print("Test cache size {}".format(len(test_cache)))
mean_loss = 0
scheduler.step()
for batch_id, data in tqdm(enumerate(trainDataLoader, 0),
total=len(trainDataLoader),
smoothing=0.9):
points, target = data
points = points.data.numpy()
# data augmentation
points = provider.random_point_dropout(points)
points[:, :, 0:3] = provider.rotate_point_cloud_z(points[:, :,
0:3])
points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :,
0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
points[:, :, 0:3] = provider.jitter_point_cloud(points[:, :, 0:3])
points = torch.Tensor(points)
target = target[:, 0]
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
optimizer.zero_grad()
classifier = classifier.train()
pred, trans_feat = classifier(points)
loss = criterion(pred, target.long(), trans_feat)
mean_loss += loss.item()
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.long().data).cpu().sum()
mean_correct.append(correct.item() / float(points.size()[0]))
loss.backward()
optimizer.step()
global_step += 1
train_instance_acc = np.mean(mean_correct)
log_string('Train Instance Accuracy: %f' % train_instance_acc)
# tensorboard
mean_loss /= len(trainDataLoader)
writer_train.add_scalar('loss', mean_loss, epoch)
writer_train.add_scalar('overall accuracy', train_instance_acc, epoch)
with torch.no_grad():
instance_acc, class_acc = test(epoch, classifier.eval(), criterion,
testDataLoader)
mean_class_acc = np.mean(class_acc[:, 2])
if (instance_acc >= best_instance_acc):
best_instance_acc = instance_acc
if (mean_class_acc >= best_class_acc):
best_class_acc = mean_class_acc
best_epoch = epoch + 1
log_string('Test Instance Accuracy: %f, Class Accuracy: %f' %
(instance_acc, mean_class_acc))
log_string('Best Instance Accuracy: %f, Class Accuracy: %f' %
(best_instance_acc, best_class_acc))
log_string('Each Class Accuracy:')
log_string('Car: %f' % class_acc[0, 2])
log_string('Cyclist: %f' % class_acc[1, 2])
log_string('Pedestrian: %f' % class_acc[2, 2])
log_string('DontCare: %f' % class_acc[3, 2])
print("")
# tensorboard
writer_test.add_scalar('overall accuracy', instance_acc, epoch)
writer_car_acc.add_scalar('accuracy of each class',
class_acc[0, 2], epoch)
writer_cyclist_acc.add_scalar('accuracy of each class',
class_acc[1, 2], epoch)
writer_pedestrian_acc.add_scalar('accuracy of each class',
class_acc[2, 2], epoch)
writer_dontcare_acc.add_scalar('accuracy of each class',
class_acc[3, 2], epoch)
if (mean_class_acc >= best_class_acc):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': best_epoch,
'instance_acc': instance_acc,
'class_acc': mean_class_acc,
'each_class_acc': class_acc[:, 2],
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
global_epoch += 1
logger.info('End of training...')
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
experiment_dir = 'log/sem_seg/' + args.log_dir
visual_dir = experiment_dir + '/visual/'
visual_dir = Path(visual_dir)
visual_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
NUM_CLASSES = 13
WITH_RGB = args.with_rgb
BATCH_SIZE = args.batch_size
NUM_POINT = args.num_point
root = 'data/stanford_indoor3d/'
TEST_DATASET_WHOLE_SCENE = ScannetDatasetWholeScene_evaluation(
root,
split='test',
with_rgb=WITH_RGB,
test_area=args.test_area,
block_points=NUM_POINT)
log_string("The number of test data is: %d" %
len(TEST_DATASET_WHOLE_SCENE))
'''MODEL LOADING'''
model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
MODEL = importlib.import_module(model_name)
classifier = MODEL.get_model(NUM_CLASSES, with_rgb=WITH_RGB).cuda()
checkpoint = torch.load(
str(experiment_dir) + '/checkpoints/best_model.pth')
classifier.load_state_dict(checkpoint['model_state_dict'])
with torch.no_grad():
scene_id = TEST_DATASET_WHOLE_SCENE.file_list
scene_id = [x[:-4] for x in scene_id]
num_batches = len(TEST_DATASET_WHOLE_SCENE)
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
log_string('---- EVALUATION WHOLE SCENE----')
for batch_idx in range(num_batches):
print("visualize %d %s ..." % (batch_idx, scene_id[batch_idx]))
total_seen_class_tmp = [0 for _ in range(NUM_CLASSES)]
total_correct_class_tmp = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class_tmp = [0 for _ in range(NUM_CLASSES)]
if args.visual:
fout = open(
os.path.join(visual_dir,
scene_id[batch_idx] + '_pred.obj'), 'w')
fout_gt = open(
os.path.join(visual_dir, scene_id[batch_idx] + '_gt.obj'),
'w')
whole_scene_data = TEST_DATASET_WHOLE_SCENE.scene_points_list[
batch_idx]
whole_scene_label = TEST_DATASET_WHOLE_SCENE.semantic_labels_list[
batch_idx]
vote_label_pool = np.zeros(
(whole_scene_label.shape[0], NUM_CLASSES))
for _ in tqdm(range(args.num_votes), total=args.num_votes):
scene_data, scene_label, scene_smpw, scene_point_index = TEST_DATASET_WHOLE_SCENE[
batch_idx]
num_blocks = scene_data.shape[0]
s_batch_num = (num_blocks + BATCH_SIZE - 1) // BATCH_SIZE
if WITH_RGB:
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 6))
else:
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 3))
batch_label = np.zeros((BATCH_SIZE, NUM_POINT))
batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT))
batch_smpw = np.zeros((BATCH_SIZE, NUM_POINT))
for sbatch in range(s_batch_num):
start_idx = sbatch * BATCH_SIZE
end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks)
real_batch_size = end_idx - start_idx
batch_data[0:real_batch_size,
...] = scene_data[start_idx:end_idx, ...]
batch_label[0:real_batch_size,
...] = scene_label[start_idx:end_idx, ...]
batch_point_index[0:real_batch_size,
...] = scene_point_index[
start_idx:end_idx, ...]
batch_smpw[0:real_batch_size,
...] = scene_smpw[start_idx:end_idx, ...]
if WITH_RGB:
batch_data[:, :, 3:6] /= 1.0
aug_data = batch_data
aug_data[:, :, :3] = provider.rotate_point_cloud_z(
aug_data[:, :, :3])
aug_data = torch.Tensor(aug_data)
aug_data = aug_data.float().cuda()
aug_data = aug_data.transpose(2, 1)
seg_pred, _ = classifier(aug_data)
batch_pred_label = seg_pred.contiguous().cpu().data.max(
2)[1].numpy()
vote_label_pool = add_vote(
vote_label_pool, batch_point_index[0:real_batch_size,
...],
batch_pred_label[0:real_batch_size, ...],
batch_smpw[0:real_batch_size, ...])
pred_label = np.argmax(vote_label_pool, 1)
for l in range(NUM_CLASSES):
total_seen_class_tmp[l] += np.sum((whole_scene_label == l))
total_correct_class_tmp[l] += np.sum((pred_label == l) &
(whole_scene_label == l))
total_iou_deno_class_tmp[l] += np.sum(
((pred_label == l) | (whole_scene_label == l)))
total_seen_class[l] += total_seen_class_tmp[l]
total_correct_class[l] += total_correct_class_tmp[l]
total_iou_deno_class[l] += total_iou_deno_class_tmp[l]
iou_map = np.array(total_correct_class_tmp) / (
np.array(total_iou_deno_class_tmp, dtype=np.float) + 1e-6)
print(iou_map)
arr = np.array(total_seen_class_tmp)
tmp_iou = np.mean(iou_map[arr != 0])
log_string('Mean IoU of %s: %.4f' % (scene_id[batch_idx], tmp_iou))
print('----------------------------')
filename = os.path.join(visual_dir, scene_id[batch_idx] + '.txt')
with open(filename, 'w') as pl_save:
for i in pred_label:
pl_save.write(str(int(i)) + '\n')
pl_save.close()
for i in range(whole_scene_label.shape[0]):
color = g_label2color[pred_label[i]]
color_gt = g_label2color[whole_scene_label[i]]
if args.visual:
fout.write(
'v %f %f %f %d %d %d\n' %
(whole_scene_data[i, 0], whole_scene_data[i, 1],
whole_scene_data[i, 2], color[0], color[1], color[2]))
fout_gt.write(
'v %f %f %f %d %d %d\n' %
(whole_scene_data[i, 0], whole_scene_data[i, 1],
whole_scene_data[i, 2], color_gt[0], color_gt[1],
color_gt[2]))
if args.visual:
fout.close()
fout_gt.close()
IoU = np.array(total_correct_class) / (
np.array(total_iou_deno_class, dtype=np.float) + 1e-6)
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])),
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
log_string('eval point avg class IoU: %f' % np.mean(IoU))
log_string('eval whole scene point avg class acc: %f' % (np.mean(
np.array(total_correct_class) /
(np.array(total_seen_class, dtype=np.float) + 1e-6))))
log_string('eval whole scene point accuracy: %f' %
(np.sum(total_correct_class) /
float(np.sum(total_seen_class) + 1e-6)))
print("Done!")
def eval_one_epoch(sess, ops, num_votes=1):
is_training = False
num_scans = len(TEST_DATASET_WHOLE_SCENE)
log_string(str(datetime.now()))
log_string('----PREPARING PREDICTIONS----')
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
labelweights = np.zeros(NUM_CLASSES)
full_points_names = TEST_DATASET_WHOLE_SCENE.points_name
for batch_idx in range(num_scans):
start = time.time()
t_seen_class = [0 for _ in range(NUM_CLASSES)]
t_correct_class = [0 for _ in range(NUM_CLASSES)]
t_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
full_points_name = full_points_names[batch_idx]
components = full_points_name.split('/')
sequence = components[-3]
points_name = components[-1]
label_name = points_name.replace('bin', 'label')
log_string("Inference sequence %s-%s [%d/%d] ..." %
(sequence, label_name.split('.')[0], batch_idx, num_scans))
full_save_dir = os.path.join(SAVE_DIR, 'sequences', sequence,
'predictions')
os.makedirs(full_save_dir, exist_ok=True)
full_label_name = os.path.join(full_save_dir, label_name)
whole_scene_label = None
for vote_idx in range(num_votes):
print("Voting [%d/%d] ..." % (vote_idx + 1, num_votes))
if DATASET == 'test':
scene_data, scene_point_index, whole_scene_data = TEST_DATASET_WHOLE_SCENE[
batch_idx]
else:
scene_data, scene_point_index, whole_scene_data, whole_scene_label = TEST_DATASET_WHOLE_SCENE[
batch_idx]
num_blocks = scene_data.shape[0]
s_batch_num = (num_blocks + BATCH_SIZE - 1) // BATCH_SIZE
if WITH_REMISSION:
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 4))
else:
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 3))
batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT))
for sbatch in range(s_batch_num):
start_idx = sbatch * BATCH_SIZE
end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks)
real_batch_size = end_idx - start_idx
batch_data[0:real_batch_size,
...] = scene_data[start_idx:end_idx, ...]
batch_point_index[0:real_batch_size,
...] = scene_point_index[start_idx:end_idx,
...]
if FLAGS.random_rotate:
batch_data[:, :, :3] = provider.rotate_point_cloud_z(
batch_data[:, :, :3])
feed_dict = {
ops['pointclouds_pl']: batch_data,
ops['is_training_pl']: is_training
}
pred_val = sess.run(ops['pred'],
feed_dict=feed_dict) # BxNxNUM_CLASSES
batch_pred_label = np.argmax(pred_val[:, :, 1:], 2) + 1 # BxN
if sbatch == 0:
vote_label_pool = np.zeros(
(whole_scene_data.shape[0], NUM_CLASSES))
vote_label_pool = add_vote(
vote_label_pool, batch_point_index[0:real_batch_size, ...],
batch_pred_label[0:real_batch_size, ...])
final_preds = np.argmax(vote_label_pool, axis=1)
final_preds = final_preds.astype(np.uint32)
print('writing %s' % full_label_name)
final_preds.tofile(full_label_name)
print('Visualuze %s-%s' % (sequence, label_name.split('.')[0]))
end = time.time()
print('Use Time %.2f s' % (end - start))
fout = open(
os.path.join(DUMP_DIR, '%s_%s_pred.obj' % (sequence, label_name)),
'w')
if DATASET != 'test':
fout_gt = open(
os.path.join(DUMP_DIR,
'%s_%s_gt.obj' % (sequence, label_name)), 'w')
for i in range(whole_scene_data.shape[0]):
color = g_label2color[final_preds[i]]
fout.write('v %f %f %f %d %d %d\n' %
(whole_scene_data[i, 0], whole_scene_data[i, 1],
whole_scene_data[i, 2], color[0], color[1], color[2]))
if DATASET != 'test':
color_gt = g_label2color[whole_scene_label[i]]
fout_gt.write('v %f %f %f %d %d %d\n' %
(whole_scene_data[i, 0], whole_scene_data[i, 1],
whole_scene_data[i, 2], color_gt[0],
color_gt[1], color_gt[2]))
fout.close()
if DATASET != 'test':
fout_gt.close()
correct = np.sum(final_preds == whole_scene_label)
seen = len(whole_scene_label)
total_correct += correct
total_seen += seen
tmp, _ = np.histogram(whole_scene_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
tem_seen = np.sum(whole_scene_label == l)
t_seen_class[l] += tem_seen
total_seen_class[l] += tem_seen
temp_correct = np.sum((final_preds == l)
& (whole_scene_label == l))
temp_iou_deno_class = np.sum((final_preds == l)
| (whole_scene_label == l))
total_correct_class[l] += temp_correct
t_correct_class[l] += temp_correct
total_iou_deno_class[l] += temp_iou_deno_class
t_iou_deno_class[l] += temp_iou_deno_class
iou = np.array(t_correct_class[1:]) / (
np.array(t_iou_deno_class[1:], dtype=np.float) + 1e-6)
arr = np.array(t_seen_class[1:])
mIoU = np.mean(iou[arr != 0])
log_string('Mean IoU of %s-%s: %.4f' %
(sequence, label_name.split('.')[0], mIoU))
print('---------------------')
if DATASET != 'test':
if batch_idx % 10 == 0:
mIoU = np.mean(
np.array(total_correct_class[1:]) /
(np.array(total_iou_deno_class[1:], dtype=np.float) +
1e-6))
log_string('eval point avg class IoU: %f' % mIoU)
log_string('eval whole scene point accuracy: %f' %
(total_correct / float(total_seen)))
log_string(
'eval whole scene point avg class acc: %f' % (np.mean(
np.array(total_correct_class) /
(np.array(total_seen_class, dtype=np.float) + 1e-6))))
labelweights = labelweights.astype(np.float32) / np.sum(
labelweights.astype(np.float32))
iou_per_class_str = '------- IoU --------\n'
for l in range(1, NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' *
(14 - len(seg_label_to_cat[l])), labelweights[l - 1],
total_correct_class[l] /
float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
print('---------------------')
print("Done!")
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('sem_seg')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
NUM_CLASSES = len(classes)
NUM_POINT = args.npoint
NUM_CHANNEL = args.nchannel
BATCH_SIZE = args.batch_size
FOLD = args.fold
TRAIN_DIR = os.path.join(BASE_DIR, args.train_dir)
TEST_DIR = os.path.join(BASE_DIR, args.test_dir)
print("start loading training data ...")
TRAIN_DATASET = PointCloudDataset(root=TRAIN_DIR,
nClasses=NUM_CLASSES,
nPoints=NUM_POINT,
split="train",
fold=FOLD)
trainDataLoader = torch.utils.data.DataLoader(
TRAIN_DATASET,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4,
pin_memory=True,
drop_last=True,
worker_init_fn=lambda x: np.random.seed(x + int(time.time())))
weights = torch.Tensor(TRAIN_DATASET.labelweights).cuda()
#weights = torch.Tensor([1/5,1/5,1/5,1/5,1/5]).cuda()
print("start loading test data ...")
TEST_DATASET = PointCloudDataset(root=TEST_DIR,
nClasses=NUM_CLASSES,
nPoints=NUM_POINT,
split="test",
fold=FOLD)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=4,
pin_memory=True,
drop_last=True)
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy(ROOT_DIR + '/models/%s.py' % args.model, str(experiment_dir))
shutil.copy(ROOT_DIR + '/models/pointnet_util.py', str(experiment_dir))
classifier = MODEL.get_model(num_class=NUM_CLASSES,
num_point=NUM_POINT,
num_channel=NUM_CHANNEL).cuda()
criterion = MODEL.get_loss().cuda()
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(
str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init)
# TRANSFER LEARNING
if ((args.log_dir is None) & (args.pretrained_dir is not None)):
MODEL_PRE = importlib.import_module("pointnet2_sem_seg")
shutil.copy(ROOT_DIR + '/models/pointnet2_sem_seg.py',
str(experiment_dir))
shutil.copy(ROOT_DIR + '/models/pointnet_util.py', str(experiment_dir))
classifier_pre = MODEL_PRE.get_model(num_class=10,
num_point=512,
num_channel=3) # Paris-Lille-3D
checkpoint = torch.load(
str(args.pretrained_dir) + '/checkpoints/best_model.pth')
classifier_pre.load_state_dict(checkpoint['model_state_dict'])
classifier_dict = classifier.state_dict()
classifier_pre_dict = classifier_pre.state_dict()
# load from pretrained_dict all the layer parameters except for the output layers (conv1, bn1, conv2)
classifier_pre_dict = {
k: v
for k, v in classifier_pre_dict.items() if k in classifier_dict
}
for k in [
'conv1.weight', 'conv1.bias', 'bn1.weight', 'bn1.bias',
'bn1.running_mean', 'bn1.running_var',
'bn1.num_batches_tracked', 'conv2.weight', 'conv2.bias'
]:
classifier_pre_dict.pop(k)
classifier_dict.update(classifier_pre_dict)
# load classifier parameters
classifier.load_state_dict(classifier_dict)
# freeze core layers
'''
for i, child in enumerate(classifier.children()):
if (i < 8):
for param in child.parameters():
param.requires_grad = False
'''
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad,
classifier.parameters()), # classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
else:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=args.learning_rate,
momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(
m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
global_epoch = 0
best_iou = 0
for epoch in range(start_epoch, args.epoch):
log_string('**** Epoch %d (%d/%s) ****' %
(global_epoch + 1, epoch + 1, args.epoch))
lr = max(
args.learning_rate * (args.lr_decay**(epoch // args.step_size)),
LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY
**(epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(
lambda x: bn_momentum_adjust(x, momentum))
num_batches = len(trainDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
'''learning one epoch'''
for i, data in tqdm(enumerate(trainDataLoader),
total=len(trainDataLoader),
smoothing=0.9):
points, target = data
points = points.data.numpy()
points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3])
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
optimizer.zero_grad()
classifier = classifier.train()
seg_pred, trans_feat = classifier(points)
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss.backward()
optimizer.step()
pred_choice = seg_pred.cpu().data.max(1)[1].numpy()
correct = np.sum(pred_choice == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
loss_sum += loss
log_string('Training mean loss: %f' % (loss_sum / num_batches))
log_string('Training accuracy: %f' %
(total_correct / float(total_seen)))
if epoch % 5 == 0:
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
with torch.no_grad():
num_batches = len(testDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
labelweights = np.zeros(NUM_CLASSES)
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_predicted_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
log_string('---- EPOCH %03d EVALUATION ----' % (global_epoch + 1))
for i, data in tqdm(enumerate(testDataLoader),
total=len(testDataLoader),
smoothing=0.9):
points, target = data
points = points.data.numpy()
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
classifier = classifier.eval()
seg_pred, trans_feat = classifier(points)
pred_val = seg_pred.contiguous().cpu().data.numpy()
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss_sum += loss
pred_val = np.argmax(pred_val, 2)
correct = np.sum((pred_val == batch_label))
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l))
total_predicted_class[l] += np.sum((pred_val == l))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l))
total_iou_deno_class[l] += np.sum(
((pred_val == l) | (batch_label == l)))
labelweights = labelweights.astype(np.float32) / np.sum(
labelweights.astype(np.float32))
mIoU = np.mean(
np.array(total_correct_class) /
(np.array(total_iou_deno_class, dtype=np.float) + 1e-6))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % (mIoU))
log_string('eval point accuracy: %f' %
(total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class) /
(np.array(total_seen_class, dtype=np.float) + 1e-6))))
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' *
(14 - len(seg_label_to_cat[l])), labelweights[l],
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
'''
print('AIUTO\n')
for l in range(NUM_CLASSES):
print('class %s %.3f: seen %d, predicted %d, correct %d, union %d \n' % (seg_label_to_cat[l],labelweights[l],total_seen_class[l],total_predicted_class[l],total_correct_class[l],total_iou_deno_class[l]))
'''
log_string('Eval mean loss: %f' % (loss_sum / num_batches))
log_string('Eval accuracy: %f' %
(total_correct / float(total_seen)))
if mIoU >= best_iou:
best_iou = mIoU
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'class_avg_iou': mIoU,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
log_string('Best mIoU: %f' % best_iou)
global_epoch += 1
def eval_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = len(TEST_DATASET)//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)]
total_correct_vox = 0
total_seen_vox = 0
total_seen_class_vox = [0 for _ in range(NUM_CLASSES)]
total_correct_class_vox = [0 for _ in range(NUM_CLASSES)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT))
labelweights = np.zeros(21)
labelweights_vox = np.zeros(21)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training}
summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'],
ops['loss'], ops['pred']], feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum((pred_val == batch_label) & (batch_label>0) & (batch_smpw>0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_label>0) & (batch_smpw>0))
loss_sum += loss_val
tmp,_ = np.histogram(batch_label,range(22))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label==l) & (batch_smpw>0))
total_correct_class[l] += np.sum((pred_val==l) & (batch_label==l) & (batch_smpw>0))
for b in range(batch_label.shape[0]):
_, uvlabel, _ = pc_util.point_cloud_label_to_surface_voxel_label_fast(aug_data[b,batch_smpw[b,:]>0,:], np.concatenate((np.expand_dims(batch_label[b,batch_smpw[b,:]>0],1),np.expand_dims(pred_val[b,batch_smpw[b,:]>0],1)),axis=1), res=0.02)
total_correct_vox += np.sum((uvlabel[:,0]==uvlabel[:,1])&(uvlabel[:,0]>0))
total_seen_vox += np.sum(uvlabel[:,0]>0)
tmp,_ = np.histogram(uvlabel[:,0],range(22))
labelweights_vox += tmp
for l in range(NUM_CLASSES):
total_seen_class_vox[l] += np.sum(uvlabel[:,0]==l)
total_correct_class_vox[l] += np.sum((uvlabel[:,0]==l) & (uvlabel[:,1]==l))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point accuracy vox: %f'% (total_correct_vox / float(total_seen_vox)))
log_string('eval point avg class acc vox: %f' % (np.mean(np.array(total_correct_class_vox[1:])/(np.array(total_seen_class_vox[1:],dtype=np.float)+1e-6))))
log_string('eval point accuracy: %f'% (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(np.array(total_correct_class[1:])/(np.array(total_seen_class[1:],dtype=np.float)+1e-6))))
labelweights_vox = labelweights_vox[1:].astype(np.float32)/np.sum(labelweights_vox[1:].astype(np.float32))
caliweights = np.array([0.388,0.357,0.038,0.033,0.017,0.02,0.016,0.025,0.002,0.002,0.002,0.007,0.006,0.022,0.004,0.0004,0.003,0.002,0.024,0.029])
log_string('eval point calibrated average acc: %f' % (np.average(np.array(total_correct_class[1:])/(np.array(total_seen_class[1:],dtype=np.float)+1e-6),weights=caliweights)))
per_class_str = 'vox based --------'
for l in range(1,NUM_CLASSES):
per_class_str += 'class %d weight: %f, acc: %f; ' % (l,labelweights_vox[l-1],total_correct_class[l]/float(total_seen_class[l]))
log_string(per_class_str)
EPOCH_CNT += 1
return total_correct/float(total_seen)
def train_one_epoch(sess, ops, train_writer, copy_from_model0_op,
copy_to_model0_op):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = len(TRAIN_DATASET) / (BATCH_SIZE // 2)
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
batch_idx_suncg = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * (BATCH_SIZE // 2)
end_idx = (batch_idx + 1) * (BATCH_SIZE // 2)
while DATA_QUEUE_SUN.empty():
pass
temp_batch_data, batch_smpw = DATA_QUEUE_SUN.get()
SUNCG_DATASET.check_gone(
temp_batch_data[-1],
batch_smpw[-1]) ## Only give 12288 points to voxel
batch_data = []
for i in range(len(temp_batch_data)):
batch_data.append(np.zeros((BATCH_SIZE, NUM_POINT_MORE[i], 3)))
batch_data[i][0:BATCH_SIZE // 2, :, :] = temp_batch_data[i]
pred_val = []
counter = 0
for i in range(len(NUM_POINT)):
for j in range(NUM_REP[i]):
sess.run(copy_to_model0_op[i])
feed_dict = {}
feed_dict[ops['is_training_pl']] = False
feed_dict[ops['data_select']] = i
feed_dict[ops['pointclouds_pl' +
str(i + 1)]] = batch_data[counter]
temp_pred_val = sess.run(ops['pred'],
feed_dict=add_empty(
feed_dict, i, ops))
pred_val.append(
np.squeeze(
np.argmax(temp_pred_val[0:BATCH_SIZE // 2, ...], 2)))
counter += 1
### Combine with other sources here
batch_data_extra, batch_label_extra, batch_smpw_extra = SUNCG_DATASET.ready(
temp_batch_data, pred_val, batch_smpw, TRAIN_DATASET.labelweights)
while DATA_QUEUE.empty():
pass
batch_data, batch_label, batch_smpw = DATA_QUEUE.get()
shuffled_data = shuffle_data([batch_data, batch_label, batch_smpw])
batch_data, batch_label, batch_smpw = shuffled_data[0], shuffled_data[
1], shuffled_data[2]
shuffled_data = shuffle_data(
[batch_data_extra, batch_label_extra, batch_smpw_extra])
batch_data_extra, batch_label_extra, batch_smpw_extra = shuffled_data[
0], shuffled_data[1], shuffled_data[2]
### Combine data
counter = 0
for i in range(len(NUM_POINT)):
for j in range(NUM_REP[i]):
if j == 0:
sess.run(copy_to_model0_op[i])
batch_data_temp = np.concatenate(
[batch_data[counter], batch_data_extra[counter]], 0)
batch_label_temp = np.concatenate(
[batch_label[counter], batch_label_extra[counter]], 0)
batch_smpw_temp = np.concatenate(
[batch_smpw[counter], batch_smpw_extra[counter]], 0)
aug_data = provider.rotate_point_cloud_z(batch_data_temp)
feed_dict = {
ops['pointclouds_pl' + str(i + 1)]: aug_data,
ops['labels_pl' + str(i + 1)]: batch_label_temp,
ops['smpws_pl' + str(i + 1)]: batch_smpw_temp,
ops['data_select']: i,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val = sess.run(
[
ops['merged'], ops['step'], ops['train_op'],
ops['loss'], ops['pred']
],
feed_dict=add_empty(feed_dict, i, ops))
sess.run(copy_from_model0_op[i])
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label_temp)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT[i])
loss_sum += loss_val
counter += 1
if (batch_idx + 1) % 10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx + 1, num_batches))
log_string('mean loss: %f' %
(loss_sum / 10.0 / float(len(NUM_POINT))))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0
def eval_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = len(TEST_DATASET)/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)]
total_correct_vox = 0
total_seen_vox = 0
total_seen_class_vox = [0 for _ in range(NUM_CLASSES)]
total_correct_class_vox = [0 for _ in range(NUM_CLASSES)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT))
labelweights = np.zeros(21)
labelweights_vox = np.zeros(21)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training}
summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'],
ops['loss'], ops['pred']], feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum((pred_val == batch_label) & (batch_label>0) & (batch_smpw>0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_label>0) & (batch_smpw>0))
loss_sum += loss_val
tmp,_ = np.histogram(batch_label,range(22))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label==l) & (batch_smpw>0))
total_correct_class[l] += np.sum((pred_val==l) & (batch_label==l) & (batch_smpw>0))
for b in xrange(batch_label.shape[0]):
_, uvlabel, _ = pc_util.point_cloud_label_to_surface_voxel_label_fast(aug_data[b,batch_smpw[b,:]>0,:], np.concatenate((np.expand_dims(batch_label[b,batch_smpw[b,:]>0],1),np.expand_dims(pred_val[b,batch_smpw[b,:]>0],1)),axis=1), res=0.02)
total_correct_vox += np.sum((uvlabel[:,0]==uvlabel[:,1])&(uvlabel[:,0]>0))
total_seen_vox += np.sum(uvlabel[:,0]>0)
tmp,_ = np.histogram(uvlabel[:,0],range(22))
labelweights_vox += tmp
for l in range(NUM_CLASSES):
total_seen_class_vox[l] += np.sum(uvlabel[:,0]==l)
total_correct_class_vox[l] += np.sum((uvlabel[:,0]==l) & (uvlabel[:,1]==l))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point accuracy vox: %f'% (total_correct_vox / float(total_seen_vox)))
log_string('eval point avg class acc vox: %f' % (np.mean(np.array(total_correct_class_vox[1:])/(np.array(total_seen_class_vox[1:],dtype=np.float)+1e-6))))
log_string('eval point accuracy: %f'% (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(np.array(total_correct_class[1:])/(np.array(total_seen_class[1:],dtype=np.float)+1e-6))))
labelweights_vox = labelweights_vox[1:].astype(np.float32)/np.sum(labelweights_vox[1:].astype(np.float32))
caliweights = np.array([0.388,0.357,0.038,0.033,0.017,0.02,0.016,0.025,0.002,0.002,0.002,0.007,0.006,0.022,0.004,0.0004,0.003,0.002,0.024,0.029])
log_string('eval point calibrated average acc: %f' % (np.average(np.array(total_correct_class[1:])/(np.array(total_seen_class[1:],dtype=np.float)+1e-6),weights=caliweights)))
per_class_str = 'vox based --------'
for l in range(1,NUM_CLASSES):
per_class_str += 'class %d weight: %f, acc: %f; ' % (l,labelweights_vox[l-1],total_correct_class[l]/float(total_seen_class[l]))
log_string(per_class_str)
EPOCH_CNT += 1
return total_correct/float(total_seen)
def eval_all_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(test_data))
num_batches = int(test_idxs.shape[0] / BATCH_SIZE)
#test_idxs = np.arange(0, len(TEST_DATASET))
#current_data = test_data[:,0:NUM_POINT,:]
#current_label = np.squeeze(test_label)
#num_batches = int(current_data.shape[0]/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)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d ALL EVALUATION ----' % (EPOCH_CNT))
labelweights = np.zeros(13)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
#batch_data, batch_label, batch_smpw = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
#batch_data = current_data[start_idx:end_idx,...]
#batch_label = current_label[start_idx:end_idx,...]
batch_data, batch_label, _ = my_get_batch(test_data, test_label,
test_idxs, start_idx,
end_idx)
#batch_smpw = np.ones(batch_label.shape)
batch_smpw = (batch_data[:, :, 0] >= -0.25) & (
batch_data[:, :, 0] <= 0.25) & (batch_data[:, :, 1] >= -0.25) & (
batch_data[:, :, 1] <= 0.25)
batch_smpw = batch_smpw.astype(np.float32)
batch_data[:, :, :3] = provider.rotate_point_cloud_z(
batch_data[:, :, :3])
#aug_data = provider.rotate_point_cloud(batch_data)
bandwidth = BANDWIDTH
feed_dict = {
ops['pointclouds_pl']: batch_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training
}
summary, step, loss_val, pred_val, boundary_loss = sess.run(
[
ops['merged'], ops['step'], ops['loss'], ops['pred'],
ops['boundary_loss']
],
feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum(
(pred_val == batch_label) &
(batch_smpw > 0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_smpw > 0))
loss_sum += loss_val
tmp, _ = np.histogram(batch_label, range(14))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l)
& (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l)
& (batch_smpw > 0))
total_iou_deno_class[l] += np.sum((
(pred_val == l) | (batch_label == l)) & (batch_smpw > 0))
mIoU = np.mean(
np.array(total_correct_class[:]) /
(np.array(total_iou_deno_class[:], dtype=np.float) + 1e-8))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % (mIoU))
log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class[:]) /
(np.array(total_seen_class[:], dtype=np.float) + 1e-8))))
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %d, acc: %f \n' % (
l, total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
EPOCH_CNT += 1
return mIoU
def eval_one_epoch(sess, ops, test_writer):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = int(len(TEST_DATASET)/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)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----' % (EPOCH_CNT))
labelweights = np.zeros(NUM_CLASSES)
for batch_idx in tqdm(range(num_batches),total=num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx+1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx)
batch_data[:, :, :3] = provider.rotate_point_cloud_z(batch_data[:, :, :3])
feed_dict = {ops['pointclouds_pl']: batch_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training}
summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'],
ops['loss'], ops['pred']], feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum((pred_val == batch_label) & (batch_label > 0) & (
batch_smpw > 0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_label > 0) & (batch_smpw > 0))
loss_sum += loss_val
tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l)& (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l)& (batch_label == l) & (batch_smpw > 0))
total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l)) & (batch_smpw > 0))
labelweights = labelweights[1:].astype(
np.float32)/np.sum(labelweights[1:].astype(np.float32))
mIoU = np.mean(np.array(
total_correct_class[1:])/(np.array(total_iou_deno_class[1:], dtype=np.float)+1e-6))
log_string('Eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('Eval point avg class IoU: %f' % (mIoU))
log_string('Eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('Eval point avg class acc: %f' % (np.mean(np.array(
total_correct_class[1:])/(np.array(total_seen_class[1:], dtype=np.float)+1e-6))))
iou_per_class_str = '------- IoU --------\n'
for l in range(1, NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (seg_label_to_cat[l]+' '*(14-len(
seg_label_to_cat[l])), labelweights[l-1], total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
EPOCH_CNT += 1
return mIoU
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('sem_seg')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/s3dis/stanford_indoor3d/'
NUM_CLASSES = 13
NUM_POINT = args.npoint
BATCH_SIZE = args.batch_size
print("start loading training data ...")
TRAIN_DATASET = S3DISDataset(split='train',
data_root=root,
num_point=NUM_POINT,
test_area=args.test_area,
block_size=1.0,
sample_rate=1.0,
transform=None)
print("start loading test data ...")
TEST_DATASET = S3DISDataset(split='test',
data_root=root,
num_point=NUM_POINT,
test_area=args.test_area,
block_size=1.0,
sample_rate=1.0,
transform=None)
trainDataLoader = torch.utils.data.DataLoader(
TRAIN_DATASET,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=10,
pin_memory=True,
drop_last=True,
worker_init_fn=lambda x: np.random.seed(x + int(time.time())))
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=10,
pin_memory=True,
drop_last=True)
weights = torch.Tensor(TRAIN_DATASET.labelweights).cuda()
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET))
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(experiment_dir))
shutil.copy('models/pointnet2_utils.py', str(experiment_dir))
classifier = MODEL.get_model(NUM_CLASSES).cuda()
criterion = MODEL.get_loss().cuda()
classifier.apply(inplace_relu)
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(
str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init)
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
else:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=args.learning_rate,
momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(
m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
global_epoch = 0
best_iou = 0
for epoch in range(start_epoch, args.epoch):
'''Train on chopped scenes'''
log_string('**** Epoch %d (%d/%s) ****' %
(global_epoch + 1, epoch + 1, args.epoch))
lr = max(
args.learning_rate * (args.lr_decay**(epoch // args.step_size)),
LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY
**(epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(
lambda x: bn_momentum_adjust(x, momentum))
num_batches = len(trainDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
classifier = classifier.train()
for i, (points, target) in tqdm(enumerate(trainDataLoader),
total=len(trainDataLoader),
smoothing=0.9):
optimizer.zero_grad()
points = points.data.numpy()
points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3])
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points)
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss.backward()
optimizer.step()
pred_choice = seg_pred.cpu().data.max(1)[1].numpy()
correct = np.sum(pred_choice == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
loss_sum += loss
log_string('Training mean loss: %f' % (loss_sum / num_batches))
log_string('Training accuracy: %f' %
(total_correct / float(total_seen)))
if epoch % 5 == 0:
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
'''Evaluate on chopped scenes'''
with torch.no_grad():
num_batches = len(testDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
labelweights = np.zeros(NUM_CLASSES)
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
classifier = classifier.eval()
log_string('---- EPOCH %03d EVALUATION ----' % (global_epoch + 1))
for i, (points, target) in tqdm(enumerate(testDataLoader),
total=len(testDataLoader),
smoothing=0.9):
points = points.data.numpy()
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points)
pred_val = seg_pred.contiguous().cpu().data.numpy()
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss_sum += loss
pred_val = np.argmax(pred_val, 2)
correct = np.sum((pred_val == batch_label))
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l))
total_iou_deno_class[l] += np.sum(
((pred_val == l) | (batch_label == l)))
labelweights = labelweights.astype(np.float32) / np.sum(
labelweights.astype(np.float32))
mIoU = np.mean(
np.array(total_correct_class) /
(np.array(total_iou_deno_class, dtype=np.float) + 1e-6))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % (mIoU))
log_string('eval point accuracy: %f' %
(total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class) /
(np.array(total_seen_class, dtype=np.float) + 1e-6))))
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' *
(14 - len(seg_label_to_cat[l])), labelweights[l - 1],
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
log_string('Eval mean loss: %f' % (loss_sum / num_batches))
log_string('Eval accuracy: %f' %
(total_correct / float(total_seen)))
if mIoU >= best_iou:
best_iou = mIoU
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'class_avg_iou': mIoU,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
log_string('Best mIoU: %f' % best_iou)
global_epoch += 1
def train_one_epoch(sess, ops, train_writer, epoch=None):
""" ops: dict mapping from string to tf ops """
is_training = True
#current_data, current_label, _ = provider.shuffle_data(train_data[:,0:NUM_POINT,:], train_label)
#current_data = current_data[:2*num_test_data,...]
#current_label = current_label[:2*num_test_data,...]
train_idxs = np.arange(0, len(train_data))
np.random.shuffle(train_idxs)
train_idxs = train_idxs[:2 * num_test_data]
# Shuffle train samples
#train_idxs = np.arange(0, len(TRAIN_DATASET))
#train_idxs = np.arange(0, train_data.shape[0])
#np.random.shuffle(train_idxs)
#num_batches = int(current_data.shape[0]/BATCH_SIZE)
num_batches = int(train_idxs.shape[0] / BATCH_SIZE)
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
total_iou_deno = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
#batch_data, batch_label, batch_smpw = get_batch_wdp(TRAIN_DATASET, train_idxs, start_idx, end_idx)
#batch_data, batch_label, batch_smpw = get_batch(train_data, train_idxs, start_idx, end_idx)
#batch_data = current_data[start_idx:end_idx,...]
#batch_label = current_label[start_idx:end_idx,...]
batch_data, batch_label, batch_smpw_tmp = my_get_batch(
train_data, train_label, train_idxs, start_idx, end_idx)
#batch_smpw = np.ones(batch_label.shape)
batch_smpw = (batch_data[:, :, 0] >= -0.25) & (
batch_data[:, :, 0] <= 0.25) & (batch_data[:, :, 1] >= -0.25) & (
batch_data[:, :, 1] <= 0.25)
batch_smpw = batch_smpw.astype(np.float32)
batch_smpw = batch_smpw * batch_smpw_tmp
# Augment batched point clouds by rotation
batch_data[:, :, :3] = provider.rotate_point_cloud_z(
batch_data[:, :, :3])
#aug_data = provider.rotate_point_cloud(batch_data)
feed_dict = {
ops['pointclouds_pl']: batch_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, pred_val = sess.run([
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['pred']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2)
correct = np.sum(pred_val == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
iou_deno = 0
for l in range(NUM_CLASSES):
iou_deno += np.sum((pred_val == l) | (batch_label == l))
total_iou_deno += iou_deno
loss_sum += loss_val
if (batch_idx + 1) % 10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx + 1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
log_string('total IoU: %f' %
(total_correct / float(total_iou_deno)))
total_correct = 0
total_seen = 0
loss_sum = 0
total_iou_deno = 0
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('part_seg')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/plant_data/'
TRAIN_DATASET = PartNormalDataset(root=root,
npoints=args.npoint,
split='trainval',
normal_channel=args.normal)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
TEST_DATASET = PartNormalDataset(root=root,
npoints=args.npoint,
split='test',
normal_channel=args.normal)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET))
num_classes = 16
num_part = 50
num_part = 3
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(experiment_dir))
shutil.copy('models/pointnet_util.py', str(experiment_dir))
classifier = MODEL.get_model(num_part,
normal_channel=args.normal) #.cuda()
criterion = MODEL.get_loss().cuda()
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
# try:
checkpoint = torch.load(
str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
num_part = 3
seg_classes = {'plant': [0, 1, 2]}
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
classifier.reinit_lastlayer(num_part) # for transfer learning
classifier = classifier.cuda()
# except:
# log_string('No existing model, starting training from scratch...')
# start_epoch = 0
# classifier = classifier.apply(weights_init)
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
else:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=args.learning_rate,
momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(
m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
best_acc = 0
global_epoch = 0
best_class_avg_iou = 0
best_inctance_avg_iou = 0
for epoch in range(start_epoch, args.epoch):
log_string('Epoch %d (%d/%s):' %
(global_epoch + 1, epoch + 1, args.epoch))
'''Adjust learning rate and BN momentum'''
lr = max(
args.learning_rate * (args.lr_decay**(epoch // args.step_size)),
LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
mean_correct = []
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY
**(epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(
lambda x: bn_momentum_adjust(x, momentum))
'''learning one epoch'''
for i, data in tqdm(enumerate(trainDataLoader),
total=len(trainDataLoader),
smoothing=0.9):
points, label, target = data
points = points.data.numpy()
if np.random.random(1)[0] > 0.5:
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :,
0:3])
points[:, :, 0:3] = provider.rotate_point_cloud(
points[:, :, 0:3]) ## added _z
elif np.random.random(1)[0] > 0.5:
points[:, :, 0:3] = provider.rotate_point_cloud_z(points[:, :,
0:3])
# points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :,
0:3])
elif np.random.random(1)[0] > 0.5:
# points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :,
0:3])
# if np.random.random(1)[0] > 0.5:
# if np.random.random(1)[0] > 0.5:
# points[:,:, 0:3] = provider.shift_point_cloud(points[:,:, 0:3])
# else:
# points[:,:, 0:3] = provider.rotate_point_cloud_z(points[:,:, 0:3]) ## added _z
points = torch.Tensor(points)
points, label, target = points.float().cuda(), label.long().cuda(
), target.long().cuda()
points = points.transpose(2, 1)
optimizer.zero_grad()
classifier = classifier.train()
seg_pred, trans_feat = classifier(
points, to_categorical(label, num_classes))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
pred_choice = seg_pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
mean_correct.append(correct.item() /
(args.batch_size * args.npoint))
loss = criterion(seg_pred, target, trans_feat)
loss.backward()
optimizer.step()
train_instance_acc = np.mean(mean_correct)
log_string('Train accuracy is: %.5f' % train_instance_acc)
with torch.no_grad():
test_metrics = {}
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(num_part)]
total_correct_class = [0 for _ in range(num_part)]
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
for batch_id, (points, label,
target) in tqdm(enumerate(testDataLoader),
total=len(testDataLoader),
smoothing=0.9):
cur_batch_size, NUM_POINT, _ = points.size()
points, label, target = points.float().cuda(), label.long(
).cuda(), target.long().cuda()
points = points.transpose(2, 1)
classifier = classifier.eval()
seg_pred, _ = classifier(points,
to_categorical(label, num_classes))
cur_pred_val = seg_pred.cpu().data.numpy()
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros(
(cur_batch_size, NUM_POINT)).astype(np.int32)
target = target.cpu().data.numpy()
for i in range(cur_batch_size):
cat = seg_label_to_cat[target[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 == target)
total_correct += correct
total_seen += (cur_batch_size * NUM_POINT)
for l in range(num_part):
total_seen_class[l] += np.sum(target == l)
total_correct_class[l] += (np.sum((cur_pred_val == l)
& (target == l)))
for i in range(cur_batch_size):
segp = cur_pred_val[i, :]
segl = target[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])
mean_shape_ious = np.mean(list(shape_ious.values()))
test_metrics['accuracy'] = total_correct / float(total_seen)
test_metrics['class_avg_accuracy'] = 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 %f' %
(cat + ' ' * (14 - len(cat)), shape_ious[cat]))
test_metrics['class_avg_iou'] = mean_shape_ious
test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious)
log_string(
'Epoch %d test Accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f'
%
(epoch + 1, test_metrics['accuracy'],
test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou']))
if (test_metrics['inctance_avg_iou'] >= best_inctance_avg_iou):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'train_acc': train_instance_acc,
'test_acc': test_metrics['accuracy'],
'class_avg_iou': test_metrics['class_avg_iou'],
'inctance_avg_iou': test_metrics['inctance_avg_iou'],
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
if test_metrics['class_avg_iou'] > best_class_avg_iou:
best_class_avg_iou = test_metrics['class_avg_iou']
if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou:
best_inctance_avg_iou = test_metrics['inctance_avg_iou']
log_string('Best accuracy is: %.5f' % best_acc)
log_string('Best class avg mIOU is: %.5f' % best_class_avg_iou)
log_string('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou)
global_epoch += 1
def eval_one_epoch(sess, ops, test_writer, out_filename):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = int(len(TEST_DATASET) / 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)]
total_correct_vox = 0
total_seen_vox = 0
total_seen_class_vox = [0 for _ in range(NUM_CLASSES)]
total_correct_class_vox = [0 for _ in range(NUM_CLASSES)]
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----' % (EPOCH_CNT))
labelweights = np.zeros(21)
labelweights_vox = np.zeros(21)
#predlabels = np.zeros(num_batches*BATCH_SIZE)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(
TEST_DATASET, test_idxs, start_idx, end_idx)
aug_data = provider.rotate_point_cloud_z(batch_data)
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['smpws_pl']: batch_smpw,
ops['is_training_pl']: is_training
}
summary, step, loss_val, pred_val = sess.run(
[ops['merged'], ops['step'], ops['loss'], ops['pred']],
feed_dict=feed_dict)
#print(pred_val)
#predlabels[batch_idx] = pred_val
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
#print(pred_val) #predval is all labels for one image
#print(pred_val.shape)
#print(batch_label)
#print(batch_label.shape)
fname = out_filename + str(batch_idx) + "_pred.txt"
gt_fname = out_filename + str(batch_idx) + "_gt.txt"
fout = open(fname, 'w')
gtfout = open(gt_fname, 'w')
for i in range(NUM_POINT):
color = label2color[pred_val[0, i]]
gtcolor = label2color[batch_label[0, i]]
fout.write('%f %f %f %d %d %d\n' %
(batch_data[0, i, 0], batch_data[0, i, 1],
batch_data[0, i, 2], color[0], color[1], color[2]))
gtfout.write(
'%f %f %f %d %d %d\n' %
(batch_data[0, i, 0], batch_data[0, i, 1], batch_data[0, i, 2],
gtcolor[0], gtcolor[1], gtcolor[2]))
#writes ground truth and prediction files by converting each class label into a unique color.
correct = np.sum(
(pred_val == batch_label) & (batch_label > 0) &
(batch_smpw > 0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_label > 0) & (batch_smpw > 0))
loss_sum += loss_val
tmp, _ = np.histogram(batch_label, range(22))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l)
& (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l)
& (batch_smpw > 0))
for b in xrange(batch_label.shape[0]):
_, uvlabel, _ = pc_util.point_cloud_label_to_surface_voxel_label_fast(
aug_data[b, batch_smpw[b, :] > 0, :],
np.concatenate(
(np.expand_dims(batch_label[b, batch_smpw[b, :] > 0], 1),
np.expand_dims(pred_val[b, batch_smpw[b, :] > 0], 1)),
axis=1),
res=0.02)
total_correct_vox += np.sum((uvlabel[:, 0] == uvlabel[:, 1])
& (uvlabel[:, 0] > 0))
total_seen_vox += np.sum(uvlabel[:, 0] > 0)
tmp, _ = np.histogram(uvlabel[:, 0], range(22))
labelweights_vox += tmp
for l in range(NUM_CLASSES):
total_seen_class_vox[l] += np.sum(uvlabel[:, 0] == l)
total_correct_class_vox[l] += np.sum((uvlabel[:, 0] == l)
& (uvlabel[:, 1] == l))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point accuracy vox: %f' %
(total_correct_vox / float(total_seen_vox)))
log_string('eval point avg class acc vox: %f' % (np.mean(
np.array(total_correct_class_vox[1:]) /
(np.array(total_seen_class_vox[1:], dtype=np.float) + 1e-6))))
log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class[1:]) /
(np.array(total_seen_class[1:], dtype=np.float) + 1e-6))))
labelweights_vox = labelweights_vox[1:].astype(np.float32) / np.sum(
labelweights_vox[1:].astype(np.float32))
caliweights = np.array([
0.388, 0.357, 0.038, 0.033, 0.017, 0.02, 0.016, 0.025, 0.002, 0.002,
0.002, 0.007, 0.006, 0.022, 0.004, 0.0004, 0.003, 0.002, 0.024, 0.029
])
#caliweights correspond to the frequency of each class in the data. These particular weights are set up for the scannet data as seen in table 7 at:
#http://openaccess.thecvf.com/content_cvpr_2017/papers/Dai_ScanNet_Richly-Annotated_3D_CVPR_2017_paper.pdf
log_string(
'eval point calibrated average acc: %f' %
(np.average(np.array(total_correct_class[1:]) /
(np.array(total_seen_class[1:], dtype=np.float) + 1e-6),
weights=caliweights)))
per_class_str = 'vox based --------'
for l in range(1, NUM_CLASSES):
per_class_str += 'class %d weight: %f, acc: %f; ' % (
l, labelweights_vox[l - 1],
total_correct_class[l] / float(total_seen_class[l]))
log_string(per_class_str)
EPOCH_CNT += 1
return total_correct / float(total_seen)
def eval_scene_one_epoch(sess, ops, test_writer, num_classes):
""" ops: dict mapping from string to tf ops """
global EPOCH_CNT
is_training = False
test_idxs = np.arange(0, len(TEST_DATASET))
num_batches = int(len(TEST_DATASET) / 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)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
total_cross_matrix = np.zeros((NUM_CLASSES, NUM_CLASSES))
log_string(str(datetime.now()))
log_string('---- EPOCH %03d EVALUATION ----' % (EPOCH_CNT))
labelweights = np.zeros(21)
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
batch_data, batch_label, batch_smpw = get_batch(
TEST_DATASET, test_idxs, start_idx, end_idx)
batch_label_onehot = np.eye(num_classes)[batch_label]
external_batch_scene_encode = np.max(batch_label_onehot, axis=1)
aug_data = provider.rotate_point_cloud_z(batch_data)
#aug_data = provider.rotate_point_cloud(batch_data)
bandwidth = BANDWIDTH
feed_dict = {
ops['pointclouds_pl']: aug_data,
ops['labels_pl']: batch_label,
ops['labels_onehot_pl']: batch_label_onehot,
ops['smpws_pl']: batch_smpw,
ops['external_scene_encode_pl']: external_batch_scene_encode,
ops['is_training_pl']: is_training,
ops['cos_loss_weight']: 1.0
}
summary, step, loss_val, pred_val = sess.run(
[ops['merged'], ops['step'], ops['loss'], ops['pred']],
feed_dict=feed_dict)
test_writer.add_summary(summary, step)
pred_val = np.argmax(pred_val, 2) # BxN
correct = np.sum(
(pred_val == batch_label) & (batch_label > 0) &
(batch_smpw > 0)) # evaluate only on 20 categories but not unknown
total_correct += correct
total_seen += np.sum((batch_label > 0) & (batch_smpw > 0))
loss_sum += loss_val
tmp, _ = np.histogram(batch_label, range(22))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l)
& (batch_smpw > 0))
total_correct_class[l] += np.sum((pred_val == l)
& (batch_label == l)
& (batch_smpw > 0))
total_iou_deno_class[l] += np.sum((
(pred_val == l) | (batch_label == l)) & (batch_smpw > 0))
for m in range(NUM_CLASSES):
total_cross_matrix[l, m] += np.sum((
(pred_val == l) & (batch_label == m)) & (batch_smpw > 0))
mIoU = np.mean(
np.array(total_correct_class[1:]) /
(np.array(total_iou_deno_class[1:], dtype=np.float) + 1e-6))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % (mIoU))
log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (np.mean(
np.array(total_correct_class[1:]) /
(np.array(total_seen_class[1:], dtype=np.float) + 1e-6))))
iou_per_class_str = '------- IoU --------\n'
for l in range(1, NUM_CLASSES):
iou_per_class_str += 'class %d, acc: %f \n' % (
l, total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
EPOCH_CNT += 1
return mIoU, total_cross_matrix
def train_one_epoch(sess, ops, train_writer):
""" ops: dict mapping from string to tf ops """
is_training = True
# Shuffle train samples
train_idxs = np.arange(0, len(TRAIN_DATASET))
np.random.shuffle(train_idxs)
num_batches = len(TRAIN_DATASET) / (BATCH_SIZE // 2)
log_string(str(datetime.now()))
total_correct = 0
total_seen = 0
loss_sum = 0
loss_sum1 = 0
loss_sum2 = 0
for batch_idx in range(num_batches):
start_idx = batch_idx * (BATCH_SIZE // 2)
end_idx = (batch_idx + 1) * (BATCH_SIZE // 2)
### Get input from other process
batch_data, batch_smpw = SUNCG_DATASET.wait_other()
###Convert it to voxel
batch_data_norm = pc_normalize_batch(batch_data)
batch_data_temp = pc_util.point_cloud_label_to_volume_batch_exact(
batch_data_norm, vsize=V_SIZE, flatten=True)
batch_data_vol = np.zeros((BATCH_SIZE, V_SIZE, V_SIZE, V_SIZE, 1))
batch_data_vol[0:BATCH_SIZE // 2, :, :, :, :] = batch_data_temp
feed_dict = {
ops['pointclouds_pl']: batch_data_vol,
ops['is_training_pl']: False
}
pred_val = sess.run(ops['pred'], feed_dict=feed_dict)
pred_val = np.expand_dims(np.argmax(pred_val, 4), -1)
pred_val = pred_val[0:BATCH_SIZE // 2, :, :, :, :]
###Convert it back to pc
pred_val = np.clip(
pc_util.volume_topc_batch_exact(pred_val, batch_data_norm) - 1,
a_min=0,
a_max=None) ### Clip the label in case of Nan in training
batch_data_extra, batch_label_extra, batch_smpw_extra = SUNCG_DATASET.ready(
batch_data, np.squeeze(pred_val), batch_smpw,
TRAIN_DATASET.labelweights)
batch_data, batch_label, batch_smpw = get_batch_wdp(
TRAIN_DATASET, train_idxs, start_idx, end_idx)
batch_data = np.concatenate([batch_data, batch_data_extra], 0)
batch_label = np.concatenate([batch_label, batch_label_extra], 0)
batch_smpw = np.concatenate([batch_smpw, batch_smpw_extra], 0)
# Augment batched point clouds by rotation
aug_data = provider.rotate_point_cloud_z(batch_data)
###Convert it to voxel
aug_data_vol, batch_label_vol, batch_smpw_vol = pc_util.point_cloud_label_to_volume_batch(
pc_normalize_batch(aug_data),
batch_label + 1,
batch_smpw,
vsize=V_SIZE,
flatten=True)
feed_dict = {
ops['pointclouds_pl']: aug_data_vol,
ops['labels_pl']: batch_label_vol,
ops['smpws_pl']: batch_smpw_vol,
ops['is_training_pl']: is_training,
}
summary, step, _, loss_val, loss_val1, loss_val2, pred_val = sess.run(
[
ops['merged'], ops['step'], ops['train_op'], ops['loss'],
ops['loss1'], ops['loss2'], ops['pred']
],
feed_dict=feed_dict)
train_writer.add_summary(summary, step)
### Change the voxel back to pc
pred_val = np.argmax(pred_val, 4)
pred_val, batch_label, batch_smpw, _, _ = pc_util.volume_topc_batch(
pred_val, batch_label_vol, batch_smpw_vol)
for i in range(len(pred_val)):
pred_val[i] -= 1
for i in range(len(batch_label)):
batch_label[i] -= 1
for i in range(len(pred_val)):
correct = np.sum(pred_val[i] == batch_label[i])
total_correct += correct
total_seen += pred_val[i].shape[0]
loss_sum += loss_val
loss_sum1 += loss_val1
loss_sum2 += loss_val2
if (batch_idx + 1) % 10 == 0:
log_string(' -- %03d / %03d --' % (batch_idx + 1, num_batches))
log_string('mean loss: %f' % (loss_sum / 10))
log_string('mean loss1: %f' % (loss_sum1 / 10))
log_string('mean loss2: %f' % (loss_sum2 / 10))
log_string('accuracy: %f' % (total_correct / float(total_seen)))
total_correct = 0
total_seen = 0
loss_sum = 0