test_scannetv2.py

import argparse
import collections
from math import sqrt
import numpy as np
import os
import sys
import importlib
from sklearn.cluster import DBSCAN
import collections
import time

import torch
from torch.utils import data
import spconv
import scipy.stats as stats
from plyfile import PlyData, PlyElement
import pointgroup_ops
from torch_scatter import scatter_min, scatter_mean, scatter_max, scatter
# import sstnet

import evaluation
import utils

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(BASE_DIR, 'modules'))              # !!!! for importlib
sys.path.append(os.path.join(BASE_DIR, 'modules/model'))        # !!!! for importlib
sys.path.append(os.path.join(BASE_DIR, 'modules/datasets'))     # !!!! for importlib

def get_parser():
    parser = argparse.ArgumentParser(description="Point Cloud Instance Segmentation")
    parser.add_argument("--config", type=str, default="config/default.yaml", help="path to config file")
    ### pretrain
    parser.add_argument("--pretrain", type=str, default="", help="path to pretrain model")
    ### split
    parser.add_argument("--split", type=str, default="val", help="dataset split to test")
    ### semantic only
    parser.add_argument("--semantic", action="store_true", help="only evaluate semantic segmentation")
    ### log file path
    parser.add_argument("--log-file", type=str, default=None, help="log_file path")
    ### test srcipt operation
    parser.add_argument("--eval", action="store_true", help="evaluate or not")
    parser.add_argument("--save", action="store_true", help="save results or not")
    parser.add_argument("--visual", type=str, default=None, help="visual path, give to save visualization results")

    args_cfg = parser.parse_args()

    return args_cfg


def init():
    args = get_parser()
    cfg = utils.Config.fromfile(args.config)
    cfg.pretrain = args.pretrain
    cfg.semantic = args.semantic
    cfg.dataset.task = args.split  # change tasks ！！！！！！！！！！！！！
    cfg.data.visual = args.visual
    cfg.data.eval = args.eval
    cfg.data.save = args.save

    utils.set_random_seed(cfg.data.test_seed)

    #### get logger file
    params_dict = dict(
        epoch=cfg.data.test_epoch,
        optim=cfg.optimizer.type,
        lr=cfg.optimizer.lr,
        scheduler=cfg.lr_scheduler.type
    )
    if "test" in args.split:
        params_dict["suffix"] = "test"

    log_dir, logger = utils.collect_logger(
        prefix=os.path.splitext(args.config.split("/")[-1])[0],     # the name of the yaml file
        log_name="test_{}".format(time.time()),
        log_file=args.log_file,
        # **params_dict
    )

    logger.info("************************ Start Logging ************************")

    # log the config
    logger.info(cfg)

    global result_dir
    result_dir = os.path.join(
        log_dir,
        "result",
        "epoch_{}".format(cfg.data.test_epoch),
        args.split)
    # os.makedirs(os.path.join(result_dir, "predicted_masks"), exist_ok=True)

    global semantic_label_idx
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    semantic_label_idx = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]).cuda()

    return logger, cfg


def test(model, cfg, logger):
    
    logger.info(">>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>")

    epoch = cfg.data.test_epoch
    semantic = cfg.semantic

    cfg.dataset.test_mode = True
    cfg.dataloader.batch_size = 1 # !!! test with a single scene
    cfg.dataloader.num_workers = 2
    
    ####################################################
    
    
    datasets = importlib.import_module(cfg.dataset.type)
    test_dataset = datasets.ScanNetV2Inst_spg(cfg.dataset, test_mode=True)

    test_dataset.aug_flag = False
    test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset,
                                                  batch_size=cfg.dataloader.batch_size,
                                                  shuffle=False,
                                                  num_workers=cfg.dataloader.num_workers,
                                                  collate_fn=test_dataset.collate_fn,
                                                  pin_memory=True,
                                                  drop_last=False)
    
    with torch.no_grad():
        model = model.eval() ##########

        # init timer to calculate time
        timer = utils.Timer()

        # define evaluator
        # get the real data root
        data_root = os.path.join(os.path.dirname(__file__), cfg.dataset.data_root)
        
        label_root = os.path.join(data_root, cfg.dataset.task + "_gt")

        
        point_sem_evaluator = evaluation.ScanNetSemanticEvaluator(label_root, logger=logger)
        mid_sem_evaluator = evaluation.ScanNetSemanticEvaluator(label_root, logger=logger)
        sem_evaluator = evaluation.ScanNetSemanticEvaluator(label_root, logger=logger)
        inst_evaluator = evaluation.ScanNetInstanceEvaluator(label_root, logger=logger)
        
        for i, batch in enumerate(test_dataloader):
            
            torch.cuda.empty_cache() # (empty cuda cache, Optional)

            ##### prepare input and forward
            coords = batch["locs"].cuda()                           # [N, 1 + 3], long, cuda, dimension 0 for batch_idx
            locs_offset = batch["locs_offset"].cuda()               # [B, 3], long, cuda
            voxel_coords = batch["voxel_locs"].cuda()               # [M, 1 + 3], long, cuda
            p2v_map = batch["p2v_map"].cuda()                       # [N], int, cuda
            v2p_map = batch["v2p_map"].cuda()                       # [M, 1 + maxActive], int, cuda

            coords_float = batch["locs_float"].cuda()               # [N, 3], float32, cuda
            feats = batch["feats"].cuda()                           # [N, C], float32, cuda

            batch_offsets = batch["offsets"].cuda()                 # [B + 1], int, cuda
            superpoint = batch["superpoint"].cuda()                 # [N], long, cuda
            GIs = batch["GIs"]                       # igraph
            is1ins_labels = batch["is1ins_labels"].float().cuda()           # [NE], float32, cuda

            edge_u_list = batch["edge_u_list"].cuda()
            edge_v_list = batch["edge_v_list"].cuda()

            scene_list = batch["scene_list"]
            spatial_shape = batch["spatial_shape"]

            superpoint_cenetr_xyz = scatter(coords_float, superpoint, dim=0, reduce='mean')
            extra_data = {
                "superpoint": superpoint,
                "GIs": GIs,
                "edge_u_list": edge_u_list,
                "edge_v_list": edge_v_list,
                "superpoint_cenetr_xyz": superpoint_cenetr_xyz,
            }

            if cfg.model.use_coords:
                feats = torch.cat((feats, coords_float), 1)

            voxel_feats = pointgroup_ops.voxelization(feats, v2p_map, cfg.data.mode)    # [M, C]

            input_ = spconv.SparseConvTensor(voxel_feats,
                                             voxel_coords.int(),
                                             spatial_shape,
                                             cfg.dataloader.batch_size)

            ret = model(input_,             # SparseConvTensor
                        p2v_map,            # [N], int, cuda
                        extra_data)         # dict

            semantic_scores = ret["semantic_scores"]    # [N, nClass] float32, cuda

            sp_semantic_scores = ret['sp_semantic_scores']

            pred_sp_offset_vectors = ret['pred_sp_offset_vectors'] 

            pred_sp_occupancy = ret['pred_sp_occupancy']

            pred_sp_instance_size = ret['pred_sp_ins_size'] 

            test_scene_name = batch["scene_list"][0]
            

            ################################# test #########################################

            test_scene_name = batch["scene_list"][0]

            superpoint = superpoint.cpu().detach().numpy()

            ##### point-level semantic segmentation evaluation
            semantic_pred = semantic_scores.max(1)[1]  # [N]
            inputs = [{"scene_name": test_scene_name}]
            outputs = [{"semantic_pred": semantic_pred}]
            point_sem_evaluator.process(inputs, outputs) # semantic evaluation

            ##### middle-level semantic segmentation evaluation
            middle_level_semantic_pred = np.zeros(len(superpoint))
            point_semantic_pred = semantic_pred.cpu().detach().numpy()
            for spID in np.unique(superpoint):
                spMask = np.where(superpoint == spID)[0]
                sp_sem_label = stats.mode(point_semantic_pred[spMask])[0][0]
                middle_level_semantic_pred[spMask] = sp_sem_label
            
            middle_level_semantic_pred = middle_level_semantic_pred.astype('int')
            middle_level_semantic_pred = torch.from_numpy(middle_level_semantic_pred)

            inputs = [{"scene_name": test_scene_name}]
            outputs = [{"semantic_pred": middle_level_semantic_pred}]
            mid_sem_evaluator.process(inputs, outputs)

            ##### superpoint-level semantic segmentation evaluation
            sp_semantic_pred = sp_semantic_scores.max(1)[1]
            sp_semantic_pred = sp_semantic_pred.cpu().detach().numpy()
            
            assert len(sp_semantic_pred) == (superpoint.max() + 1)
            assert len(coords) == len(superpoint)
            point_level_semantic_pred = np.zeros(len(superpoint))

            for spID in np.unique(superpoint):
                spMask = np.where(superpoint == spID)[0]
                point_level_semantic_pred[spMask] = sp_semantic_pred[spID]
            
            point_level_semantic_pred = point_level_semantic_pred.astype('int')
            point_level_semantic_pred = torch.from_numpy(point_level_semantic_pred)

            inputs = [{"scene_name": test_scene_name}]
            outputs = [{"semantic_pred": point_level_semantic_pred}]
            sem_evaluator.process(inputs, outputs) # semantic evaluation
            
            ##### instance segmentation evaluation
            xyz_origin = coords_float.cpu().detach().numpy()
            graph = test_dataset.get_scene_graph(test_scene_name)
            pred_sp_offset_vectors = pred_sp_offset_vectors.cpu().detach().numpy()
            pred_sp_occupancy = pred_sp_occupancy.cpu().detach().numpy()
            pred_sp_instance_size = pred_sp_instance_size.cpu().detach().numpy()

            pred_info = {}
            pred_info["scene_name"] = test_scene_name
            pred_info["conf"], pred_info["label_id"], pred_info["mask"] = clustering_in_graph(test_scene_name, xyz_origin, superpoint, graph, 
                                                                            sp_semantic_pred, pred_sp_offset_vectors, pred_sp_occupancy, pred_sp_instance_size)

            inputs = [{"scene_name": test_scene_name}]
            inst_evaluator.process(inputs, [pred_info]) # instance evaluation

            

        inst_evaluator.evaluate(prec_rec=False)
        logger.info("point semantic evaluation")
        point_sem_evaluator.evaluate()
        logger.info("middle-level semantic evalution")
        mid_sem_evaluator.evaluate()
        logger.info("superpoint semantic evaluation")
        sem_evaluator.evaluate()





def clustering_in_graph(scene_name, xyz_origin, superpoint, graph, sp_semnatic_pred, pred_sp_offset_vectors, pred_sp_occupancy, pred_sp_ins_size):

    print(scene_name)
    assert len(xyz_origin) == len(superpoint)
    assert len(np.unique(superpoint)) == (superpoint.max() + 1) == len(sp_semnatic_pred) == len(pred_sp_offset_vectors)

    semantic_ind2label = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
    instanceSeg_valid_semantic_label = np.array([3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])

    # show_superpoint_instance_center(scene_name, xyz_origin, superpoint, sp_semnatic_pred, pred_sp_offset_vectors)
    ######################################
    superpoint_feat = dict()
    def get_superpoint_feature(spID):
        
        if spID in superpoint_feat:
            return superpoint_feat[spID]

        superpoint_mask = (superpoint == spID)

        sp_semantic_label = sp_semnatic_pred[spID]

        superpoint_center = xyz_origin[superpoint_mask].mean(0)
        sp_instance_center = superpoint_center + pred_sp_offset_vectors[spID]

        superpoint_feat[spID] = {'superpoint_pred_label':sp_semantic_label, 'sp_instance_center':sp_instance_center, 'superpoint_mask':superpoint_mask}
        return superpoint_feat[spID]

    spID_list = np.unique(superpoint)
    superpoint_visited = { spID:False for spID in spID_list }

    def BFS(spID):
        nonlocal superpoint_visited
        superpoint_visited[spID] = True
        queue = collections.deque()
        queue.append(spID)
        group_superpoint = set()
        group_superpoint.add(spID)

        group_mask = np.zeros(len(xyz_origin)).astype(bool)

        initial_spID_feat = get_superpoint_feature(spID)
        semantic_label = initial_spID_feat['superpoint_pred_label']
        group_mask = group_mask | initial_spID_feat['superpoint_mask']

        while queue:
            cur_spID = queue.popleft()

            cur_spID_feat = get_superpoint_feature(cur_spID)
            
            # for neighbor_spID in graph[cur_spID]:
            for neighbor_spID in graph.neighbors(vertex=cur_spID, mode='all'): # igraph https://igraph.org/python/api/latest/igraph._igraph.GraphBase.html#neighbors
                neighbor_spID_feat = get_superpoint_feature(neighbor_spID)

                if (neighbor_spID_feat['superpoint_pred_label'] == semantic_label) and (superpoint_visited[neighbor_spID] == False):

                    if np.linalg.norm((cur_spID_feat['sp_instance_center']-neighbor_spID_feat['sp_instance_center']), ord=2) < 0.25 * pred_sp_ins_size[spID]: ### key threshold
                        group_superpoint.add(neighbor_spID)
                        group_mask = group_mask | neighbor_spID_feat['superpoint_mask'] ####
                        superpoint_visited[neighbor_spID] = True
                        queue.append(neighbor_spID)

        return list(group_superpoint), group_mask
    
    #################################################################
    
    primary_instance_list = []
    fragment_list = []

    def get_group_pred_occupancy(group_sp_list):
        group_sp_list = np.array(group_sp_list)
        group_pred_occupancy = pred_sp_occupancy[group_sp_list]
        group_pred_occupancy = np.exp(group_pred_occupancy).mean() 
        return group_pred_occupancy

    
    def get_group_instance_center(group_list):
        instance_center = np.zeros(3)
        group_point_n = 0
        for spID in group_list:
            spID_feat = get_superpoint_feature(spID)
            instance_center += spID_feat['sp_instance_center'] * spID_feat['superpoint_mask'].sum()
            group_point_n += spID_feat['superpoint_mask'].sum()
        
        return instance_center / group_point_n

    def get_group_instance_size(group_sp_list):
        group_sp_list = np.array(group_sp_list)
        return np.mean(pred_sp_ins_size[group_sp_list])

    bfs_mask_list = []

    for spID in spID_list:
        spID_feat = get_superpoint_feature(spID)
        superpoint_pred_label = spID_feat['superpoint_pred_label']
        if (semantic_ind2label[superpoint_pred_label] not in instanceSeg_valid_semantic_label) or (superpoint_visited[spID] == True):
            continue

        group_sp_list, group_mask = BFS(spID)
        ############
        bfs_mask_list.append(group_mask.astype(int))

        group_pred_occupancy = get_group_pred_occupancy(group_sp_list)
        high_thre = 0.3 * group_pred_occupancy

        group_xyz = xyz_origin[group_mask]
        xyz = group_xyz * 50 # 0.02 cm
        xyz = torch.from_numpy(xyz).long()
        xyz = torch.cat([torch.LongTensor(xyz.shape[0], 1).fill_(0), xyz], 1)
        voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(xyz, 1, 4)
        group_voxel_num = voxel_locs.shape[0] 

        group_n = group_mask.sum()

        if group_voxel_num < high_thre:
            fragment_center = get_group_instance_center(group_sp_list)
            fragment = {'mask':group_mask, 'classLabel':superpoint_pred_label, 'instance_center':fragment_center, 
                        'absorbed':False, 'group_sp_list':group_sp_list, 'group_n': group_n}
            fragment_list.append(fragment)
        else:
            r_voxel = 0.02 * sqrt(group_pred_occupancy) # 2cm = 0.02m
            r_size = 0.01 * sqrt(group_n)
            r_ins_size = get_group_instance_size(group_sp_list)
            r_set = max(r_size, r_voxel, r_ins_size)



            primary_instance_center = get_group_instance_center(group_sp_list)
            primary_instance = {'mask':group_mask, 'classLabel':superpoint_pred_label, 'instance_center':primary_instance_center, 
                                'r_set':r_set, 'group_sp_list':group_sp_list, 'group_n': group_n}
            primary_instance_list.append(primary_instance)



    for fi, fragment in enumerate(fragment_list):
        index, dis_min = -1, float('inf')

        for i, primary_instance in enumerate(primary_instance_list):
            dis = fragment['instance_center'] - primary_instance['instance_center']
            assert dis.shape==(3,)
            dis = np.linalg.norm(dis, ord=2)
            if fragment['classLabel']==primary_instance['classLabel'] and dis < dis_min:
                index = i
                dis_min = dis
        
        closest_primary = primary_instance_list[index]
        if dis_min < closest_primary['r_set'] :
            _ins_mask = fragment['mask'] | closest_primary['mask']
            _center = get_group_instance_center(fragment['group_sp_list'] + closest_primary['group_sp_list'])

            _r_voxel = 0.02 * sqrt(get_group_pred_occupancy(fragment['group_sp_list'] + closest_primary['group_sp_list']))
            _r_size = 0.01 * sqrt(_ins_mask.sum())
            _r_ins_size = get_group_instance_size(fragment['group_sp_list'] + closest_primary['group_sp_list'])
            _r_set = max(_r_voxel, _r_size, closest_primary['r_set'], _r_ins_size)

            closest_primary['mask'] = _ins_mask
            closest_primary['instance_center'] = _center
            closest_primary['r_set'] = _r_set
            closest_primary['group_n'] = _ins_mask.sum()
            closest_primary['group_sp_list'] += fragment['group_sp_list']
            fragment['absorbed'] = True
    
    ############################ return result ##########################
    conf, label_id, ins_mask_list = [], [], []
    for primary_instance in primary_instance_list:

        _conf = primary_instance['group_n'] / get_group_pred_occupancy(primary_instance['group_sp_list'])
        _conf = min(_conf, 1)
        conf.append(_conf)
        label_id.append(semantic_ind2label[primary_instance['classLabel']])
        ins_mask_list.append(primary_instance['mask'].astype(int))  

    
    # show_instance_segmentation_result(scene_name, xyz_origin, ins_mask_list, name='final')
    
    return np.array(conf), np.array(label_id), np.array(ins_mask_list)




def show_instance_segmentation_result(scene_name, xyz, instance_mask_list:list, name=''):
    instance_num = len(instance_mask_list)
    instance_color_table = np.random.randint(low=0, high=255, size=(instance_num, 3))
    color = np.zeros((len(xyz), 3))

    for i, instance_mask in enumerate(instance_mask_list):
        instance_mask = instance_mask.astype(bool)
        color[instance_mask] = instance_color_table[i]

    xyz_rgb = np.concatenate((xyz, color), axis=1)
    vertex = np.array([tuple(i) for i in xyz_rgb], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),('blue', 'u1')])
    d = PlyElement.describe(vertex, 'vertex')
    plydata = PlyData([d])
    save_path = r'' # !!!
    plydata.write(os.path.join(save_path, scene_name+f'_{name}.ply'))


def show_superpoint_instance_center(scene_name, xyz, superpoint, sp_semnatic_pred, pred_sp_offset_vectors):
    superpoint_num = len(np.unique(superpoint))
    color_table = np.random.randint(low=0, high=255, size=(20, 3))
    color = np.zeros((superpoint_num, 3))
    superpoint_instance_center = np.zeros((superpoint_num, 3))

    for spID in np.unique(superpoint):
        superpoint_mask = (superpoint == spID)

        sp_semantic_label = sp_semnatic_pred[spID]

        superpoint_center = xyz[superpoint_mask].mean(0)
        sp_instance_center = superpoint_center + pred_sp_offset_vectors[spID]

        superpoint_instance_center[spID] = sp_instance_center
        color[spID] = color_table[sp_semantic_label]
    
    xyz_rgb = np.concatenate((superpoint_instance_center, color), axis=1)
    vertex = np.array([tuple(i) for i in xyz_rgb], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),('blue', 'u1')])
    d = PlyElement.describe(vertex, 'vertex')
    plydata = PlyData([d])
    save_path = r'' # !!!
    plydata.write(os.path.join(save_path, scene_name+'_ins_center.ply'))




if __name__ == "__main__":
    logger, cfg = init()

    ##### model
    logger.info("=> creating model ...")
    logger.info(f"Classes: {cfg.model.classes}")

    # model = xxx.build_model(cfg.model)
    # create model
    model = importlib.import_module(cfg.model.type)
    model = model.Network(cfg.model)

    use_cuda = torch.cuda.is_available()
    logger.info(f"cuda available: {use_cuda}")
    assert use_cuda
    model = model.cuda()

    # logger.info(model)
    logger.info(f"#classifier parameters (model): {sum([x.nelement() for x in model.parameters()])}")

    ##### load model
    utils.load_checkpoint(
        model, 
        cfg.pretrain,
        strict=False,
    )  # resume from the latest epoch, or specify the epoch to restore

    ##### evaluate
    test(model, cfg, logger)