204 lines
9.2 KiB
Python
Executable File
204 lines
9.2 KiB
Python
Executable File
"""
|
|
Author: Benny
|
|
Date: Nov 2019
|
|
"""
|
|
import argparse
|
|
import os
|
|
from data_utils.S3DISDataLoader import ScannetDatasetWholeScene
|
|
from data_utils.indoor3d_util import g_label2color
|
|
import torch
|
|
import logging
|
|
from pathlib import Path
|
|
import sys
|
|
import importlib
|
|
from tqdm import tqdm
|
|
import provider
|
|
import numpy as np
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
ROOT_DIR = BASE_DIR
|
|
sys.path.append(os.path.join(ROOT_DIR, 'models'))
|
|
|
|
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
|
|
|
|
|
|
def parse_args():
|
|
'''PARAMETERS'''
|
|
parser = argparse.ArgumentParser('Model')
|
|
parser.add_argument('--batch_size', type=int, default=32, help='batch size in testing [default: 32]')
|
|
parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
|
|
parser.add_argument('--num_point', type=int, default=4096, help='point number [default: 4096]')
|
|
parser.add_argument('--log_dir', type=str, required=True, help='experiment root')
|
|
parser.add_argument('--visual', action='store_true', default=False, help='visualize result [default: False]')
|
|
parser.add_argument('--test_area', type=int, default=5, help='area for testing, option: 1-6 [default: 5]')
|
|
parser.add_argument('--num_votes', type=int, default=3, help='aggregate segmentation scores with voting [default: 5]')
|
|
return parser.parse_args()
|
|
|
|
|
|
def add_vote(vote_label_pool, point_idx, pred_label, weight):
|
|
B = pred_label.shape[0]
|
|
N = pred_label.shape[1]
|
|
for b in range(B):
|
|
for n in range(N):
|
|
if weight[b, n] != 0 and not np.isinf(weight[b, n]):
|
|
vote_label_pool[int(point_idx[b, n]), int(pred_label[b, n])] += 1
|
|
return vote_label_pool
|
|
|
|
|
|
def main(args):
|
|
def log_string(str):
|
|
logger.info(str)
|
|
print(str)
|
|
|
|
'''HYPER PARAMETER'''
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
|
experiment_dir = 'log/sem_seg/' + args.log_dir
|
|
visual_dir = experiment_dir + '/visual/'
|
|
visual_dir = Path(visual_dir)
|
|
visual_dir.mkdir(exist_ok=True)
|
|
|
|
'''LOG'''
|
|
args = parse_args()
|
|
logger = logging.getLogger("Model")
|
|
logger.setLevel(logging.INFO)
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
|
|
file_handler.setLevel(logging.INFO)
|
|
file_handler.setFormatter(formatter)
|
|
logger.addHandler(file_handler)
|
|
log_string('PARAMETER ...')
|
|
log_string(args)
|
|
|
|
NUM_CLASSES = 13
|
|
BATCH_SIZE = args.batch_size
|
|
NUM_POINT = args.num_point
|
|
|
|
root = 'data/s3dis/stanford_indoor3d/'
|
|
|
|
TEST_DATASET_WHOLE_SCENE = ScannetDatasetWholeScene(root, split='test', test_area=args.test_area, block_points=NUM_POINT)
|
|
log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE))
|
|
|
|
'''MODEL LOADING'''
|
|
model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
|
|
MODEL = importlib.import_module(model_name)
|
|
classifier = MODEL.get_model(NUM_CLASSES).cuda()
|
|
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
|
|
classifier.load_state_dict(checkpoint['model_state_dict'])
|
|
classifier = classifier.eval()
|
|
|
|
with torch.no_grad():
|
|
scene_id = TEST_DATASET_WHOLE_SCENE.file_list
|
|
scene_id = [x[:-4] for x in scene_id]
|
|
num_batches = len(TEST_DATASET_WHOLE_SCENE)
|
|
|
|
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('---- EVALUATION WHOLE SCENE----')
|
|
|
|
for batch_idx in range(num_batches):
|
|
print("Inference [%d/%d] %s ..." % (batch_idx + 1, num_batches, scene_id[batch_idx]))
|
|
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)]
|
|
if args.visual:
|
|
fout = open(os.path.join(visual_dir, scene_id[batch_idx] + '_pred.obj'), 'w')
|
|
fout_gt = open(os.path.join(visual_dir, scene_id[batch_idx] + '_gt.obj'), 'w')
|
|
|
|
whole_scene_data = TEST_DATASET_WHOLE_SCENE.scene_points_list[batch_idx]
|
|
whole_scene_label = TEST_DATASET_WHOLE_SCENE.semantic_labels_list[batch_idx]
|
|
vote_label_pool = np.zeros((whole_scene_label.shape[0], NUM_CLASSES))
|
|
for _ in tqdm(range(args.num_votes), total=args.num_votes):
|
|
scene_data, scene_label, scene_smpw, scene_point_index = TEST_DATASET_WHOLE_SCENE[batch_idx]
|
|
num_blocks = scene_data.shape[0]
|
|
s_batch_num = (num_blocks + BATCH_SIZE - 1) // BATCH_SIZE
|
|
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 9))
|
|
|
|
batch_label = np.zeros((BATCH_SIZE, NUM_POINT))
|
|
batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT))
|
|
batch_smpw = np.zeros((BATCH_SIZE, NUM_POINT))
|
|
|
|
for sbatch in range(s_batch_num):
|
|
start_idx = sbatch * BATCH_SIZE
|
|
end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks)
|
|
real_batch_size = end_idx - start_idx
|
|
batch_data[0:real_batch_size, ...] = scene_data[start_idx:end_idx, ...]
|
|
batch_label[0:real_batch_size, ...] = scene_label[start_idx:end_idx, ...]
|
|
batch_point_index[0:real_batch_size, ...] = scene_point_index[start_idx:end_idx, ...]
|
|
batch_smpw[0:real_batch_size, ...] = scene_smpw[start_idx:end_idx, ...]
|
|
batch_data[:, :, 3:6] /= 1.0
|
|
|
|
torch_data = torch.Tensor(batch_data)
|
|
torch_data = torch_data.float().cuda()
|
|
torch_data = torch_data.transpose(2, 1)
|
|
seg_pred, _ = classifier(torch_data)
|
|
batch_pred_label = seg_pred.contiguous().cpu().data.max(2)[1].numpy()
|
|
|
|
vote_label_pool = add_vote(vote_label_pool, batch_point_index[0:real_batch_size, ...],
|
|
batch_pred_label[0:real_batch_size, ...],
|
|
batch_smpw[0:real_batch_size, ...])
|
|
|
|
pred_label = np.argmax(vote_label_pool, 1)
|
|
|
|
for l in range(NUM_CLASSES):
|
|
total_seen_class_tmp[l] += np.sum((whole_scene_label == l))
|
|
total_correct_class_tmp[l] += np.sum((pred_label == l) & (whole_scene_label == l))
|
|
total_iou_deno_class_tmp[l] += np.sum(((pred_label == l) | (whole_scene_label == 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 %s: %.4f' % (scene_id[batch_idx], tmp_iou))
|
|
print('----------------------------')
|
|
|
|
filename = os.path.join(visual_dir, scene_id[batch_idx] + '.txt')
|
|
with open(filename, 'w') as pl_save:
|
|
for i in pred_label:
|
|
pl_save.write(str(int(i)) + '\n')
|
|
pl_save.close()
|
|
for i in range(whole_scene_label.shape[0]):
|
|
color = g_label2color[pred_label[i]]
|
|
color_gt = g_label2color[whole_scene_label[i]]
|
|
if args.visual:
|
|
fout.write('v %f %f %f %d %d %d\n' % (
|
|
whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color[0], color[1],
|
|
color[2]))
|
|
fout_gt.write(
|
|
'v %f %f %f %d %d %d\n' % (
|
|
whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color_gt[0],
|
|
color_gt[1], color_gt[2]))
|
|
if args.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] + ' ' * (14 - len(seg_label_to_cat[l])),
|
|
total_correct_class[l] / float(total_iou_deno_class[l]))
|
|
log_string(iou_per_class_str)
|
|
log_string('eval point avg class IoU: %f' % np.mean(IoU))
|
|
log_string('eval 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('eval whole scene point accuracy: %f' % (
|
|
np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6)))
|
|
|
|
print("Done!")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_args()
|
|
main(args)
|