mirror of
https://github.com/botastic/SoftGroup.git
synced 2025-10-16 11:45:42 +00:00
341 lines
16 KiB
Python
341 lines
16 KiB
Python
import torch
|
|
import time
|
|
import numpy as np
|
|
import random
|
|
import os
|
|
|
|
from util.config import cfg
|
|
cfg.task = 'test'
|
|
from util.log import logger
|
|
import util.utils as utils
|
|
import util.eval_s3dis as eval
|
|
|
|
def init():
|
|
global result_dir
|
|
result_dir = os.path.join(cfg.exp_path, 'result', cfg.split, cfg.test_area)
|
|
backup_dir = os.path.join(result_dir, 'backup_files')
|
|
os.makedirs(backup_dir, exist_ok=True)
|
|
os.makedirs(os.path.join(result_dir, 'predicted_masks'), exist_ok=True)
|
|
os.system('cp test.py {}'.format(backup_dir))
|
|
os.system('cp {} {}'.format(cfg.model_dir, backup_dir))
|
|
os.system('cp {} {}'.format(cfg.dataset_dir, backup_dir))
|
|
os.system('cp {} {}'.format(cfg.config, backup_dir))
|
|
|
|
global semantic_label_idx
|
|
# semantic_label_idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]
|
|
semantic_label_idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
|
|
|
|
logger.info(cfg)
|
|
|
|
random.seed(cfg.test_seed)
|
|
np.random.seed(cfg.test_seed)
|
|
torch.manual_seed(cfg.test_seed)
|
|
torch.cuda.manual_seed_all(cfg.test_seed)
|
|
|
|
|
|
def test(model, model_fn, data_name, epoch):
|
|
logger.info('>>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>')
|
|
|
|
if cfg.dataset == 'scannetv2':
|
|
if data_name == 'scannet':
|
|
from data.scannetv2_inst import Dataset
|
|
dataset = Dataset(test=True)
|
|
dataset.testLoader()
|
|
else:
|
|
print("Error: no data loader - " + data_name)
|
|
exit(0)
|
|
elif cfg.dataset == 's3dis':
|
|
if data_name == 's3dis':
|
|
import data.s3dis_inst
|
|
dataset = data.s3dis_inst.Dataset(test=True)
|
|
dataset.testLoader()
|
|
|
|
dataloader = dataset.test_data_loader
|
|
total = 0
|
|
|
|
with torch.no_grad():
|
|
model = model.eval()
|
|
|
|
total_end1 = 0.
|
|
matches = {}
|
|
nclusters = 0
|
|
for i_b, batch in enumerate(dataloader):
|
|
test_scene_name = dataset.test_file_names[int(batch['id'][0])].split('/')[-1][:-17]
|
|
# inference
|
|
start1 = time.time()
|
|
preds = model_fn(batch, model, epoch)
|
|
end1 = time.time() - start1
|
|
|
|
# decode results for evaluation
|
|
N = batch['feats'].shape[0]
|
|
semantic_scores = preds['semantic'] # (N, nClass=20) float32, cuda
|
|
semantic_pred = semantic_scores.max(1)[1] # (N) long, cuda
|
|
pt_offsets = preds['pt_offsets'] # (N, 3), float32, cuda
|
|
if (epoch > cfg.prepare_epochs) and not cfg.semantic_only:
|
|
scores = preds['score'] # (nProposal, 1) float, cuda
|
|
# scores_pred = torch.sigmoid(scores.view(-1))
|
|
|
|
scores_batch_idxs, proposals_idx, proposals_offset, mask_scores = preds['proposals']
|
|
cls_scores = preds['cls_score'].softmax(1)
|
|
slice_inds = torch.arange(cls_scores.size(0), dtype=torch.long, device=cls_scores.device)
|
|
cls_scores_new, cls_pred = cls_scores[:, :-1].max(1)
|
|
|
|
cluster_scores_list = []
|
|
clusters_list = []
|
|
cluster_semantic_id_list = []
|
|
for i in range(13):
|
|
if i < 2:
|
|
proposals_pred = semantic_pred == i
|
|
proposals_pred = proposals_pred[None, :].int()
|
|
semantic_id = cls_scores.new_tensor([semantic_label_idx[i]], dtype=torch.long)
|
|
scores_pred = cls_scores.new_tensor([1.], dtype=torch.float32)
|
|
else:
|
|
|
|
# arg_score = cls_pred == i
|
|
score_inds = (cls_scores[:, i] > 0.001)
|
|
cls_scores_new = cls_scores[:, i]
|
|
scores_pred = scores[slice_inds, i]
|
|
scores_pred = scores_pred.clamp(0, 1) * cls_scores_new
|
|
# scores_pred = cls_scores_new
|
|
# mask_cls_pred = cls_pred[scores_batch_idxs.long()]
|
|
mask_slice_inds = torch.arange(scores_batch_idxs.size(0), dtype=torch.long, device=scores_batch_idxs.device)
|
|
mask_scores_new = mask_scores[:, i]
|
|
# proposals_idx: (sumNPoint, 2), int, cpu, [:, 0] for cluster_id, [:, 1] for corresponding point idxs in N
|
|
# proposals_offset: (nProposal + 1), int, cpu
|
|
proposals_pred = torch.zeros((proposals_offset.shape[0] - 1, N), dtype=torch.int, device=scores_pred.device)
|
|
# (nProposal, N), int, cuda
|
|
|
|
# outlier filtering
|
|
test_mask_score_thre = getattr(cfg, 'test_mask_score_thre', -0.5)
|
|
_mask = mask_scores_new > test_mask_score_thre
|
|
proposals_pred[proposals_idx[_mask][:, 0].long(), proposals_idx[_mask][:, 1].long()] = 1
|
|
|
|
# bg filtering
|
|
# import pdb; pdb.set_trace()
|
|
# pos_inds = (cls_pred != cfg.classes - 2)
|
|
# proposals_pred = proposals_pred[pos_inds]
|
|
# scores_pred = scores_pred[pos_inds]
|
|
# cls_pred = cls_pred[pos_inds]
|
|
# import pdb; pdb.set_trace()
|
|
semantic_id = cls_scores.new_full(cls_scores_new.size(), semantic_label_idx[i], dtype=torch.long)
|
|
|
|
|
|
semantic_id1 = torch.tensor(semantic_label_idx, device=scores_pred.device) \
|
|
[semantic_pred[proposals_idx[:, 1][proposals_offset[:-1].long()].long()]] # (nProposal), long
|
|
# semantic_id_idx = semantic_pred[proposals_idx[:, 1][proposals_offset[:-1].long()].long()]
|
|
|
|
proposals_pred = proposals_pred[score_inds]
|
|
scores_pred = scores_pred[score_inds]
|
|
semantic_id = semantic_id[score_inds]
|
|
|
|
# score threshold
|
|
score_mask = (scores_pred > cfg.TEST_SCORE_THRESH)
|
|
scores_pred = scores_pred[score_mask]
|
|
proposals_pred = proposals_pred[score_mask]
|
|
semantic_id = semantic_id[score_mask]
|
|
# semantic_id_idx = semantic_id_idx[score_mask]
|
|
|
|
# npoint threshold
|
|
proposals_pointnum = proposals_pred.sum(1)
|
|
npoint_mask = (proposals_pointnum >= cfg.TEST_NPOINT_THRESH)
|
|
scores_pred = scores_pred[npoint_mask]
|
|
proposals_pred = proposals_pred[npoint_mask]
|
|
semantic_id = semantic_id[npoint_mask]
|
|
|
|
|
|
# nms (no need)
|
|
if getattr(cfg, 'using_NMS', False):
|
|
if semantic_id.shape[0] == 0:
|
|
pick_idxs = np.empty(0)
|
|
else:
|
|
proposals_pred_f = proposals_pred.float() # (nProposal, N), float, cuda
|
|
intersection = torch.mm(proposals_pred_f, proposals_pred_f.t()) # (nProposal, nProposal), float, cuda
|
|
proposals_pointnum = proposals_pred_f.sum(1) # (nProposal), float, cuda
|
|
proposals_pn_h = proposals_pointnum.unsqueeze(-1).repeat(1, proposals_pointnum.shape[0])
|
|
proposals_pn_v = proposals_pointnum.unsqueeze(0).repeat(proposals_pointnum.shape[0], 1)
|
|
cross_ious = intersection / (proposals_pn_h + proposals_pn_v - intersection)
|
|
pick_idxs = non_max_suppression(cross_ious.cpu().numpy(), scores_pred.cpu().numpy(), cfg.TEST_NMS_THRESH)
|
|
# int, (nCluster, N)
|
|
clusters = proposals_pred[pick_idxs]
|
|
cluster_scores = scores_pred[pick_idxs]
|
|
cluster_semantic_id = semantic_id[pick_idxs]
|
|
else:
|
|
clusters = proposals_pred
|
|
cluster_scores = scores_pred
|
|
cluster_semantic_id = semantic_id
|
|
clusters_list.append(clusters.cpu())
|
|
cluster_scores_list.append(cluster_scores)
|
|
cluster_semantic_id_list.append(cluster_semantic_id)
|
|
clusters = torch.cat(clusters_list)
|
|
cluster_scores = torch.cat(cluster_scores_list)
|
|
cluster_semantic_id = torch.cat(cluster_semantic_id_list)
|
|
# import pdb; pdb.set_trace()
|
|
|
|
nclusters = clusters.shape[0]
|
|
# if nclusters > cfg.max_clusters:
|
|
# nclusters = cfg.max_clusters
|
|
# _, topk_inds = cluster_scores.topk(cfg.max_clusters)
|
|
# clusters = clusters[topk_inds]
|
|
# cluster_scores = cluster_scores[topk_inds]
|
|
# cluster_semantic_id = cluster_semantic_id[topk_inds]
|
|
|
|
|
|
# prepare for evaluation
|
|
if cfg.eval:
|
|
pred_info = {}
|
|
pred_info['conf'] = cluster_scores.cpu().numpy()
|
|
pred_info['label_id'] = cluster_semantic_id.cpu().numpy()
|
|
pred_info['mask'] = clusters.cpu().numpy()
|
|
gt_file = os.path.join(cfg.data_root, cfg.dataset, cfg.split + '_gt', test_scene_name + '.txt')
|
|
gt2pred, pred2gt = eval.assign_instances_for_scan(test_scene_name, pred_info, gt_file)
|
|
|
|
matches[test_scene_name] = {}
|
|
matches[test_scene_name]['gt'] = gt2pred
|
|
matches[test_scene_name]['pred'] = pred2gt
|
|
|
|
if cfg.split == 'val':
|
|
if cfg.semantic_only:
|
|
matches[test_scene_name] = {}
|
|
|
|
matches[test_scene_name]['seg_gt'] = batch['labels']
|
|
matches[test_scene_name]['seg_pred'] = semantic_pred
|
|
# break
|
|
|
|
|
|
# save files
|
|
if cfg.save_semantic:
|
|
os.makedirs(os.path.join(result_dir, 'semantic'), exist_ok=True)
|
|
semantic_np = semantic_pred.cpu().numpy()
|
|
np.save(os.path.join(result_dir, 'semantic', test_scene_name + '.npy'), semantic_np)
|
|
|
|
if cfg.save_pt_offsets:
|
|
os.makedirs(os.path.join(result_dir, 'coords_offsets'), exist_ok=True)
|
|
pt_offsets_np = pt_offsets.cpu().numpy()
|
|
coords_np = batch['locs_float'].numpy()
|
|
coords_offsets = np.concatenate((coords_np, pt_offsets_np), 1) # (N, 6)
|
|
np.save(os.path.join(result_dir, 'coords_offsets', test_scene_name + '.npy'), coords_offsets)
|
|
|
|
if(cfg.save_instance):
|
|
f = open(os.path.join(result_dir, test_scene_name + '.txt'), 'w')
|
|
for proposal_id in range(nclusters):
|
|
clusters_i = clusters[proposal_id].cpu().numpy() # (N)
|
|
semantic_label = semantic_label = cluster_semantic_id[proposal_id]
|
|
score = cluster_scores[proposal_id]
|
|
f.write('predicted_masks/{}_{:03d}.npy {} {:.4f}'.format( \
|
|
test_scene_name, proposal_id, semantic_label, score))
|
|
if proposal_id < nclusters - 1:
|
|
f.write('\n')
|
|
with open(os.path.join(result_dir, 'predicted_masks', test_scene_name + '_%03d.npy' % (proposal_id)), 'wb') as mask_f:
|
|
np.save(mask_f, clusters_i)
|
|
# np.savetxt(os.path.join(result_dir, 'predicted_masks', test_scene_name + '_%03d.npy' % (proposal_id)), clusters_i, fmt='%d')
|
|
f.close()
|
|
|
|
|
|
logger.info("instance iter: {}/{} point_num: {} ncluster: {} inference time: {:.2f}s".format( \
|
|
batch['id'][0] + 1, len(dataset.test_files), N, nclusters, end1))
|
|
total_end1 += end1
|
|
# import pdb; pdb.set_trace()
|
|
# if i_b == 10: break
|
|
# break
|
|
|
|
# evaluation
|
|
if cfg.eval and not cfg.semantic_only:
|
|
ap_scores = eval.evaluate_matches(matches)
|
|
avgs = eval.compute_averages(ap_scores)
|
|
eval.print_results(avgs)
|
|
|
|
logger.info("whole set inference time: {:.2f}s, latency per frame: {:.2f}ms".format(total_end1, total_end1 / len(dataloader) * 1000))
|
|
|
|
# evaluate semantic segmantation accuracy and mIoU
|
|
if cfg.split == 'val':
|
|
seg_accuracy = evaluate_semantic_segmantation_accuracy(matches)
|
|
logger.info("semantic_segmantation_accuracy: {:.4f}".format(seg_accuracy))
|
|
miou = evaluate_semantic_segmantation_miou(matches)
|
|
logger.info("semantic_segmantation_mIoU: {:.4f}".format(miou))
|
|
|
|
def evaluate_semantic_segmantation_accuracy(matches):
|
|
seg_gt_list = []
|
|
seg_pred_list = []
|
|
for k, v in matches.items():
|
|
seg_gt_list.append(v['seg_gt'])
|
|
seg_pred_list.append(v['seg_pred'])
|
|
seg_gt_all = torch.cat(seg_gt_list, dim=0).cuda()
|
|
seg_pred_all = torch.cat(seg_pred_list, dim=0).cuda()
|
|
assert seg_gt_all.shape == seg_pred_all.shape
|
|
correct = (seg_gt_all[seg_gt_all != -100] == seg_pred_all[seg_gt_all != -100]).sum()
|
|
whole = (seg_gt_all != -100).sum()
|
|
seg_accuracy = correct.float() / whole.float()
|
|
return seg_accuracy
|
|
|
|
def evaluate_semantic_segmantation_miou(matches):
|
|
seg_gt_list = []
|
|
seg_pred_list = []
|
|
for k, v in matches.items():
|
|
seg_gt_list.append(v['seg_gt'])
|
|
seg_pred_list.append(v['seg_pred'])
|
|
seg_gt_all = torch.cat(seg_gt_list, dim=0).cuda()
|
|
seg_pred_all = torch.cat(seg_pred_list, dim=0).cuda()
|
|
assert seg_gt_all.shape == seg_pred_all.shape
|
|
pos = seg_gt_all >= 0
|
|
seg_gt_all = seg_gt_all[pos]
|
|
seg_pred_all = seg_pred_all[pos]
|
|
iou_list = []
|
|
for _index in seg_gt_all.unique():
|
|
if _index != -100:
|
|
intersection = ((seg_gt_all == _index) & (seg_pred_all == _index)).sum()
|
|
union = ((seg_gt_all == _index) | (seg_pred_all == _index)).sum()
|
|
iou = intersection.float() / union
|
|
iou_list.append(iou)
|
|
iou_tensor = torch.tensor(iou_list)
|
|
miou = iou_tensor.mean()
|
|
return miou
|
|
|
|
def non_max_suppression(ious, scores, threshold):
|
|
ixs = scores.argsort()[::-1]
|
|
pick = []
|
|
while len(ixs) > 0:
|
|
i = ixs[0]
|
|
pick.append(i)
|
|
iou = ious[i, ixs[1:]]
|
|
remove_ixs = np.where(iou > threshold)[0] + 1
|
|
ixs = np.delete(ixs, remove_ixs)
|
|
ixs = np.delete(ixs, 0)
|
|
return np.array(pick, dtype=np.int32)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
torch.backends.cudnn.enabled = False
|
|
init()
|
|
|
|
exp_name = cfg.config.split('/')[-1][:-5]
|
|
model_name = exp_name.split('_')[0]
|
|
data_name = exp_name.split('_')[-1]
|
|
|
|
logger.info('=> creating model ...')
|
|
logger.info('Classes: {}'.format(cfg.classes))
|
|
|
|
if model_name == 'softgroup':
|
|
from model.softgroup.softgroup import SoftGroup as Network
|
|
from model.softgroup.softgroup import model_fn_decorator
|
|
|
|
else:
|
|
print("Error: no model version " + model_name)
|
|
exit(0)
|
|
model = Network(cfg)
|
|
|
|
use_cuda = torch.cuda.is_available()
|
|
logger.info('cuda available: {}'.format(use_cuda))
|
|
assert use_cuda
|
|
model = model.cuda()
|
|
|
|
logger.info('#classifier parameters (model): {}'.format(sum([x.nelement() for x in model.parameters()])))
|
|
model_fn = model_fn_decorator(test=True)
|
|
|
|
# load model
|
|
utils.checkpoint_restore(cfg, model, None, cfg.exp_path, cfg.config.split('/')[-1][:-5],
|
|
use_cuda, cfg.test_epoch, dist=False, f=cfg.pretrain)
|
|
# resume from the latest epoch, or specify the epoch to restore
|
|
|
|
# evaluate
|
|
test(model, model_fn, data_name, cfg.test_epoch)
|