def train_epoch(self, epoch):
epoch_start_time = time.time()
loss_buf = []
num_train = len(self.train_loader.dataset)
num_batch = int(num_train / self.batch_size)
log_string("total training nuber: " + str(num_train) +
"total batch number: " + str(num_batch) + " .")
for iter, (pts, _) in enumerate(self.train_loader):
log_string("batch idx: " + str(iter) + "/" + str(num_batch) +
" in " + str(epoch) + "/" + str(self.epochs) +
" epoch...")
if self.gpu_mode:
pts = pts.cuda()
start_idx = (self.batch_size * iter) % num_train
end_idx = min(start_idx + self.batch_size, num_train)
batch_size_train = end_idx - start_idx
if self.dataset_name == 'arch':
if start_idx + batch_size_train == num_train:
if self.is_list_of_h5_list:
filelist_train_prev = self.seg_list[
(self.seg_list_idx - 1) % len(self.seg_list)]
filelist_train = self.seg_list[self.seg_list_idx %
len(self.seg_list)]
if filelist_train != filelist_train_prev:
self.train_loader = get_dataloader(
filelist=filelist_train,
batch_size=self.batch_size,
num_workers=self.workers)
num_train = len(self.train_loader.dataset)
self.seg_list_idx += 1
# forward
self.optimizer.zero_grad()
#input(bs, 2048, 3), output(bs, 2025,3)
output, _, _ = self.model(pts)
#print("input: " + pts.shape + ", output shape: " + output.shape)
loss = self.model.get_loss(pts, output)
# backward
loss.backward()
self.optimizer.step()
loss_buf.append(loss.detach().cpu().numpy())
# finish one epoch
epoch_time = time.time() - epoch_start_time
self.train_hist['per_epoch_time'].append(epoch_time)
self.train_hist['loss'].append(np.mean(loss_buf))
print(
f'Epoch {epoch+1}: Loss {np.mean(loss_buf)}, time {epoch_time:.4f}s'
)
return np.mean(loss_buf)
def main(opt):
experiment_id = 'Semantic_segmentation_'+ opt.encoder +'_' +opt.pre_ae_epochs + '_' + str(opt.feat_dims) + '_' + opt.dataset+'_' + str(opt.percentage)+'_percent'
LOG_FOUT = open(os.path.join(ROOT_DIR, 'LOG', experiment_id+'_train_log.txt'), 'a')
def log_string(out_str):
LOG_FOUT.write(out_str + '\n')
LOG_FOUT.flush()
print(out_str)
snapshot_root = 'snapshot/%s' %experiment_id
tensorboard_root = 'tensorboard/%s' %experiment_id
heatmap_root = 'heatmap/%s' %experiment_id
save_dir = os.path.join(ROOT_DIR, snapshot_root, 'models/')
tboard_dir = os.path.join(ROOT_DIR, tensorboard_root)
hmap_dir = os.path.join(ROOT_DIR, heatmap_root)
#create folder to save trained models
if opt.model == '':
if not os.path.exists(save_dir):
os.makedirs(save_dir)
else:
choose = input("Remove " + save_dir + " ? (y/n)")
if choose == 'y':
shutil.rmtree(save_dir)
os.makedirs(save_dir)
else:
sys.exit(0)
if not os.path.exists(tboard_dir):
os.makedirs(tboard_dir)
else:
shutil.rmtree(tboard_dir)
os.makedirs(tboard_dir)
if not os.path.exists(hmap_dir):
os.makedirs(hmap_dir)
else:
shutil.rmtree(hmap_dir)
os.makedirs(hmap_dir)
writer = SummaryWriter(log_dir = tboard_dir)
#generate part label one-hot correspondence from the catagory:
if opt.dataset == 'arch':
if opt.no_others:
class2label = {"arch":0, "column":1, "moldings":2, "floor":3, "door_window":4, "wall":5, "stairs":6, "vault":7, "roof":8}
else:
class2label = {"arch":0, "column":1, "moldings":2, "floor":3, "door_window":4, "wall":5, "stairs":6, "vault":7, "roof":8, "other":9}
seg_classes = class2label
seg_label_to_cat = {}
for i,cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
elif opt.dataset == 'arch_scene_2':
class2label = {"arch":0, "column":1, "moldings":2, "floor":3, "door_window":4, "wall":5, "stairs":6, "vault":7}
seg_classes = class2label
seg_label_to_cat = {}
for i,cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
# load the dataset
log_string('-Preparing dataset...')
data_resized=False
#train_path = os.path.join(ROOT_DIR, 'cache', 'latent_' + opt.pre_ae_epochs + '_' + opt.dataset + '_' +str(opt.feature_dims), 'features')
if opt.dataset == 'arch':
if opt.no_others:
NUM_CLASSES = 9
arch_data_dir = opt.folder if opt.folder else 'arch_no_others_1.0m_pointnet_hdf5_data'
else:
NUM_CLASSES = 10
arch_data_dir = opt.folder if opt.folder else "arch_pointcnn_hdf5_2048"
elif opt.dataset == 'arch_scene_2':
NUM_CLASSES = 8
arch_data_dir = opt.folder if opt.folder else "scene_2_1.0m_pointnet_hdf5_data"
train_filelist = os.path.join(DATA_DIR, arch_data_dir, "train_data_files.txt")
val_filelist = os.path.join(DATA_DIR, arch_data_dir, "val_data_files.txt")
# load training data
train_dataset = arch_dataloader.get_dataloader(filelist=train_filelist, is_rotated=False, split='train', batch_size=opt.batch_size, num_workers=4, num_points=opt.num_points, num_dims=opt.num_dims, group_shuffle=False, random_translate=opt.use_translate, shuffle=False, drop_last=True)
log_string("classifer set size: " + str(train_dataset.dataset.__len__()))
val_dataset = arch_dataloader.get_dataloader(filelist=val_filelist, is_rotated=False, split='train', num_points=opt.num_points, batch_size=opt.batch_size, num_workers=4, num_dims=opt.num_dims, group_shuffle=False, shuffle=False, random_translate=opt.use_translate, drop_last=False)
log_string("classifer set size: " + str(val_dataset.dataset.__len__()))
# load the model for point auto encoder
if opt.num_dims == 3:
ae_net = MultiTaskNet(opt)
if opt.ae_model != '':
ae_net = load_pretrain(ae_net, os.path.join(ROOT_DIR, opt.ae_model))
if opt.gpu_mode:
print("Let's use", torch.cuda.device_count(), "GPUs!")
ae_net = ae_net.cuda()
ae_net=ae_net.eval()
#initial segmentation model
sem_seg_net = SemSegNet(num_class=opt.n_classes, encoder=opt.encoder, symmetric_function=opt.symmetric_function, dropout=opt.dropout)
#load pretrained model
if opt.model != '':
sem_seg_net = load_pretrain(sem_seg_net, opt.model)
# load model to gpu
if opt.gpu_mode:
sem_seg_net = sem_seg_net.cuda()
# initialize optimizer
optimizer = optim.Adam([{'params': sem_seg_net.parameters(), 'initial_lr': 1e-4}], lr=0.01, weight_decay=1e-6)
scheduler = optim.lr_scheduler.StepLR(optimizer, 20, 0.5, opt.n_epochs)
# start training
n_batch = 0
# start epoch index
if opt.model != '':
start_epoch = opt.model[-7:-4]
if start_epoch[0] == '_':
start_epoch = start_epoch[1:]
start_epoch=int(start_epoch)
else:
start_epoch = 0
log_string('training start!!!')
start_time = time.time()
total_time = []
best_iou = 0
for epoch in range(start_epoch, opt.n_epochs):
train_acc_epoch, train_iou_epoch, test_acc_epoch, test_iou_epoch = [], [], [], []
loss_buf = []
sem_seg_net=sem_seg_net.train()
for iter, data in enumerate(train_dataset):
points, target = data
# use the pre-trained AE to encode the point cloud into latent features
points_ = Variable(points)
target = target.long()
if opt.gpu_mode:
points_ = points_.cuda()
_, _, latent_caps, mid_features = ae_net(points_) # (batch_size, emb_dims*2), (batch_size, 64*3, num_points)
con_code = torch.cat([latent_caps.view(-1,opt.feat_dims*opt.symmetric_function,1).repeat(1,1,opt.num_points), mid_features],1).cpu().detach().numpy()
latent_caps = torch.from_numpy(con_code).float()
if(latent_caps.size(0)<opt.batch_size):
continue
latent_caps, target = Variable(latent_caps), Variable(target)
if opt.gpu_mode:
latent_caps,target = latent_caps.cuda(), target.cuda()
# forward
optimizer.zero_grad()
#latent_caps=latent_caps.transpose(2, 1)# consider the feature vector size as the channel in the network
output_digit =sem_seg_net(latent_caps)
output_digit = output_digit.view(-1, opt.n_classes)
#batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy()
target= target.view(-1,1)[:,0]
train_loss = F.nll_loss(output_digit, target)
train_loss.backward()
optimizer.step()
pred_choice = output_digit.data.cpu().max(1)[1]
correct = pred_choice.eq(target.data.cpu()).cpu().sum()
train_acc_epoch.append(correct.item() / float(opt.batch_size*opt.num_points))
train_iou_epoch.append(correct.item() / float(2*opt.batch_size*opt.num_points-correct.item()))
n_batch= train_dataset.dataset.__len__() // opt.batch_size
loss_buf.append(train_loss.detach().cpu().numpy())
log_string('[%d: %d/%d] | train loss: %f | train accuracy: %f | train iou: %f' %(epoch+1, iter+1, n_batch, np.mean(loss_buf), correct.item()/float(opt.batch_size * opt.num_points),correct.item() / float(2*opt.batch_size*opt.num_points-correct.item())))
#update lr
if optimizer.param_groups[0]['lr'] > 1e-5:
scheduler.step()
if optimizer.param_groups[0]['lr'] < 1e-5:
for param_group in optimizer.param_groups:
param_group['lr'] = 1e-5
# save tensorboard
total_time.append(time.time()-start_time)
print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(total_time), epoch+1, total_time[0]))
writer.add_scalar('Train Loss', np.mean(loss_buf), epoch)
writer.add_scalar('Train Accuracy', np.mean(train_acc_epoch), epoch)
writer.add_scalar('Train IoU', np.mean(train_iou_epoch), epoch)
log_string('---- EPOCH %03d TRAIN ----' % (epoch + 1))
log_string('epoch %d | mean train accuracy: %f | mean train IoU: %f' %(epoch+1, np.mean(train_acc_epoch), np.mean(train_iou_epoch)))
if (epoch+1) % opt.snapshot_interval == 0 or epoch == 0:
_snapshot(save_dir, sem_seg_net, epoch + 1, opt)
with torch.no_grad():
total_correct = 0
total_seen = 0
loss_sum = 0
n_batch= val_dataset.dataset.__len__()
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)]
log_string('---- EPOCH %03d EVALUATION ----' % (epoch + 1))
#sem_seg_net=sem_seg_net.eval()
for iter, data in enumerate(val_dataset):
points, target = data
# use the pre-trained AE to encode the point cloud into latent capsules
points_ = Variable(points)
target = target.long()
if opt.gpu_mode:
points_ = points_.cuda()
_, _, latent_caps, mid_features = ae_net(points_)
con_code = torch.cat([latent_caps.view(-1,opt.feat_dims*opt.symmetric_function,1).repeat(1,1,opt.num_points), mid_features],1).cpu().detach().numpy()
latent_caps = torch.from_numpy(con_code).float()
if(latent_caps.size(0)<opt.batch_size):
continue
latent_caps, target = Variable(latent_caps), Variable(target)
if opt.gpu_mode:
latent_caps,target = latent_caps.cuda(), target.cuda()
batch_label = target.cpu().data.numpy() #(8, 4096)
#output
sem_seg_net = sem_seg_net.eval()
output=sem_seg_net(latent_caps) #([8, 4096, n_classes])
#output to prediction
pred_val = output.contiguous().cpu().data.numpy()
pred_val = np.argmax(pred_val, 2) #(8, 4096)
#convert output to calculate loss
output_digit = output.view(-1, opt.n_classes) #([32768, 10])
target= target.view(-1,1)[:,0] #[32768]
#print("target shape: " + str(target.shape))
loss = F.nll_loss(output_digit, target)
loss_sum +=loss
correct = np.sum((pred_val == batch_label))
total_correct += correct
total_seen += (opt.batch_size * opt.num_points)
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('epoch %d | eval mean loss: %f' % (epoch+1,loss_sum / float(n_batch)))
log_string('epoch %d | eval point avg class IoU: %f' % (epoch+1, mIoU))
log_string('epoch %d | eval point accuracy: %f' % (epoch+1, total_correct / float(total_seen)))
log_string('epoch %d | eval point avg class acc: %f' % (epoch+1,
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 += 'epoch %d | class %s weight: %.3f, IoU: %.3f \n' % (epoch+1,
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)
writer.add_scalar('Eval Loss', loss_sum / float(n_batch), epoch)
writer.add_scalar('Eval mIoU', mIoU, epoch)
writer.add_scalar('Eval Point accuracy', total_correct / float(total_seen), epoch)
if mIoU >= best_iou:
best_iou = mIoU
_snapshot(save_dir, sem_seg_net, 'best', opt)
log_string('Saving model....')
log_string('Best mIoU: %f at epoch %d' % (best_iou, epoch+1))
log_string("Training finish!... save training results")
def main(opt):
USE_CUDA = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
label2color = {
0: [255, 255, 255], # white
1: [0, 0, 255], # blue
2: [128, 0, 0], # maroon
3: [255, 0, 255], # fuchisia
4: [0, 128, 0], # green
5: [255, 0, 0], # red
6: [128, 0, 128], # purple
7: [0, 0, 128], # navy
8: [128, 128, 0], # olive
9: [128, 128, 128]
}
experiment_id = opt.seg_model.split('/')[-3] + '_test'
n_epoch = opt.seg_model.split('/')[-1][-7:-4]
LOG_FOUT = open(
os.path.join(ROOT_DIR, 'LOG',
experiment_id + '_at_epoch_' + n_epoch + '_log.txt'), 'a')
def log_string(out_str):
LOG_FOUT.write(out_str + '\n')
LOG_FOUT.flush()
print(out_str)
output_root = 'output/%s' % experiment_id
save_dir = os.path.join(ROOT_DIR, output_root, 'at_' + n_epoch)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
if opt.visual:
if opt.dataset == 'arch':
fout = open(os.path.join(save_dir, 'Scene_A_pred.obj'), 'w')
fout_gt = open(os.path.join(save_dir, 'Scene_A_gt.obj'), 'w')
elif opt.dataset == 'arch_scene_2':
fout = open(os.path.join(save_dir, 'Scene_2_pred.obj'), 'w')
fout_gt = open(os.path.join(save_dir, 'Scene_2_gt.obj'), 'w')
# generate part label one-hot correspondence from the catagory:
if opt.dataset == 's3dis':
classes = [
'ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter'
]
class2label = {cls: i for i, cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
elif opt.dataset == 'arch':
if opt.no_others:
class2label = {
"arch": 0,
"column": 1,
"moldings": 2,
"floor": 3,
"door_window": 4,
"wall": 5,
"stairs": 6,
"vault": 7,
"roof": 8
}
else:
class2label = {
"arch": 0,
"column": 1,
"moldings": 2,
"floor": 3,
"door_window": 4,
"wall": 5,
"stairs": 6,
"vault": 7,
"roof": 8,
"other": 9
}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
elif opt.dataset == 'arch_scene_2':
class2label = {
"arch": 0,
"column": 1,
"moldings": 2,
"floor": 3,
"door_window": 4,
"wall": 5,
"stairs": 6,
"vault": 7
}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
# load the model for point auto encoder
if opt.num_dims == 3:
ae_net = DGCNN_FoldNet(opt)
else:
ae_net = DGCNN_FoldNet_feat(opt)
if opt.ae_model != '':
ae_net = load_pretrain(ae_net, os.path.join(ROOT_DIR, opt.ae_model))
if USE_CUDA:
print("Let's use", torch.cuda.device_count(), "GPUs!")
ae_net = ae_net.cuda()
ae_net = ae_net.eval()
# load the model for capsule wised part segmentation
if opt.feat_dims == 512:
sem_seg_net = SemSegNet(num_class=opt.n_classes,
encoder=opt.encoder,
dropout=opt.dropout,
feat_dims=True,
with_rgb=False)
elif opt.feat_dims == 1024:
sem_seg_net = SemSegNet(num_class=opt.n_classes,
encoder=opt.encoder,
dropout=opt.dropout)
if opt.seg_model != '':
sem_seg_net = load_pretrain(sem_seg_net,
os.path.join(ROOT_DIR, opt.seg_model))
if USE_CUDA:
sem_seg_net = sem_seg_net.cuda()
sem_seg_net = sem_seg_net.eval()
# load dataset
if opt.dataset == 's3dis':
print('-Preparing Loading s3dis test dataset...')
root = os.path.join(DATA_DIR, 'stanford_indoor3d')
NUM_CLASSES = 13
NUM_POINT = opt.num_points
BATCH_SIZE = opt.batch_size
dataset = S3DISDataset(split='test',
data_root=root,
num_point=NUM_POINT,
rgb=False,
test_area=5,
block_size=1.0,
sample_rate=1.0,
transform=None)
log_string("start loading test data ...")
dataLoader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=4,
pin_memory=True,
drop_last=True,
worker_init_fn=lambda x: np.random.seed(x + int(time.time())))
log_string("classifer set size: " + str(dataloader.dataset.__len__()))
elif opt.dataset == 'arch':
print('-Preparing Loading ArCH evaluation dataset...')
if opt.no_others:
NUM_CLASSES = 9
arch_dir = opt.folder if opt.folder else 'arch_no_others_1.0m_pointnet_hdf5_data'
else:
NUM_CLASSES = 10
arch_dir = opt.folder if opt.folder else "arch_pointcnn_hdf5_2048"
log_string('-Now loading test ArCH dataset...')
filelist = os.path.join(DATA_DIR, arch_dir, "test_data_files.txt")
# load test data
dataloader = arch_dataloader.get_dataloader(
filelist=filelist,
num_dims=opt.num_dims,
batch_size=opt.batch_size,
num_points=opt.num_points,
num_workers=4,
group_shuffle=False,
shuffle=False,
random_rotate=False,
random_jitter=False,
random_translate=opt.use_translate,
drop_last=False)
log_string("classifer set size: " + str(dataloader.dataset.__len__()))
elif opt.dataset == 'arch_scene_2':
NUM_CLASSES = 8
print('-Preparing Loading ArCH evaluation dataset...')
log_string('-Now loading test ArCH dataset...')
filelist = os.path.join(DATA_DIR, "scene_2_1.0m_pointnet_hdf5_data",
"test_data_files.txt")
# load test data
dataloader = arch_dataloader.get_dataloader(
filelist=filelist,
batch_size=opt.batch_size,
num_points=opt.num_points,
num_workers=4,
group_shuffle=False,
shuffle=False,
random_rotate=False,
random_jitter=False,
random_translate=opt.use_translate,
drop_last=False)
log_string("classifer set size: " + str(dataloader.dataset.__len__()))
correct_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)]
points_collector = []
pd_labels_collector = []
gt_labels_collector = []
for batch_id, data in enumerate(dataloader):
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)]
points, target = data
if (points.size(0) < opt.batch_size):
break
# use the pre-trained AE to encode the point cloud into latent capsules
points_ = Variable(points)
target = target.long()
if USE_CUDA:
points_ = points_.cuda()
_, latent_caps, mid_features = ae_net(points_)
con_code = torch.cat([
latent_caps.view(-1, opt.feat_dims, 1).repeat(
1, 1, opt.num_points), mid_features
], 1).cpu().detach().numpy()
latent_caps = torch.from_numpy(con_code).float()
# predict the part class per capsule
latent_caps, target = Variable(latent_caps), Variable(target)
if USE_CUDA:
latent_caps, target = latent_caps.cuda(), target.cuda()
output = sem_seg_net(latent_caps)
output_digit = output.view(-1, opt.n_classes)
#batch_label = target.cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
pred_choice = output_digit.data.cpu().max(1)[1]
correct = pred_choice.eq(target.data.cpu()).cpu().sum()
correct_sum = correct_sum + correct.item()
target = target.cpu().data.numpy()
pred_choice = pred_choice.cpu().data.numpy()
# calculate the accuracy with the GT
for l in range(NUM_CLASSES):
total_seen_class_tmp[l] += np.sum((target == l))
total_correct_class_tmp[l] += np.sum((pred_choice == l)
& (target == l))
total_iou_deno_class_tmp[l] += np.sum(
((pred_choice == l) | (target == 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 batch %d in Scene_A: %.4f' % (batch_id+1,tmp_iou))
#log_string('----------------------------')
points_collector.extend(points.reshape((-1, 3)).numpy())
pd_labels_collector.extend(pred_choice)
gt_labels_collector.extend(target)
if opt.visual:
log_string('Writing results...')
sparse_labels = np.array(pd_labels_collector).astype(int).flatten()
if opt.dataset == 'arch':
pd_label_filename = save_dir + '/Scene_A_pd_labels.txt'
gt_label_filename = save_dir + '/Scene_A_gt_labels.txt'
elif opt.dataset == 'arch_scene_2':
pd_label_filename = save_dir + '/Scene_2_pd_labels.txt'
gt_label_filename = save_dir + '/Scene_2_gt_labels.txt'
#np.savetxt(pd_label_filename, sparse_labels, fmt='%d', delimiter='\n')
#log_string("Exported sparse labels to {}".format(pd_label_filename))
gt_labels = np.array(gt_labels_collector).astype(int).flatten()
#np.savetxt(gt_label_filename, gt_labels, fmt='%d', delimiter='\n')
#log_string("Exported sparse labels to {}".format(gt_label_filename))
#print(points_collector.size())
sparse_points = np.array(points_collector).reshape((-1, 3))
for i in range(gt_labels.shape[0]):
color = label2color[sparse_labels[i]]
color_gt = label2color[gt_labels[i]]
if opt.visual:
fout.write('v %f %f %f %d %d %d\n' %
(sparse_points[i, 0], sparse_points[i, 1],
sparse_points[i, 2], color[0], color[1], color[2]))
fout_gt.write('v %f %f %f %d %d %d\n' %
(sparse_points[i, 0], sparse_points[i, 1],
sparse_points[i, 2], color_gt[0], color_gt[1],
color_gt[2]))
log_string("Exported sparse pcd to {}".format(save_dir))
if opt.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] + ' ' * (12 - len(seg_label_to_cat[l])),
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
log_string('test point avg class IoU: %f' % np.mean(IoU))
log_string('test 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(
'test whole scene point accuracy: %f' %
(np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6)))
log_string('test mIoU: %f' % np.mean(
np.array(total_correct_class) /
(np.array(total_iou_deno_class, dtype=np.float) + 1e-6)))
log_string("Done!")
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--folder', '-f', help='Path to data folder')
parser.add_argument('--split',
'-s',
help='Path to data folder',
default='train')
parser.add_argument('--outpath', '-o', help='Path to output data folder')
parser.add_argument('--max_point_num',
'-m',
help='Max point number of each sample',
type=int,
default=2048)
parser.add_argument('--batch_size',
'-b',
help='Batch_size',
type=int,
default=32)
parser.add_argument('--NUM_ANGLES',
'-a',
help='Rotation angle number',
type=int,
default=4)
args = parser.parse_args()
print(args)
#read data
root = os.path.join(
ROOT_DIR, 'data', args.folder) if args.folder else os.path.join(
ROOT_DIR, 'data', 'arch3_no_others_combined_5m_4096')
max_point_num = args.max_point_num
split = args.split
train_filelist = os.path.join(root, split + '_data_files.txt')
train_dataset = arch_dataloader.get_dataloader(filelist=train_filelist,
num_points=max_point_num,
num_dims=3,
split=split,
batch_size=args.batch_size,
num_workers=8,
drop_last=True)
print("classifer set size: " + str(train_dataset.dataset.__len__()))
#output settings
output_dir = os.path.join(ROOT_DIR, 'data',
args.outpath) if args.outpath else os.path.join(
ROOT_DIR, 'data', 'rotated_%d_angle_%d' %
(args.NUM_ANGLES, max_point_num))
h5_file_size = 2112
NUM_ANGLES = args.NUM_ANGLES
batch_size = args.batch_size * args.NUM_ANGLES
batch_num = h5_file_size / batch_size
rotated_dataset = np.zeros((h5_file_size, max_point_num, 3))
rotated_labelset = np.zeros((h5_file_size), dtype=np.int32)
idx = 0
idx_h5 = 0
#rotate data
for iter, data in enumerate(train_dataset):
input_points, _ = data
rotated_data, rotated_label = rotate_pc(current_data=input_points,
NUM_CLASSES=args.NUM_ANGLES)
start = idx * batch_size
end = (idx + 1) * batch_size
rotated_dataset[start:end, ...] = rotated_data
rotated_labelset[start:end, ...] = rotated_label
idx += 1
if (idx == batch_num):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
filename_h5 = os.path.join(output_dir, '%d.h5' % idx_h5)
print('{}-Saving {}...'.format(datetime.now(), filename_h5))
file = h5py.File(filename_h5, 'w')
file.create_dataset('data', data=rotated_dataset[0:end, ...])
file.create_dataset('label', data=rotated_labelset[0:end, ...])
file.close()
idx_h5 += 1
idx = 0
def main(args):
USE_CUDA = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ae_net = DGCNN_FoldNet(args)
if args.model != '':
capsule_net = load_pretrain(ae_net, os.path.join(ROOT_DIR, args.model))
if USE_CUDA:
print("Let's use", torch.cuda.device_count(), "GPUs!")
ae_net.cuda()
if args.dataset=='s3dis':
log_string('-Preparing Loading s3dis evaluation dataset...')
classes = ['ceiling','floor','wall','beam','column','window','door','table','chair','sofa','bookcase','board','clutter']
class2label = {cls: i for i,cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i,cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
if args.save_training:
log_string('-Now loading s3dis training classifer dataset...')
split='train'
else:
log_string('-Now loading test s3dis dataset...')
split='test'
root = '../data/stanford_indoor3d/'
NUM_CLASSES = 13
NUM_POINT = args.num_points
BATCH_SIZE = args.batch_size
dataset = S3DISDataset(split=split, data_root=root, num_point=NUM_POINT, rgb=False, test_area=5, block_size=1.0, sample_rate=1.0, transform=None)
log_string("start loading test data ...")
dataLoader = torch.utils.data.DataLoader(dataset, batch_size=args.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())))
log_string("classifer set size: " + dataloader.dataset.__len__())
elif args.dataset == 'arch':
log_string('-Preparing Loading ArCH evaluation dataset...')
data_root = os.path.join(ROOT_DIR, 'data')
if args.save_training:
log_string('-Now loading ArCH training classifer dataset...')
filelist = os.path.join(data_root, 'arch_pointcnn_hdf5_2048', "train_data_files.txt")
else:
log_string('-Now loading test ArCH dataset...')
filelist = os.path.join(data_root, 'arch_pointcnn_hdf5_2048', "test_data_files.txt")
# load training data
dataloader = get_dataloader(filelist=filelist, batch_size=args.batch_size,
num_workers=4, group_shuffle=False,shuffle=True)
log_string("segmentation training dataset size: " + str(dataloader.dataset.__len__()))
# init saving process
#pcd = PointCloud()
data_size=0
dataset_main_path=os.path.abspath(os.path.join(ROOT_DIR, 'cache'))
experiment_name = 'latent_' + args.pre_ae_epochs + '_' + args.dataset + '_' + str(args.latent_vec_size)
out_file_path=os.path.join(dataset_main_path, experiment_name, 'features')
if not os.path.exists(out_file_path):
os.makedirs(out_file_path);
# process for 'shapenet_part' or 'shapenet_core13'
ae_net.eval()
n = 0
feature_set = []
label_set = []
for batch_id, data in enumerate(dataloader):
#for batch_id, data in (enumerate(dataloader), total=len(dataloader)):
points, sem_label= data
if(points.size(0)<args.batch_size):
break
points = Variable(points)
#points = points.transpose(2, 1)
if USE_CUDA:
points = points.cuda()
sem_label = sem_label.cuda()
_, code, mid_features = ae_net(points)
con_code = torch.cat([code.view(-1,args.latent_vec_size,1).repeat(1,1,args.num_points), mid_features],1)
# write the output latent and cls into file
feature_set.append(con_code.transpose(2,1).detach().cpu().numpy())
label_set.append(sem_label.cpu().detach().numpy())
if ((batch_id+1)*args.batch_size % 256) == 0 or (batch_id+1)==len(dataloader):
feature_set = np.concatenate(feature_set, axis=0)
label_set = np.concatenate(label_set, axis=0)
if args.save_training:
out_file_name=out_file_path+"/saved_train_with_sem_label_"+ str(n)+".h5"
else:
out_file_name=out_file_path+"/saved_test_with_sem_label_"+ str(n)+".h5"
if os.path.exists(out_file_name):
os.remove(out_file_name)
fw = h5py.File(out_file_name, 'w', libver='latest')
fw['data'] = feature_set
fw['label_seg'] = label_set
fw.close()
n+=1
log_string('accumalate of batch %d, and datasize is %s, label size is: %s ' % ((batch_id), str(feature_set.shape), str(label_set.shape)))
print('latent set {n} for segmentation saved')
feature_set = []
label_set = []
print("finish generating latent set for SVM.")
def __init__(self, args):
self.dataset_name = args.dataset
self.epochs = args.epochs
self.batch_size = args.batch_size
#self.data_dir = os.path.join(ROOT_DIR, 'data')
self.snapshot_interval = args.snapshot_interval
self.gpu_mode = args.gpu_mode
self.model_path = args.model_path
self.split = args.split
self.num_workers = args.num_workers
# create output directory and files
file = [f for f in args.model_path.split('/')]
if args.experiment_name != None:
self.experiment_id = 'Reconstruct_' + args.experiment_name
elif file[-2] == 'models':
self.experiment_id = file[-3]
else:
self.experiment_id = "Reconstruct" + time.strftime('%m%d%H%M%S')
snapshot_root = 'snapshot/%s' % self.experiment_id
tensorboard_root = 'tensorboard/%s' % self.experiment_id
self.save_dir = os.path.join(ROOT_DIR, snapshot_root, 'models/')
self.tboard_dir = os.path.join(ROOT_DIR, tensorboard_root)
#chenck arguments
if self.model_path == '':
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
else:
choose = input("Remove " + self.save_dir + " ? (y/n)")
if choose == 'y':
shutil.rmtree(self.save_dir)
os.makedirs(self.save_dir)
else:
sys.exit(0)
if not os.path.exists(self.tboard_dir):
os.makedirs(self.tboard_dir)
else:
shutil.rmtree(self.tboard_dir)
os.makedirs(self.tboard_dir)
sys.stdout = Logger(os.path.join(ROOT_DIR, 'LOG',
'ae_network_log.txt'))
self.writer = SummaryWriter(log_dir=self.tboard_dir)
print(str(args))
# initial dataset by dataloader
print('-Preparing dataset...')
if self.dataset_name == 'arch':
# initial dataset filelist
print('-Preparing dataset file list...')
if self.split == 'train':
filelist = os.path.join(DATA_DIR, 'arch_pointcnn_hdf5_4096',
"train_data_files.txt")
else:
filelist = os.path.join(DATA_DIR, 'arch_pointcnn_hdf5_4096',
"test_data_files.txt")
self.is_list_of_h5_list = not is_h5_list(filelist)
if self.is_list_of_h5_list:
self.seg_list = load_seg_list(filelist)
self.seg_list_idx = 0
filepath = self.seg_list[self.seg_list_idx]
self.seg_list_idx += 1
else:
filepath = filelist
print('-Now loading ArCH dataset...')
self.train_loader = get_dataloader(
filelist=filepath,
batch_size=args.batch_size,
num_workers=args.workers,
group_shuffle=True,
random_rotate=args.use_rotate,
random_jitter=args.use_jitter,
random_translate=args.use_translate,
shuffle=True)
print("training set size: ", self.train_loader.dataset.__len__())
elif self.dataset_name == 'all_arch':
# initial dataset filelist
if args.no_others:
arch_data_dir = args.folder if args.folder else 'arch_no_others_pointcnn_hdf5_' + str(
args.num_points)
else:
arch_data_dir = args.folder if args.folder else 'arch_pointcnn_hdf5_' + str(
args.num_points)
print('-Preparing dataset file list...')
filelist = os.path.join(DATA_DIR, arch_data_dir,
"train_data_files.txt")
print('-Now loading ArCH dataset...')
self.train_loader = get_dataloader(
filelist=filelist,
batch_size=args.batch_size,
num_workers=args.workers,
num_points=args.num_points,
group_shuffle=False,
random_rotate=args.use_rotate,
random_jitter=args.use_jitter,
random_translate=args.use_translate,
shuffle=True,
drop_last=True)
print("training set size: ", self.train_loader.dataset.__len__())
elif self.dataset_name == 'shapenetcorev2':
print('-Loading ShapeNetCore dataset...')
self.train_loader = get_shapenet_dataloader(
root=DATA_DIR,
dataset_name=self.dataset_name,
split='train',
batch_size=args.batch_size,
num_workers=args.workers,
num_points=args.num_points,
shuffle=True,
random_translate=args.use_translate,
random_rotate=args.use_rotate,
random_jitter=args.use_jitter)
print("training set size: ", self.train_loader.dataset.__len__())
#initial model
if args.num_dims == 3:
self.model = DGCNN_FoldNet(args)
else:
self.model = DGCNN_FoldNet_feat(args)
#load pretrained model
if args.model_path != '':
self._load_pretrain(args.model_path)
# load model to gpu
if self.gpu_mode:
self.model = self.model.cuda()
# initialize optimizer
self.parameter = self.model.parameters()
self.optimizer = optim.Adam(self.parameter,
lr=0.0001 * 16 / args.batch_size,
betas=(0.9, 0.999),
weight_decay=1e-6)