def train(args): # Prepare inputs print('{}-Preparing datasets...'.format(datetime.now())) is_list_of_h5_list = not data_utils.is_h5_list(args.filelist) if is_list_of_h5_list: seg_list = data_utils.load_seg_list(args.filelist) seg_list_idx = 0 filelist_train = seg_list[seg_list_idx] seg_list_idx = seg_list_idx + 1 else: filelist_train = args.filelist data_train, labels, data_num_train, label_train, _ = data_utils.load_seg( filelist_train) print("Done", data_train.shape) THREADS = 4 BATCH_SIZE = args.batchsize USE_CUDA = True N_CLASSES = 50 EPOCHS = 200 MILESTONES = [60, 120, 180] print("Creating network...") net = Net(input_channels=1, output_channels=N_CLASSES) net.cuda() print("parameters", count_parameters(net)) ds = PartNormalDataset(data_train, data_num_train, label_train, net.config, npoints=args.npoints, shape_labels=labels) train_loader = torch.utils.data.DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=THREADS, collate_fn=tree_collate) optimizer = torch.optim.Adam(net.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, MILESTONES) # create the model folder time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S') root_folder = os.path.join( args.savedir, "Net_b{}_pts{}_{}".format(args.batchsize, args.npoints, time_string)) os.makedirs(root_folder, exist_ok=True) # create the log file logs = open(os.path.join(root_folder, "log.txt"), "w") for epoch in range(EPOCHS): scheduler.step() cm = np.zeros((N_CLASSES, N_CLASSES)) t = tqdm(train_loader, ncols=120, desc="Epoch {}".format(epoch)) for pts, features, seg, tree, labels in t: if USE_CUDA: features = features.cuda() pts = pts.cuda() for l_id in range(len(tree)): tree[l_id]["points"] = tree[l_id]["points"].cuda() tree[l_id]["indices"] = tree[l_id]["indices"].cuda() seg = seg.cuda() optimizer.zero_grad() outputs = net(features, pts, tree) loss = F.cross_entropy(outputs.view(-1, N_CLASSES), seg.view(-1)) loss.backward() optimizer.step() output_np = np.argmax(outputs.cpu().detach().numpy(), axis=2).copy() target_np = seg.cpu().numpy().copy() cm_ = confusion_matrix(target_np.ravel(), output_np.ravel(), labels=list(range(N_CLASSES))) cm += cm_ oa = "{:.3f}".format(metrics.stats_overall_accuracy(cm)) aa = "{:.3f}".format(metrics.stats_accuracy_per_class(cm)[0]) t.set_postfix(OA=oa, AA=aa) # save the model torch.save(net.state_dict(), os.path.join(root_folder, "state_dict.pth")) # write the logs logs.write("{} {} {} \n".format(epoch, oa, aa)) logs.flush() logs.close()
def test_multiple(args): THREADS = 4 BATCH_SIZE = args.batchsize USE_CUDA = True N_CLASSES = 50 args.data_folder = os.path.join(args.rootdir, "test_data") # create the output folders output_folder = args.savedir + '_predictions_multi_{}'.format(args.ntree) category_list = [ (category, int(label_num)) for (category, label_num) in [line.split() for line in open(args.category, 'r')] ] offset = 0 category_range = dict() for category, category_label_seg_max in category_list: category_range[category] = (offset, offset + category_label_seg_max) offset = offset + category_label_seg_max folder = os.path.join(output_folder, category) if not os.path.exists(folder): os.makedirs(folder) input_filelist = [] output_filelist = [] output_ply_filelist = [] for category in sorted(os.listdir(args.data_folder)): data_category_folder = os.path.join(args.data_folder, category) for filename in sorted(os.listdir(data_category_folder)): input_filelist.append( os.path.join(args.data_folder, category, filename)) output_filelist.append( os.path.join(output_folder, category, filename[0:-3] + 'seg')) output_ply_filelist.append( os.path.join(output_folder + '_ply', category, filename[0:-3] + 'ply')) # Prepare inputs print('{}-Preparing datasets...'.format(datetime.now())) data, label, data_num, label_test, _ = data_utils.load_seg( args.filelist_val) # no segmentation labels net = Net(input_channels=1, output_channels=N_CLASSES) net.load_state_dict( torch.load(os.path.join(args.savedir, "state_dict.pth"))) net.cuda() net.eval() ds = PartNormalDataset(data, data_num, label_test, net.config, npoints=args.npoints) test_loader = torch.utils.data.DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=THREADS, collate_fn=tree_collate) cm = np.zeros((N_CLASSES, N_CLASSES)) t = tqdm(test_loader, ncols=120) with torch.no_grad(): count = 0 for shape_id in tqdm(range(len(ds)), ncols=120): segmentation_ = None batches = [] if args.ntree <= args.batchsize: batch = [] for tree_id in range(args.ntree): batch.append(ds.__getitem__(shape_id)) batches.append(batch) else: for i in range(math.ceil(args.ntree / args.batchsize)): bs = min(args.batchsize, args.ntree - i * args.batchsize) batch = [] for tree_id in range(bs): batch.append(ds.__getitem__(shape_id)) batches.append(batch) for batch in batches: pts, features, seg, tree = tree_collate(batch) if USE_CUDA: features = features.cuda() pts = pts.cuda() for l_id in range(len(tree)): tree[l_id]["points"] = tree[l_id]["points"].cuda() tree[l_id]["indices"] = tree[l_id]["indices"].cuda() outputs = net(features, pts, tree) for i in range(pts.size(0)): pts_src = pts[i].cpu().numpy() # pts_dest point_num = data_num[count] pts_dest = data[count] pts_dest = pts_dest[:point_num] object_label = label[count] category = category_list[object_label][0] label_start, label_end = category_range[category] seg_ = outputs[i][:, label_start:label_end].cpu().numpy() seg_ = nearest_correspondance(pts_src, pts_dest, seg_) if segmentation_ is None: segmentation_ = seg_ else: segmentation_ += seg_ segmentation_ = np.argmax(segmentation_, axis=1) # save labels np.savetxt(output_filelist[count], segmentation_, fmt="%i") if args.ply: data_utils.save_ply_property(pts_dest, segmentation_, 6, output_ply_filelist[count]) count += 1
def test(args, flist_test): N_CLASSES = 13 # create the network print("Creating network...") if args.nocolor: net = Net(input_channels=1, output_channels=N_CLASSES) else: net = Net(input_channels=3, output_channels=N_CLASSES) net.load_state_dict(torch.load(os.path.join(args.savedir, "state_dict.pth"))) net.cuda() net.eval() print("parameters", count_parameters(net)) for filename in flist_test: print(filename) ds = PartDatasetTest(filename, args.rootdir, block_size=args.blocksize, min_pick_per_point= args.npick, npoints= args.npoints, test_step=args.test_step, nocolor=args.nocolor ) loader = torch.utils.data.DataLoader(ds, batch_size=args.batchsize, shuffle=False, num_workers=args.threads ) xyzrgb = ds.xyzrgb[:,:3] scores = np.zeros((xyzrgb.shape[0], N_CLASSES)) with torch.no_grad(): t = tqdm(loader, ncols=80) for pts, features, indices in t: features = features.cuda() pts = pts.cuda() outputs = net(features, pts) outputs_np = outputs.cpu().numpy().reshape((-1, N_CLASSES)) scores[indices.cpu().numpy().ravel()] += outputs_np mask = np.logical_not(scores.sum(1)==0) scores = scores[mask] pts_src = xyzrgb[mask] # create the scores for all points scores = nearest_correspondance(pts_src, xyzrgb, scores, K=1) # compute softmax scores = scores - scores.max(axis=1)[:,None] scores = np.exp(scores) / np.exp(scores).sum(1)[:,None] scores = np.nan_to_num(scores) os.makedirs(os.path.join(args.savedir, filename), exist_ok=True) # saving labels save_fname = os.path.join(args.savedir, filename, "pred.txt") scores = scores.argmax(1) np.savetxt(save_fname,scores,fmt='%d') if args.savepts: save_fname = os.path.join(args.savedir, filename, "pts.txt") xyzrgb = np.concatenate([xyzrgb, np.expand_dims(scores,1)], axis=1) np.savetxt(save_fname,xyzrgb,fmt=['%.4f','%.4f','%.4f','%d'])
def test(args): THREADS = 4 BATCH_SIZE = args.batchsize USE_CUDA = True N_CLASSES = 50 args.data_folder = os.path.join(args.rootdir, "test_data") # create the output folders output_folder = args.savedir + '_predictions' category_list = [ (category, int(label_num)) for (category, label_num) in [line.split() for line in open(args.category, 'r')] ] offset = 0 category_range = dict() for category, category_label_seg_max in category_list: category_range[category] = (offset, offset + category_label_seg_max) offset = offset + category_label_seg_max folder = os.path.join(output_folder, category) if not os.path.exists(folder): os.makedirs(folder) input_filelist = [] output_filelist = [] output_ply_filelist = [] for category in sorted(os.listdir(args.data_folder)): data_category_folder = os.path.join(args.data_folder, category) for filename in sorted(os.listdir(data_category_folder)): input_filelist.append( os.path.join(args.data_folder, category, filename)) output_filelist.append( os.path.join(output_folder, category, filename[0:-3] + 'seg')) output_ply_filelist.append( os.path.join(output_folder + '_ply', category, filename[0:-3] + 'ply')) # Prepare inputs print('{}-Preparing datasets...'.format(datetime.now())) data, label, data_num, label_test, _ = data_utils.load_seg( args.filelist_val) # no segmentation labels net = Net(input_channels=1, output_channels=N_CLASSES) net.load_state_dict( torch.load(os.path.join(args.savedir, "state_dict.pth"))) net.cuda() net.eval() ds = PartNormalDataset(data, data_num, label_test, net.config, npoints=args.npoints) test_loader = torch.utils.data.DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=THREADS, collate_fn=tree_collate) cm = np.zeros((N_CLASSES, N_CLASSES)) t = tqdm(test_loader, ncols=120) with torch.no_grad(): count = 0 for pts, features, seg, tree in t: if USE_CUDA: features = features.cuda() pts = pts.cuda() for l_id in range(len(tree)): tree[l_id]["points"] = tree[l_id]["points"].cuda() tree[l_id]["indices"] = tree[l_id]["indices"].cuda() outputs = net(features, pts, tree) # save results for i in range(pts.size(0)): # pts_src pts_src = pts[i].cpu().numpy() # pts_dest point_num = data_num[count + i] pts_dest = data[count + i] pts_dest = pts_dest[:point_num] object_label = label[count + i] category = category_list[object_label][0] label_start, label_end = category_range[category] seg_ = outputs[i][:, label_start:label_end].cpu().numpy() seg_ = np.argmax(seg_, axis=1) seg_ = nearest_correspondance(pts_src, pts_dest, seg_) # save labels np.savetxt(output_filelist[count + i], seg_, fmt="%i") if args.ply: data_utils.save_ply_property( pts_dest, seg_, 6, output_ply_filelist[count + i]) count += pts.size(0) output_np = np.argmax(outputs.cpu().detach().numpy(), axis=2).copy() target_np = seg.cpu().numpy().copy() cm_ = confusion_matrix(target_np.ravel(), output_np.ravel(), labels=list(range(N_CLASSES))) cm += cm_ oa = "{:.3f}".format(metrics.stats_overall_accuracy(cm)) aa = "{:.3f}".format(metrics.stats_accuracy_per_class(cm)[0]) t.set_postfix(OA=oa, AA=aa)
def train(args, flist_train, flist_test): N_CLASSES = 13 # create the network print("Creating network...") if args.nocolor: net = Net(input_channels=1, output_channels=N_CLASSES) else: net = Net(input_channels=3, output_channels=N_CLASSES) net.cuda() print("parameters", count_parameters(net)) print("Creating dataloader and optimizer...") ds = PartDatasetTrainVal(flist_train, args.rootdir, training=True, block_size=args.blocksize, npoints=args.npoints,iteration_number=args.batchsize*args.iter, nocolor=args.nocolor) train_loader = torch.utils.data.DataLoader(ds, batch_size=args.batchsize, shuffle=True, num_workers=args.threads ) ds_val = PartDatasetTrainVal(flist_test, args.rootdir, training=False, block_size=args.blocksize, npoints=args.npoints, nocolor=args.nocolor) test_loader = torch.utils.data.DataLoader(ds_val, batch_size=args.batchsize, shuffle=False, num_workers=args.threads ) optimizer = torch.optim.Adam(net.parameters(), lr=1e-3) print("done") # create the root folder print("Creating results folder") time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S') root_folder = os.path.join(args.savedir, "s3dis_area{}_{}_nocolor{}_{}".format(args.area, args.npoints, args.nocolor, time_string)) os.makedirs(root_folder, exist_ok=True) print("done at", root_folder) # create the log file logs = open(os.path.join(root_folder, "log.txt"), "w") # iterate over epochs for epoch in range(50): ####### # training net.train() train_loss = 0 cm = np.zeros((N_CLASSES, N_CLASSES)) t = tqdm(train_loader, ncols=100, desc="Epoch {}".format(epoch)) for pts, features, seg in t: features = features.cuda() pts = pts.cuda() seg = seg.cuda() optimizer.zero_grad() outputs = net(features, pts) loss = F.cross_entropy(outputs.view(-1, N_CLASSES), seg.view(-1)) loss.backward() optimizer.step() output_np = np.argmax(outputs.cpu().detach().numpy(), axis=2).copy() target_np = seg.cpu().numpy().copy() cm_ = confusion_matrix(target_np.ravel(), output_np.ravel(), labels=list(range(N_CLASSES))) cm += cm_ oa = f"{metrics.stats_overall_accuracy(cm):.5f}" aa = f"{metrics.stats_accuracy_per_class(cm)[0]:.5f}" iou = f"{metrics.stats_iou_per_class(cm)[0]:.5f}" train_loss += loss.detach().cpu().item() t.set_postfix(OA=wblue(oa), AA=wblue(aa), IOU=wblue(iou), LOSS=wblue(f"{train_loss/cm.sum():.4e}")) ###### ## validation net.eval() cm_test = np.zeros((N_CLASSES, N_CLASSES)) test_loss = 0 t = tqdm(test_loader, ncols=80, desc=" Test epoch {}".format(epoch)) with torch.no_grad(): for pts, features, seg in t: features = features.cuda() pts = pts.cuda() seg = seg.cuda() outputs = net(features, pts) loss = F.cross_entropy(outputs.view(-1, N_CLASSES), seg.view(-1)) output_np = np.argmax(outputs.cpu().detach().numpy(), axis=2).copy() target_np = seg.cpu().numpy().copy() cm_ = confusion_matrix(target_np.ravel(), output_np.ravel(), labels=list(range(N_CLASSES))) cm_test += cm_ oa_val = f"{metrics.stats_overall_accuracy(cm_test):.5f}" aa_val = f"{metrics.stats_accuracy_per_class(cm_test)[0]:.5f}" iou_val = f"{metrics.stats_iou_per_class(cm_test)[0]:.5f}" test_loss += loss.detach().cpu().item() t.set_postfix(OA=wgreen(oa_val), AA=wgreen(aa_val), IOU=wgreen(iou_val), LOSS=wgreen(f"{test_loss/cm_test.sum():.4e}")) # save the model torch.save(net.state_dict(), os.path.join(root_folder, "state_dict.pth")) # write the logs logs.write(f"{epoch} {oa} {aa} {iou} {oa_val} {aa_val} {iou_val}\n") logs.flush() logs.close()