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)