def test(args, flist_test, model_folder, info_class):
nb_class = info_class['nb_class']
# create the network
print("Creating network...")
net, features = get_model(nb_class, args)
state = torch.load(model_folder / "state_dict.pth")
net.load_state_dict(state['state_dict'])
net.cuda()
net.eval()
print(f"Number of parameters in the model: {count_parameters(net):,}")
for filename in flist_test:
print(filename)
ds_tst = PartDatasetTest(filename,
args.rootdir,
block_size=args.blocksize,
npoints=args.npoints,
test_step=args.test_step,
features=features,
labels=args.test_labels)
tst_loader = torch.utils.data.DataLoader(ds_tst,
batch_size=args.batchsize,
shuffle=False,
num_workers=args.num_workers)
xyz = ds_tst.xyzni[:, :3]
scores = np.zeros((xyz.shape[0], nb_class))
total_time = 0
iter_nb = 0
with torch.no_grad():
t = tqdm(tst_loader, ncols=150)
for pts, features, indices in t:
t1 = time.time()
features = features.cuda()
pts = pts.cuda()
outputs = net(features, pts)
t2 = time.time()
outputs_np = outputs.cpu().numpy().reshape((-1, nb_class))
scores[indices.cpu().numpy().ravel()] += outputs_np
iter_nb += 1
total_time += (t2 - t1)
t.set_postfix(
time=f"{total_time / (iter_nb * args.batchsize):05e}")
mask = np.logical_not(scores.sum(1) == 0)
scores = scores[mask]
pts_src = xyz[mask]
# create the scores for all points
scores = nearest_correspondance(pts_src, xyz, 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)
scores = scores.argmax(1)
# Compute confusion matrix
if args.test_labels:
tst_logs = InformationLogger(model_folder, 'tst')
lbl = ds_tst.labels[:, :]
cm = confusion_matrix(lbl.ravel(),
scores.ravel(),
labels=list(range(nb_class)))
cl_acc = metrics.stats_accuracy_per_class(cm)
cl_iou = metrics.stats_iou_per_class(cm)
cl_fscore = metrics.stats_f1score_per_class(cm)
print(f"Stats for test dataset:")
print_metric('Test', 'Accuracy', cl_acc)
print_metric('Test', 'iou', cl_iou)
print_metric('Test', 'F1-Score', cl_fscore)
tst_avg_score = {
'loss': -1,
'acc': cl_acc[0],
'iou': cl_iou[0],
'fscore': [0]
}
tst_class_score = {
'acc': cl_acc[1],
'iou': cl_iou[1],
'fscore': cl_fscore[1]
}
tst_logs.add_metric_values(tst_avg_score, -1)
tst_logs.add_class_scores(tst_class_score, -1)
# write error file.
# error2ply(model_folder / f"{filename}_error.ply", xyz=xyz, labels=lbl, prediction=scores, info_class=info_class['class_info'])
if args.savepts:
# Save predictions
out_folder = model_folder / 'tst'
out_folder.mkdir(exist_ok=True)
prediction2ply(model_folder / f"{filename}_predictions.ply",
xyz=xyz,
prediction=scores,
info_class=info_class['class_info'])
def train(args):
THREADS = 4
USE_CUDA = True
N_CLASSES = 50
EPOCHS = 200
MILESTONES = [60, 120]
shapenet_labels = [
['Airplane', 4],
['Bag', 2],
['Cap', 2],
['Car', 4],
['Chair', 4],
['Earphone', 3],
['Guitar', 3],
['Knife', 2],
['Lamp', 4],
['Laptop', 2],
['Motorbike', 6],
['Mug', 2],
['Pistol', 3],
['Rocket', 3],
['Skateboard', 3],
['Table', 3],
]
category_range = []
count = 0
for element in shapenet_labels:
part_start = count
count += element[1]
part_end = count
category_range.append([part_start, part_end])
# 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)
print("Computing class weights (if needed, 1 otherwise)...")
if args.weighted:
frequences = []
for i in range(len(shapenet_labels)):
frequences.append((labels == i).sum())
frequences = np.array(frequences)
frequences = frequences.mean() / frequences
else:
frequences = [1 for _ in range(len(shapenet_labels))]
weights = torch.FloatTensor(frequences)
if USE_CUDA:
weights = weights.cuda()
print("Done")
print("Creating network...")
net = get_model(args.model, input_channels=1, output_channels=N_CLASSES)
net.cuda()
print("parameters", count_parameters(net))
ds = PartNormalDataset(data_train,
data_num_train,
label_train,
npoints=args.npoints)
train_loader = torch.utils.data.DataLoader(ds,
batch_size=args.batchsize,
shuffle=True,
num_workers=THREADS)
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,
"{}_b{}_pts{}_weighted{}_{}".format(args.model, args.batchsize,
args.npoints, args.weighted,
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, indices in t:
if USE_CUDA:
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 = 0
for i in range(pts.size(0)):
# get the number of part for the shape
object_label = labels[indices[i]]
part_start, part_end = category_range[object_label]
part_nbr = part_end - part_start
loss = loss + weights[object_label] * F.cross_entropy(
outputs[i, :, part_start:part_end].view(-1, part_nbr),
seg[i].view(-1) - part_start)
loss.backward()
optimizer.step()
outputs_np = outputs.cpu().detach().numpy()
for i in range(pts.size(0)):
# get the number of part for the shape
object_label = labels[indices[i]]
part_start, part_end = category_range[object_label]
part_nbr = part_end - part_start
outputs_np[i, :, :part_start] = -1e7
outputs_np[i, :, part_end:] = -1e7
output_np = np.argmax(outputs_np, 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 train(args, dataset_dict, info_class):
nb_class = info_class['nb_class']
print("Creating network...")
net, features = get_model(nb_class, args)
net.cuda()
print(f"Number of parameters in the model: {count_parameters(net):,}")
print("Creating dataloader and optimizer...", end="")
ds_trn = PartDatasetTrainVal(filelist=dataset_dict['trn'],
folder=args.rootdir,
training=True,
block_size=args.blocksize,
npoints=args.npoints,
iteration_number=args.batchsize * args.iter,
features=features,
class_info=info_class['class_info'])
train_loader = torch.utils.data.DataLoader(ds_trn,
batch_size=args.batchsize,
shuffle=True,
num_workers=args.num_workers)
ds_val = PartDatasetTrainVal(filelist=dataset_dict['val'],
folder=args.rootdir,
training=False,
block_size=args.blocksize,
npoints=args.npoints,
iteration_number=args.batchsize *
args.val_iter,
features=features,
class_info=info_class['class_info'])
val_loader = torch.utils.data.DataLoader(ds_val,
batch_size=args.batchsize,
shuffle=False,
num_workers=args.num_workers)
optimizer = torch.optim.Adam(net.parameters(), lr=float(args.lr))
print("done")
# create the root folder
print("Creating results folder...", end="")
time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
root_folder = Path(
f"{args.savedir}/{args.model}_{args.npoints}_drop{args.drop}_{time_string}"
)
root_folder.mkdir(exist_ok=True)
args_dict = vars(args)
args_dict['data'] = dataset_dict
#write_config(root_folder, args_dict)
print("done at", root_folder)
# create the log file
trn_logs = InformationLogger(root_folder, 'trn')
val_logs = InformationLogger(root_folder, 'val')
# iterate over epochs
for epoch in range(args.nepochs):
#######
# training
net.train()
train_loss = 0
cm = np.zeros((nb_class, nb_class))
t = tqdm(train_loader, ncols=150, 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, nb_class), 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(nb_class)))
cm += cm_
oa = f"{metrics.stats_overall_accuracy(cm):.4f}"
acc = metrics.stats_accuracy_per_class(cm)
iou = metrics.stats_iou_per_class(cm)
train_loss += loss.detach().cpu().item()
t.set_postfix(OA=wblue(oa),
AA=wblue(f"{acc[0]:.4f}"),
IOU=wblue(f"{iou[0]:.4f}"),
LOSS=wblue(f"{train_loss / cm.sum():.4e}"))
fscore = metrics.stats_f1score_per_class(cm)
trn_metrics_values = {
'loss': f"{train_loss / cm.sum():.4e}",
'acc': acc[0],
'iou': iou[0],
'fscore': fscore[0]
}
trn_class_score = {'acc': acc[1], 'iou': iou[1], 'fscore': fscore[1]}
trn_logs.add_metric_values(trn_metrics_values, epoch)
trn_logs.add_class_scores(trn_class_score, epoch)
print_metric('Training', 'F1-Score', fscore)
######
# validation
net.eval()
cm_val = np.zeros((nb_class, nb_class))
val_loss = 0
t = tqdm(val_loader,
ncols=150,
desc=" Validation 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, nb_class),
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(nb_class)))
cm_val += cm_
oa_val = f"{metrics.stats_overall_accuracy(cm_val):.4f}"
acc_val = metrics.stats_accuracy_per_class(cm_val)
iou_val = metrics.stats_iou_per_class(cm_val)
val_loss += loss.detach().cpu().item()
t.set_postfix(OA=wgreen(oa_val),
AA=wgreen(f"{acc_val[0]:.4f}"),
IOU=wgreen(f"{iou_val[0]:.4f}"),
LOSS=wgreen(f"{val_loss / cm_val.sum():.4e}"))
fscore_val = metrics.stats_f1score_per_class(cm_val)
# save the model
state = {
'epoch': epoch,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
'args': args
}
torch.save(state, root_folder / "state_dict.pth")
# write the logs
val_metrics_values = {
'loss': f"{val_loss / cm_val.sum():.4e}",
'acc': acc_val[0],
'iou': iou_val[0],
'fscore': fscore_val[0]
}
val_class_score = {
'acc': acc_val[1],
'iou': iou_val[1],
'fscore': fscore_val[1]
}
val_logs.add_metric_values(val_metrics_values, epoch)
val_logs.add_class_scores(val_class_score, epoch)
print_metric('Validation', 'F1-Score', fscore_val)
return root_folder
def apply(epoch, training=False):
error = 0
cm = np.zeros((N_LABELS, N_LABELS))
if training:
t = tqdm(train_loader, desc="Epoch " + str(epoch), ncols=130)
for pts, features, targets, indices in t:
if args.cuda:
features = features.cuda()
pts = pts.cuda()
targets = targets.cuda()
optimizer.zero_grad()
outputs = net(features, pts)
targets = targets.view(-1)
loss = F.cross_entropy(outputs, targets)
loss.backward()
optimizer.step()
output_np = np.argmax(outputs.cpu().detach().numpy(), axis=1)
target_np = targets.cpu().numpy()
cm_ = confusion_matrix(target_np.ravel(),
output_np.ravel(),
labels=list(range(N_LABELS)))
cm += cm_
error += loss.item()
# scores
oa = "{:.5f}".format(metrics.stats_overall_accuracy(cm))
aa = "{:.5f}".format(metrics.stats_accuracy_per_class(cm)[0])
aiou = "{:.5f}".format(metrics.stats_iou_per_class(cm)[0])
aloss = "{:.5e}".format(error / cm.sum())
t.set_postfix(OA=oa, AA=aa, AIOU=aiou, ALoss=aloss)
else:
predictions = np.zeros((test_data.shape[0], N_LABELS), dtype=float)
t = tqdm(test_loader, desc=" Test " + str(epoch), ncols=100)
for pts, features, targets, indices in t:
if args.cuda:
features = features.cuda()
pts = pts.cuda()
targets = targets.cuda()
outputs = net(features, pts)
targets = targets.view(-1)
loss = F.cross_entropy(outputs, targets)
outputs_np = outputs.cpu().detach().numpy()
for i in range(indices.size(0)):
predictions[indices[i]] += outputs_np[i]
# l_ = np.argmax(outputs_np[i])
# predictions[indices[i],l_] += 1
error += loss.item()
if args.ntree == 1:
pred_labels = np.argmax(outputs_np, axis=1)
cm_ = confusion_matrix(targets.cpu().numpy(),
pred_labels,
labels=list(range(N_LABELS)))
cm += cm_
# scores
oa = "{:.5f}".format(metrics.stats_overall_accuracy(cm))
aa = "{:.5f}".format(
metrics.stats_accuracy_per_class(cm)[0])
aiou = "{:.5f}".format(metrics.stats_iou_per_class(cm)[0])
aloss = "{:.5e}".format(error / cm.sum())
t.set_postfix(OA=oa, AA=aa, AIOU=aiou, ALoss=aloss)
predictions = np.argmax(predictions, axis=1)
cm = confusion_matrix(test_labels,
predictions,
labels=list(range(N_LABELS)))
oa = "{:.5f}".format(metrics.stats_overall_accuracy(cm))
aa = "{:.5f}".format(metrics.stats_accuracy_per_class(cm)[0])
aiou = "{:.5f}".format(metrics.stats_iou_per_class(cm)[0])
aloss = "{:.5e}".format(error / cm.sum())
print("Predictions", "loss", aloss, "OA", oa, "AA", aa, "IOU",
aiou)
return aloss, oa, aa, aiou