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 main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--load_ckpt',
'-l',
default=
'log/seg/shellconv_seg_shapenet_2019-08-06-14-42-34/ckpts/epoch-326',
help='Path to a check point file for load')
parser.add_argument('--model',
'-m',
default='shellconv',
help='Model to use')
parser.add_argument('--setting',
'-x',
default='seg_shapenet',
help='Setting to use')
parser.add_argument('--repeat_num',
'-r',
help='Repeat number',
type=int,
default=1)
parser.add_argument('--save_ply',
'-s',
help='Save results as ply',
default=False)
args = parser.parse_args()
print(args)
model = importlib.import_module(args.model)
setting_path = os.path.join(os.path.dirname(__file__), args.model)
sys.path.append(setting_path)
setting = importlib.import_module(args.setting)
sample_num = setting.sample_num
output_folder = setting.data_folder + '_pred_shellnet_' + str(
args.repeat_num)
category_list = [(category, int(label_num)) for (
category,
label_num) in [line.split() for line in open(setting.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(setting.data_folder)):
data_category_folder = os.path.join(setting.data_folder, category)
for filename in sorted(os.listdir(data_category_folder)):
input_filelist.append(
os.path.join(setting.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, _, _ = data_utils.load_seg(setting.filelist_val)
batch_num = data.shape[0]
max_point_num = data.shape[1]
batch_size = args.repeat_num * math.ceil(data.shape[1] / sample_num)
print('{}-{:d} testing batches.'.format(datetime.now(), batch_num))
######################################################################
# Placeholders
indices = tf.placeholder(tf.int32,
shape=(batch_size, None, 2),
name="indices")
is_training = tf.placeholder(tf.bool, name='is_training')
pts_fts = tf.placeholder(tf.float32,
shape=(None, max_point_num, setting.data_dim),
name='pts_fts')
######################################################################
######################################################################
pts_fts_sampled = tf.gather_nd(pts_fts,
indices=indices,
name='pts_fts_sampled')
if setting.data_dim > 3:
points_sampled, _ = tf.split(pts_fts_sampled,
[3, setting.data_dim - 3],
axis=-1,
name='split_points_features')
else:
points_sampled = pts_fts_sampled
logits_op = model.get_model(points_sampled,
is_training,
setting.sconv_params,
setting.sdconv_params,
setting.fc_params,
sampling=setting.sampling,
weight_decay=setting.weight_decay,
bn_decay=None,
part_num=setting.num_class)
probs_op = tf.nn.softmax(logits_op, name='probs')
saver = tf.train.Saver()
parameter_num = np.sum(
[np.prod(v.shape.as_list()) for v in tf.trainable_variables()])
print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num))
with tf.Session() as sess:
# Load the model
saver.restore(sess, args.load_ckpt)
print('{}-Checkpoint loaded from {}!'.format(datetime.now(),
args.load_ckpt))
indices_batch_indices = np.tile(
np.reshape(np.arange(batch_size), (batch_size, 1, 1)),
(1, sample_num, 1))
for batch_idx in range(batch_num):
points_batch = data[[batch_idx] * batch_size, ...]
object_label = label[batch_idx]
point_num = data_num[batch_idx]
category = category_list[object_label][0]
label_start, label_end = category_range[category]
tile_num = math.ceil((sample_num * batch_size) / point_num)
indices_shuffle = np.tile(np.arange(point_num),
tile_num)[0:sample_num * batch_size]
np.random.shuffle(indices_shuffle)
indices_batch_shuffle = np.reshape(indices_shuffle,
(batch_size, sample_num, 1))
indices_batch = np.concatenate(
(indices_batch_indices, indices_batch_shuffle), axis=2)
probs = sess.run(
[probs_op],
feed_dict={
pts_fts: points_batch,
indices: indices_batch,
is_training: False,
})
probs_2d = np.reshape(probs, (sample_num * batch_size, -1))
predictions = [(-1, 0.0)] * point_num
for idx in range(sample_num * batch_size):
point_idx = indices_shuffle[idx]
probs = probs_2d[idx, label_start:label_end]
confidence = np.amax(probs)
seg_idx = np.argmax(probs)
if confidence > predictions[point_idx][1]:
predictions[point_idx] = (seg_idx, confidence)
labels = []
with open(output_filelist[batch_idx], 'w') as file_seg:
for seg_idx, _ in predictions:
file_seg.write('%d\n' % (seg_idx))
labels.append(seg_idx)
# read the coordinates from the txt file for verification
coordinates = [[float(value) for value in xyz.split(' ')]
for xyz in open(input_filelist[batch_idx], 'r')
if len(xyz.split(' ')) == 3]
assert (point_num == len(coordinates))
if args.save_ply:
data_utils.save_ply_property(np.array(coordinates),
np.array(labels), 6,
output_ply_filelist[batch_idx])
print('{}-[Testing]-Iter: {:06d} saved to {}'.format(
datetime.now(), batch_idx, output_filelist[batch_idx]))
sys.stdout.flush()
######################################################################
print('{}-Done!'.format(datetime.now()))
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 main():
parser = argparse.ArgumentParser()
parser.add_argument('--folder_gt',
'-g',
help='Path to ground truth folder',
required=True)
parser.add_argument('--folder_pred',
'-p',
help='Path to prediction folder',
required=True)
parser.add_argument('--folder_data',
'-d',
help='Path to point cloud data folder')
parser.add_argument('--part_avg',
'-a',
action='store_true',
help='Use part level average')
args = parser.parse_args()
print(args)
category_id_to_name = {
2691156: 'Airplane',
2773838: 'Bag',
2954340: 'Cap',
2958343: 'Car',
3001627: 'Chair',
3261776: 'Earphone',
3467517: 'Guitar',
3624134: 'Knife',
3636649: 'Lamp',
3642806: 'Laptop',
3790512: 'Motorbike',
3797390: 'Mug',
3948459: 'Pistol',
4099429: 'Rocket',
4225987: 'Skateboard',
4379243: 'Table'
}
categories = sorted(os.listdir(args.folder_gt))
label_min = sys.maxsize
for category in categories:
category_folder_gt = os.path.join(args.folder_gt, category)
filenames = sorted(os.listdir(category_folder_gt))
for filename in filenames:
filepath_gt = os.path.join(category_folder_gt, filename)
label_gt = np.loadtxt(filepath_gt).astype(np.int32)
label_min = min(label_min, np.amin(label_gt))
IoU = 0.0
total_num = 0
mIoU = 0.0
for category in categories:
category_folder_gt = os.path.join(args.folder_gt, category)
category_folder_pred = os.path.join(args.folder_pred, category)
if args.folder_data:
category_folder_data = os.path.join(args.folder_data, category)
category_folder_err = os.path.join(args.folder_pred + '_err_ply',
category)
IoU_category = 0.0
filenames = sorted(os.listdir(category_folder_gt))
for filename in filenames:
filepath_gt = os.path.join(category_folder_gt, filename)
filepath_pred = os.path.join(category_folder_pred, filename)
label_gt = np.loadtxt(filepath_gt).astype(np.int32) - label_min
label_pred = np.loadtxt(filepath_pred).astype(np.int32)
if args.folder_data:
filepath_data = os.path.join(category_folder_data,
filename[:-3] + 'pts')
filepath_err = os.path.join(category_folder_err,
filename[:-3] + 'ply')
coordinates = [[float(value) for value in xyz.split(' ')]
for xyz in open(filepath_data, 'r')
if len(xyz.split(' ')) == 3]
assert (label_gt.shape[0] == len(coordinates))
data_utils.save_ply_property(np.array(coordinates),
(label_gt == label_pred), 6,
filepath_err)
if args.part_avg:
label_max = np.amax(label_gt)
IoU_part = 0.0
for label_idx in range(label_max + 1):
locations_gt = (label_gt == label_idx)
locations_pred = (label_pred == label_idx)
I_locations = np.logical_and(locations_gt, locations_pred)
U_locations = np.logical_or(locations_gt, locations_pred)
I = np.sum(I_locations) + np.finfo(np.float32).eps
U = np.sum(U_locations) + np.finfo(np.float32).eps
IoU_part = IoU_part + I / U
IoU_sample = IoU_part / (label_max + 1)
else:
label_correct_locations = (label_gt == label_pred)
IoU_sample = np.sum(label_correct_locations) / label_gt.size
IoU_category = IoU_category + IoU_sample
IoU = IoU + IoU_category
IoU_category = IoU_category / len(filenames)
if category.isdigit():
print("IoU of %s: " % (category_id_to_name[int(category)]),
IoU_category)
else:
print("IoU of %s: " % category, IoU_category)
mIoU = mIoU + IoU_category
total_num = total_num + len(filenames)
IoU = IoU / total_num
print("IoU: ", IoU)
mIoU = mIoU / len(categories)
print("mIoU: ", mIoU)