This repo is the offical PyTorch implementation of MF-MOS.
import os
import numpy as np
import torch
from torch.utils.data import Dataset
from common.laserscan import LaserScan, SemLaserScan

import torch
import random

from collections.abc import Sequence, Iterable
from common.dataset.kitti.utils import load_poses, load_calib

# import math
# import types
# import numbers
# import warnings
# import torchvision
# from PIL import Image
# try:
# 	import accimage
# except ImportError:
# 	accimage = None

EXTENSIONS_SCAN = ['.bin']
EXTENSIONS_LABEL = ['.label']
EXTENSIONS_RESIDUAL = ['.npy']


def is_scan(filename):
    return any(filename.endswith(ext) for ext in EXTENSIONS_SCAN)


def is_label(filename):
    return any(filename.endswith(ext) for ext in EXTENSIONS_LABEL)


def is_residual(filename):
    return any(filename.endswith(ext) for ext in EXTENSIONS_RESIDUAL)


def my_collate(batch):
    data = [item[0] for item in batch]
    project_mask = [item[1] for item in batch]
    proj_labels = [item[2] for item in batch]
    data = torch.stack(data,dim=0)
    project_mask = torch.stack(project_mask,dim=0)
    proj_labels = torch.stack(proj_labels, dim=0)

    to_augment =(proj_labels == 12).nonzero()
    to_augment_unique_12 = torch.unique(to_augment[:, 0])

    to_augment = (proj_labels == 5).nonzero()
    to_augment_unique_5 = torch.unique(to_augment[:, 0])

    to_augment = (proj_labels == 8).nonzero()
    to_augment_unique_8 = torch.unique(to_augment[:, 0])

    to_augment_unique = torch.cat((to_augment_unique_5,to_augment_unique_8,to_augment_unique_12),dim=0)
    to_augment_unique = torch.unique(to_augment_unique)

    for k in to_augment_unique:
        data = torch.cat((data,torch.flip(data[k.item()], [2]).unsqueeze(0)),dim=0)
        proj_labels = torch.cat((proj_labels,torch.flip(proj_labels[k.item()], [1]).unsqueeze(0)),dim=0)
        project_mask = torch.cat((project_mask,torch.flip(project_mask[k.item()], [1]).unsqueeze(0)),dim=0)

    return data, project_mask,proj_labels


class SemanticKitti(Dataset):

    def __init__(self, root, # directory where data is
                 sequences,     # sequences for this data (e.g. [1,3,4,6])
                 labels,        # label dict: (e.g 10: "car")
                 residual_aug,  # use the data augmentation for residual maps (True or False)
                 color_map,     # colors dict bgr (e.g 10: [255, 0, 0])
                 learning_map,  # classes to learn (0 to N-1 for xentropy)
                 movable_learning_map,
                 learning_map_inv,    # inverse of previous (recover labels)
                 movable_learning_map_inv,
                 sensor,              # sensor to parse scans from
                 valid_residual_delta_t=1,  # modulation interval in data augmentation fro residual maps
                 max_points=150000,   # max number of points present in dataset
                 gt=True,             # send ground truth?
                 transform=False,
                 drop_few_static_frames=False):
        # save deats
        self.root = os.path.join(root, "sequences")
        self.sequences = sequences
        self.labels = labels
        self.residual_aug = residual_aug
        self.color_map = color_map
        self.learning_map = learning_map
        self.movable_learning_map = movable_learning_map
        self.learning_map_inv = learning_map_inv
        self.movable_learning_map_inv = movable_learning_map_inv
        self.sensor = sensor
        self.sensor_img_H = sensor["img_prop"]["height"]
        self.sensor_img_W = sensor["img_prop"]["width"]
        self.sensor_img_means = torch.tensor(sensor["img_means"], dtype=torch.float)
        self.sensor_img_stds = torch.tensor(sensor["img_stds"], dtype=torch.float)
        self.sensor_fov_up = sensor["fov_up"]
        self.sensor_fov_down = sensor["fov_down"]
        self.valid_residual_delta_t = valid_residual_delta_t
        self.max_points = max_points
        self.gt = gt
        self.transform = transform
        print("self.residual_aug {}, self.valid_residual_delta_t {}".format(self.residual_aug, self.valid_residual_delta_t))
        """
        Added stuff for dynamic object segmentation
        """
        # dictionary for mapping a dataset index to a sequence, frame_id tuple needed for using multiple frames
        self.dataset_size = 0
        self.index_mapping = {}
        dataset_index = 0
        # added this for dynamic object removal
        self.n_input_scans = sensor["n_input_scans"]  # This needs to be the same as in arch_cfg.yaml!
        self.use_residual = sensor["residual"]
        self.transform_mod = sensor["transform"]
        self.use_normal = sensor["use_normal"] if 'use_normal' in sensor.keys() else False
        """"""

        # get number of classes (can't be len(self.learning_map) because there
        # are multiple repeated entries, so the number that matters is how many
        # there are for the xentropy)
        self.nclasses = len(self.learning_map_inv)
        self.movable_nclasses = len(self.movable_learning_map_inv)

        # sanity checks

        # make sure directory exists
        if os.path.isdir(self.root):
            print("Sequences folder exists! Using sequences from %s" % self.root)
        else:
            raise ValueError("Sequences folder doesn't exist! Exiting...")

        assert(isinstance(self.labels, dict))	# make sure labels is a dict
        assert(isinstance(self.color_map, dict)) # make sure color_map is a dict
        assert(isinstance(self.learning_map, dict)) # make sure learning_map is a dict
        assert(isinstance(self.sequences, list)) # make sure sequences is a list

        if self.residual_aug:
            self.all_residaul_id = list(set([delta_t * i for i in range(self.n_input_scans) for delta_t in range(1, 4)]))
        else:
            self.all_residaul_id = [1 * i for i in range(self.n_input_scans)]

        # placeholder for filenames
        self.scan_files = {}
        self.label_files = {}
        self.poses = {}
        if self.use_residual:
            # for i in range(self.n_input_scans):
            for i in self.all_residaul_id:
                exec("self.residual_files_" + str(str(i+1)) + " = {}")

        # fill in with names, checking that all sequences are complete
        for seq in self.sequences:

            seq = '{0:02d}'.format(int(seq)) # to string
            print("parsing seq {}".format(seq))

            # get paths for each
            scan_path = os.path.join(self.root, seq, "velodyne")
            label_path = os.path.join(self.root, seq, "labels")

            if self.use_residual:
                # for i in range(self.n_input_scans):
                for i in self.all_residaul_id:
                    folder_name = "residual_images_" + str(i+1)
                    exec("residual_path_" + str(i+1) + " = os.path.join(self.root, seq, folder_name)")

            # get files
            scan_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
                    os.path.expanduser(scan_path)) for f in fn if is_scan(f)]
            label_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
                    os.path.expanduser(label_path)) for f in fn if is_label(f)]

            if self.use_residual:
                # for i in range(self.n_input_scans):
                for i in self.all_residaul_id:
                    exec("residual_files_" + str(i+1) + " = " + '[os.path.join(dp, f) for dp, dn, fn in '
                             'os.walk(os.path.expanduser(residual_path_' + str(i+1) + '))'
                             ' for f in fn if is_residual(f)]')

            ### Get poses and transform them to LiDAR coord frame for transforming point clouds
            # load poses
            pose_file = os.path.join(self.root, seq, "poses.txt")
            poses = np.array(load_poses(pose_file))
            inv_frame0 = np.linalg.inv(poses[0])

            # load calibrations
            calib_file = os.path.join(self.root, seq, "calib.txt")
            T_cam_velo = load_calib(calib_file)
            T_cam_velo = np.asarray(T_cam_velo).reshape((4, 4))
            T_velo_cam = np.linalg.inv(T_cam_velo)

            # convert kitti poses from camera coord to LiDAR coord
            new_poses = []
            for pose in poses:
                new_poses.append(T_velo_cam.dot(inv_frame0).dot(pose).dot(T_cam_velo))
            self.poses[seq] = np.array(new_poses)

            # check all scans have labels
            if self.gt:
                assert(len(scan_files) == len(label_files))

            """
            Added for dynamic object segmentation
            """
            # fill index mapper which is needed when loading several frames
            # n_used_files = max(0, len(scan_files) - self.n_input_scans + 1)  # this is used for multi-scan attach
            n_used_files = max(0, len(scan_files))  # this is used for multi residual images
            for start_index in range(n_used_files):
                self.index_mapping[dataset_index] = (seq, start_index)
                dataset_index += 1
            self.dataset_size += n_used_files
            """"""

            # extend list
            scan_files.sort()
            label_files.sort()

            self.scan_files[seq] = scan_files
            self.label_files[seq] = label_files

            if self.use_residual:
                # for i in range(self.n_input_scans):
                for i in self.all_residaul_id:
                    exec("residual_files_" + str(i+1) + ".sort()")
                    exec("self.residual_files_" + str(i+1) + "[seq]" + " = " + "residual_files_" + str(i+1))
        # print("\033[32m No model directory found.\033[0m")

        print(f"\033[32m There are {self.dataset_size} frames in total. \033[0m")
        if drop_few_static_frames:
            self.remove_few_static_frames()
            print(f"\033[32m Remove {self.total_remove} frames. \n New use {self.dataset_size} frames. \033[0m")

        print(f"\033[32m Using {self.dataset_size} scans from sequences {self.sequences}\033[0m")

    def __getitem__(self, dataset_index):

        # Get sequence and start
        seq, start_index = self.index_mapping[dataset_index]  # seq: '05' start_index: 1856
        # current_index = start_index + self.n_input_scans - 1  # this is used for multi-scan attach
        current_index = start_index  # this is used for multi residual images
        current_pose = self.poses[seq][current_index]
        proj_full = torch.Tensor()
        # index is now looping from first scan in input sequence to current scan
        # for index in range(start_index, start_index + self.n_input_scans):
        for index in range(start_index, start_index + 1):
            # get item in tensor shape
            scan_file = self.scan_files[seq][index]
            if self.gt:
                label_file = self.label_files[seq][index]

            if self.residual_aug:
                if self.transform:
                    residual_enhance_prob = random.random()
                    if residual_enhance_prob <= 0.5:
                        residual_input_scans_id = [1 * i for i in range(self.n_input_scans)]
                    elif residual_enhance_prob <= 0.75:
                        residual_input_scans_id = [2 * i for i in range(self.n_input_scans)]
                    else:
                        residual_input_scans_id = [3 * i for i in range(self.n_input_scans)]
                else:
                    residual_input_scans_id = [self.valid_residual_delta_t * i for i in range(self.n_input_scans)]
            else:
                residual_input_scans_id = [1 * i for i in range(self.n_input_scans)]

            if self.use_residual:
                # for i in range(self.n_input_scans):
                for i in residual_input_scans_id:
                    # exec("print(index, len(self.residual_files_" + str(i+1) + "[seq]), i+1, seq) if index >= len(self.residual_files_" + str(i+1) + "[seq]) else print()")
                    exec("residual_file_" + str(i+1) + " = " + "self.residual_files_" + str(i+1) + "[seq][index]")

            index_pose = self.poses[seq][index]

            # open a semantic laserscan
            DA = False
            flip_sign = False
            rot = False
            drop_points = False
            if self.transform:
                if random.random() > 0.5:
                    if random.random() > 0.5:
                            DA = True
                    if random.random() > 0.5:
                            flip_sign = True
                    if random.random() > 0.5:
                            rot = True
                    drop_points = random.uniform(0, 0.5)

            if self.gt:
                scan = SemLaserScan(self.color_map,
                                    project=True,
                                    H=self.sensor_img_H,
                                    W=self.sensor_img_W,
                                    fov_up=self.sensor_fov_up,
                                    fov_down=self.sensor_fov_down,
                                    DA=DA,
                                    flip_sign=flip_sign,
                                    drop_points=drop_points,
                                    use_normal=self.use_normal)
            else:
                scan = LaserScan(project=True,
                                 H=self.sensor_img_H,
                                 W=self.sensor_img_W,
                                 fov_up=self.sensor_fov_up,
                                 fov_down=self.sensor_fov_down,
                                 DA=DA,
                                 rot=rot,
                                 flip_sign=flip_sign,
                                 drop_points=drop_points,
                                 use_normal=self.use_normal)

            # open and obtain (transformed) scan
            scan.open_scan(scan_file, index_pose, current_pose, if_transform=self.transform_mod)

            if self.gt:
                scan.open_label(label_file)
                # map unused classes to used classes (also for projection)
                scan.sem_label = self.map(scan.sem_label, self.learning_map)
                scan.proj_sem_movable_label = scan.proj_sem_label.copy()

                scan.proj_sem_label = self.map(scan.proj_sem_label, self.learning_map)
                scan.proj_sem_movable_label = self.map(scan.proj_sem_movable_label, self.movable_learning_map)

            # make a tensor of the uncompressed data (with the max num points)
            unproj_n_points = scan.points.shape[0]
            unproj_xyz = torch.full((self.max_points, 3), -1.0, dtype=torch.float)
            unproj_xyz[:unproj_n_points] = torch.from_numpy(scan.points)
            unproj_range = torch.full([self.max_points], -1.0, dtype=torch.float)
            unproj_range[:unproj_n_points] = torch.from_numpy(scan.unproj_range)
            unproj_remissions = torch.full([self.max_points], -1.0, dtype=torch.float)
            unproj_remissions[:unproj_n_points] = torch.from_numpy(scan.remissions)
            if self.gt:
                unproj_labels = torch.full([self.max_points], -1.0, dtype=torch.int32)
                unproj_labels[:unproj_n_points] = torch.from_numpy(scan.sem_label)
            else:
                unproj_labels = []

            # get points and labels
            proj_range = torch.from_numpy(scan.proj_range).clone()
            proj_xyz = torch.from_numpy(scan.proj_xyz).clone()
            proj_remission = torch.from_numpy(scan.proj_remission).clone()
            proj_mask = torch.from_numpy(scan.proj_mask)
            if self.gt:
                proj_labels = torch.from_numpy(scan.proj_sem_label).clone()
                proj_labels = proj_labels * proj_mask

                proj_movable_labels = torch.from_numpy(scan.proj_sem_movable_label).clone()
                proj_movable_labels = proj_movable_labels * proj_mask
            else:
                proj_labels = []
                proj_movable_labels = []
            proj_x = torch.full([self.max_points], -1, dtype=torch.long)
            proj_x[:unproj_n_points] = torch.from_numpy(scan.proj_x)
            proj_y = torch.full([self.max_points], -1, dtype=torch.long)
            proj_y[:unproj_n_points] = torch.from_numpy(scan.proj_y)

            if self.use_residual:
                # for i in range(self.n_input_scans):
                for i in residual_input_scans_id:
                    exec("proj_residuals_" + str(i+1) + " = torch.Tensor(np.load(residual_file_" + str(i+1) + "))")
                    # if flip_sign:
                    #   exec("proj_residuals_" + str(i+1) + " = torch.flip(proj_residuals_" + str(i+1) + ", [-1])")


            proj = torch.cat([proj_range.unsqueeze(0).clone(),      # torch.Size([1, 64, 2048])
                              proj_xyz.clone().permute(2, 0, 1),    # torch.Size([3, 64, 2048])
                               proj_remission.unsqueeze(0).clone()]) # torch.Size([1, 64, 2048])
            proj = (proj - self.sensor_img_means[:, None, None]) / self.sensor_img_stds[:, None, None]

            proj_full = torch.cat([proj_full, proj])

        if self.use_normal:
            proj_full = torch.cat([proj_full, torch.from_numpy(scan.normal_map).clone().permute(2, 0, 1)]) # 5 + 3 = 8 channel
            # proj_full = torch.cat([proj_full, proj_xyz.clone().permute(2, 0, 1)]) # 5 + 3 = 8 channel

        # add residual channel
        if self.use_residual:
            # for i in range(self.n_input_scans):
            for i in residual_input_scans_id:
                proj_full = torch.cat([proj_full, torch.unsqueeze(eval("proj_residuals_" + str(i+1)), 0)])

        proj_full = proj_full * proj_mask.float()

        # get name and sequence
        path_norm = os.path.normpath(scan_file)
        path_split = path_norm.split(os.sep)
        path_seq = path_split[-3]
        path_name = path_split[-1].replace(".bin", ".label")

        ### A simple version Cutline operation by JiadaiSun ###
        # if int(seq) in [0, 1, 2, 3, 4, 5, 6, 7, 9, 10]: # only for training sequences
        # 	if random.random() > 0.5:
        # 		cutline = random.randint(self.sensor_img_W / 4, self.sensor_img_W * 3 / 4) # [0 -- 2048] ==> [512 - 1536]
        # 		proj_full = torch.cat((proj_full[:, :, cutline:], proj_full[:, :, :cutline]), dim=2)
        # 		proj_mask = torch.cat((proj_mask[:, cutline:], proj_mask[:, :cutline]), dim=1)
        # 		proj_labels = torch.cat((proj_labels[:, cutline:], proj_labels[:, :cutline]), dim=1)
        # else:
        # 	print("### In validation, don not use cutline augmentation ###")

        # if int(seq) in [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 30, 31, 32, 33, 34, 40]: # only for training sequences
        # 	if random.random() > 0.5:
        # 		cutline = random.randint(self.sensor_img_W / 4, self.sensor_img_W * 3 / 4) # [0 -- 2048] ==> [512 - 1536]
        # 		proj_full = torch.cat((proj_full[:, :, cutline:], proj_full[:, :, :cutline]), dim=2)  # torch.Size([13, 64, 2048])
        # 		proj_mask = torch.cat((proj_mask[:, cutline:], proj_mask[:, :cutline]), dim=1)        # torch.Size([64, 2048])
        # 		proj_labels = torch.cat((proj_labels[:, cutline:], proj_labels[:, :cutline]), dim=1)  # torch.Size([64, 2048])
        # 	if random.random() > 0.5:
        # 		proj_full = torch.flip(proj_full, [-1])
        # 		proj_mask = torch.flip(proj_mask, [-1])
        # 		proj_labels = torch.flip(proj_labels, [-1])

        return proj_full, proj_mask, (proj_labels, proj_movable_labels), unproj_labels, path_seq, path_name, proj_x, proj_y, proj_range, \
                     unproj_range, proj_xyz, unproj_xyz, proj_remission, unproj_remissions, unproj_n_points

    def __len__(self):
        return self.dataset_size

    @staticmethod
    def map(label, mapdict):
        # put label from original values to xentropy
        # or vice-versa, depending on dictionary values
        # make learning map a lookup table
        maxkey = 0
        for key, data in mapdict.items():
            if isinstance(data, list):
                nel = len(data)
            else:
                nel = 1
            if key > maxkey:
                maxkey = key
        # +100 hack making lut bigger just in case there are unknown labels
        if nel > 1:
            lut = np.zeros((maxkey + 100, nel), dtype=np.int32)
        else:
            lut = np.zeros((maxkey + 100), dtype=np.int32)
        for key, data in mapdict.items():
            try:
                lut[key] = data
            except IndexError:
                print("Wrong key ", key)
        # do the mapping
        return lut[label]

    def remove_few_static_frames(self):
        # Developed by Jiadai Sun 2021-11-07
        # This function is used to clear some frames, because too many static frames will lead to a long training time
        # There are several main dicts that need to be modified and processed
        #   self.scan_files, self.label_files, self.residual_files_1 ....8
        #   self.poses, self.index_mapping
        #   self.dataset_size (int number)

        remove_mapping_path = os.path.join(os.path.dirname(__file__), "../../../config/train_split_dynamic_pointnumber.txt")
        with open(remove_mapping_path) as fd:
            lines = fd.readlines()
            lines = [line.strip() for line in lines]

        pending_dict = {}
        for line in lines:
            if line != '':
                seq, fid, _ = line.split()
                if int(seq) in self.sequences:
                    if seq in pending_dict.keys():
                        if fid in pending_dict[seq]:
                            raise ValueError(f"!!!! Duplicate {fid} in seq {seq} in .txt file")
                        pending_dict[seq].append(fid)
                    else:
                        pending_dict[seq] = [fid]

        self.total_remove = 0
        self.index_mapping = {} # Reinitialize
        dataset_index = 0
        self.dataset_size = 0
        for seq in self.sequences:
            seq = '{0:02d}'.format(int(seq))
            if seq in pending_dict.keys():
                raw_len = len(self.scan_files[seq])

                # lidar scan files
                scan_files = self.scan_files[seq]
                useful_scan_paths = [path for path in scan_files if os.path.split(path)[-1][:-4] in pending_dict[seq]]
                self.scan_files[seq] = useful_scan_paths

                # label_files
                label_files = self.label_files[seq]
                useful_label_paths = [path for path in label_files if os.path.split(path)[-1][:-6] in pending_dict[seq]]
                self.label_files[seq] = useful_label_paths

                # poses_file
                self.poses[seq] = self.poses[seq][list(map(int, pending_dict[seq]))]

                assert len(useful_scan_paths) == len(useful_label_paths)
                assert len(useful_scan_paths) == self.poses[seq].shape[0]

                # the index_mapping and dataset_size is used in dataloader __getitem__
                n_used_files = max(0, len(useful_scan_paths))  # this is used for multi residual images
                for start_index in range(n_used_files):
                    self.index_mapping[dataset_index] = (seq, start_index)
                    dataset_index += 1
                self.dataset_size += n_used_files

                # More elegant implementation
                if self.use_residual:
                    # for i in range(self.n_input_scans):
                    for i in self.all_residaul_id:
                        tmp_residuals = eval(f"self.residual_files_{i+1}[\'{seq}\']")
                        tmp_pending_list = eval(f"pending_dict[\'{seq}\']")
                        tmp_usefuls = [path for path in tmp_residuals if os.path.split(path)[-1][:-4] in tmp_pending_list]
                        exec(f"self.residual_files_{i+1}[\'{seq}\'] = tmp_usefuls")
                        new_len = len(eval(f"self.residual_files_{i+1}[\'{seq}\']"))
                        print(f"  Drop residual_images_{i+1} in seq{seq}: {len(tmp_residuals)} -> {new_len}")
                        # if i >= 2:
                        #     exec(f"assert len(self.residual_files_{i-1}[\'{seq}\']) == len(self.residual_files_{i}[\'{seq}\'])")

                new_len = len(self.scan_files[seq])
                print(f"Seq {seq} drop {raw_len - new_len}: {raw_len} -> {new_len}")
                self.total_remove += raw_len - new_len


class Parser():
    # standard conv, BN, relu
    def __init__(self,
                 root,              # directory for data
                 train_sequences,   # sequences to train
                 valid_sequences,   # sequences to validate.
                 test_sequences,    # sequences to test (if none, don't get)
                 split,             # split (train, valid, test)
                 labels,            # labels in data
                 residual_aug,      # the data augmentation for residual maps
                 color_map,         # color for each label
                 learning_map,      # mapping for training labels
                 movable_learning_map,
                 learning_map_inv,  # recover labels from xentropy
                 movable_learning_map_inv,
                 sensor,            # sensor to use
                 max_points,        # max points in each scan in entire dataset
                 batch_size,        # batch size for train and val
                 workers,           # threads to load data
                 valid_residual_delta_t=1,  # modulation interval in data augmentation fro residual maps
                 gt=True,           # get gt?
                 shuffle_train=False):  # shuffle training set?
        super(Parser, self).__init__()

        # if I am training, get the dataset
        self.root = root
        self.train_sequences = train_sequences
        self.valid_sequences = valid_sequences
        self.test_sequences = test_sequences
        self.split = split
        self.labels = labels
        self.residual_aug = residual_aug
        self.valid_residual_delta_t = valid_residual_delta_t
        self.color_map = color_map
        self.learning_map = learning_map
        self.movable_learning_map = movable_learning_map
        self.learning_map_inv = learning_map_inv
        self.movable_learning_map_inv = movable_learning_map_inv
        self.sensor = sensor
        self.max_points = max_points
        self.batch_size = batch_size
        self.workers = workers
        self.gt = gt
        self.shuffle_train = shuffle_train

        # number of classes that matters is the one for xentropy
        self.nclasses = len(self.learning_map_inv)
        self.movable_nclasses = len(self.movable_learning_map_inv)

        # Data loading code
        if self.split == 'train':
            self.train_dataset = SemanticKitti(root=self.root,
                                               sequences=self.train_sequences,
                                               labels=self.labels,
                                               residual_aug=self.residual_aug,
                                               valid_residual_delta_t=self.valid_residual_delta_t,
                                               color_map=self.color_map,
                                               learning_map=self.learning_map,
                                               movable_learning_map=self.movable_learning_map,
                                               learning_map_inv=self.learning_map_inv,
                                               movable_learning_map_inv=self.movable_learning_map_inv,
                                               sensor=self.sensor,
                                               max_points=max_points,
                                               transform=True,
                                               gt=self.gt,
                                               drop_few_static_frames=True)

            shuffle_train = True
            if torch.cuda.is_available() and torch.cuda.device_count() < 2:
                self.train_sampler = None
            else:
                self.train_sampler = torch.utils.data.distributed.DistributedSampler(self.train_dataset)
                shuffle_train = False
            # self.trainloader = torch.utils.data.DataLoader(self.train_dataset,
            #                                                batch_size=self.batch_size,
            #                                                shuffle=self.shuffle_train,
            #                                                # shuffle=False,
            #                                                num_workers=self.workers,
            #                                                pin_memory=True,
            #                                                drop_last=True)
            self.trainloader = torch.utils.data.DataLoader(self.train_dataset,
                                                           batch_size=self.batch_size,
                                                           shuffle=shuffle_train,
                                                           num_workers=self.workers,
                                                           pin_memory=True,
                                                           drop_last=True,
                                                           sampler=self.train_sampler)

            assert len(self.trainloader) > 0
            self.trainiter = iter(self.trainloader)

            self.valid_dataset = SemanticKitti(root=self.root,
                                               sequences=self.valid_sequences,
                                               labels=self.labels,
                                               residual_aug=self.residual_aug,
                                               valid_residual_delta_t=self.valid_residual_delta_t,
                                               color_map=self.color_map,
                                               learning_map=self.learning_map,
                                               movable_learning_map=self.movable_learning_map,
                                               learning_map_inv=self.learning_map_inv,
                                               movable_learning_map_inv=self.movable_learning_map_inv,
                                               sensor=self.sensor,
                                               max_points=max_points,
                                               gt=self.gt)

            self.validloader = torch.utils.data.DataLoader(self.valid_dataset,
                                                           batch_size=self.batch_size,
                                                           shuffle=False,
                                                           num_workers=self.workers,
                                                           pin_memory=True,
                                                           drop_last=True)
            assert len(self.validloader) > 0
            self.validiter = iter(self.validloader)

        if self.split == 'valid':
            self.valid_dataset = SemanticKitti(root=self.root,
                                               sequences=self.valid_sequences,
                                               labels=self.labels,
                                               residual_aug=self.residual_aug,
                                               valid_residual_delta_t=self.valid_residual_delta_t,
                                               color_map=self.color_map,
                                               learning_map=self.learning_map,
                                               movable_learning_map=self.movable_learning_map,
                                               learning_map_inv=self.learning_map_inv,
                                               movable_learning_map_inv=self.movable_learning_map_inv,
                                               sensor=self.sensor,
                                               max_points=max_points,
                                               gt=self.gt)

            self.validloader = torch.utils.data.DataLoader(self.valid_dataset,
                                                           batch_size=self.batch_size,
                                                           shuffle=False,
                                                           num_workers=self.workers,
                                                           pin_memory=True,
                                                           drop_last=True)
            assert len(self.validloader) > 0
            self.validiter = iter(self.validloader)

        if self.split == 'test':
            if self.test_sequences:
                self.test_dataset = SemanticKitti(root=self.root,
                                                  sequences=self.test_sequences,
                                                  labels=self.labels,
                                                  residual_aug=self.residual_aug,
                                                  valid_residual_delta_t=self.valid_residual_delta_t,
                                                  color_map=self.color_map,
                                                  learning_map=self.learning_map,
                                                  learning_map_inv=self.learning_map_inv,
                                                  movable_learning_map_inv=self.movable_learning_map_inv,
                                                  sensor=self.sensor,
                                                  max_points=max_points,
                                                  gt=False)

                self.testloader = torch.utils.data.DataLoader(self.test_dataset,
                                                              batch_size=self.batch_size,
                                                              shuffle=False,
                                                              num_workers=self.workers,
                                                              pin_memory=True,
                                                              drop_last=True)
                assert len(self.testloader) > 0
                self.testiter = iter(self.testloader)

    def get_train_batch(self):
        scans = self.trainiter.next()
        return scans

    def get_train_set(self):
        return self.trainloader

    def get_valid_batch(self):
        scans = self.validiter.next()
        return scans

    def get_valid_set(self):
        return self.validloader

    def get_test_batch(self):
        scans = self.testiter.next()
        return scans

    def get_test_set(self):
        return self.testloader

    def get_train_size(self):
        return len(self.trainloader)

    def get_valid_size(self):
        return len(self.validloader)

    def get_test_size(self):
        return len(self.testloader)

    def get_n_classes(self, movable=False):
        if not movable:
            return self.nclasses
        else:
            return self.movable_nclasses

    def get_original_class_string(self, idx):
        return self.labels[idx]

    def get_xentropy_class_string(self, idx, movable=False):
        if not movable:
            return self.labels[self.learning_map_inv[idx]]
        else:
            return self.labels[self.movable_learning_map_inv[idx]]

    def to_original(self, label, movable=False):
        # put label in original values
        if not movable:
            return SemanticKitti.map(label, self.learning_map_inv)
        else:
            return SemanticKitti.map(label, self.movable_learning_map_inv)

    def to_xentropy(self, label, movable=False):
        # put label in xentropy values
        if not movable:
            return SemanticKitti.map(label, self.learning_map)
        else:
            return SemanticKitti.map(label, self.movable_learning_map)

    def to_color(self, label, movable=False):
        if not movable:
            # put label in original values
            label = SemanticKitti.map(label, self.learning_map_inv)
        else:
            # put label in original values
            label = SemanticKitti.map(label, self.movable_learning_map_inv)
        # put label in color
        return SemanticKitti.map(label, self.color_map)
common/dataset/kitti/parser.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
import open3d as o3d

import numpy as np
import random
from scipy.spatial.transform import Rotation as R

# import math
# import time

class LaserScan:
    """ Class that contains LaserScan with x,y,z,r,intensity """
    EXTENSIONS_SCAN = ['.bin']

    def __init__(self, project=False, H=64, W=1024, fov_up=3.0, fov_down=-25.0, DA=False,
                 flip_sign=False, rot=False, drop_points=False, use_normal=False):
        self.project = project
        self.proj_H = H
        self.proj_W = W
        self.proj_fov_up = fov_up
        self.proj_fov_down = fov_down
        self.DA = DA
        self.flip_sign = flip_sign
        self.rot = rot
        self.drop_points = drop_points
        self.use_normal = use_normal

        self.reset()

    def reset(self):
        """ Reset scan members. """
        self.points = np.zeros((0, 3), dtype=np.float32)  # [m, 3]: x, y, z
        self.remissions = np.zeros((0, 1), dtype=np.float32)  # [m ,1]: remission

        # projected range image - [H,W] range (-1 is no data)
        self.proj_range = np.full((self.proj_H, self.proj_W), -1, dtype=np.float32)

        # unprojected range (list of depths for each point)
        self.unproj_range = np.zeros((0, 1), dtype=np.float32)

        # projected point cloud xyz - [H,W,3] xyz coord (-1 is no data)
        self.proj_xyz = np.full((self.proj_H, self.proj_W, 3), -1, dtype=np.float32)

        # projected remission - [H,W] intensity (-1 is no data)
        self.proj_remission = np.full((self.proj_H, self.proj_W), -1, dtype=np.float32)

        # projected index (for each pixel, what I am in the pointcloud)
        # [H,W] index (-1 is no data)
        self.proj_idx = np.full((self.proj_H, self.proj_W), -1, dtype=np.int32)

        # for each point, where it is in the range image
        self.proj_x = np.zeros((0, 1), dtype=np.int32)  # [m, 1]: x
        self.proj_y = np.zeros((0, 1), dtype=np.int32)  # [m, 1]: y

        # mask containing for each pixel, if it contains a point or not
        self.proj_mask = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W] mask
        # if self.use_normal:
        #     self.proj_vertex = np.full((self.proj_H, self.proj_W, 3), -1, dtype=np.float32)

    def size(self):
        """ Return the size of the point cloud. """
        return self.points.shape[0]

    def __len__(self):
        return self.size()

    def open_scan(self, filename, from_pose, to_pose, if_transform=True):
        """ Open raw scan and fill in attributes
        """
        # reset just in case there was an open structure
        self.reset()

        # check filename is string
        if not isinstance(filename, str):
            raise TypeError("Filename should be string type, "
                            "but was {type}".format(type=str(type(filename))))

        # check extension is a laserscan
        if not any(filename.endswith(ext) for ext in self.EXTENSIONS_SCAN):
            raise RuntimeError("Filename extension is not valid scan file.")

        # if all goes well, open pointcloud
        scan = np.fromfile(filename, dtype=np.float32)
        scan = scan.reshape((-1, 4))

        # put in attribute
        points = scan[:, 0:3]  # get xyz
        """
        Transform points to current frame for dynamic object segmentation
        """
        hom_points = np.ones(scan.shape)
        hom_points[:, :-1] = points
        if if_transform:
            points_transformed = np.linalg.inv(to_pose).dot(from_pose).dot(hom_points.T).T
        else:
            points_transformed = hom_points
        """"""
        points = points_transformed[:, :3]
        remissions = scan[:, 3]  # get remission
        if self.drop_points is not False:
            self.points_to_drop = np.random.randint(0, len(points)-1,int(len(points)*self.drop_points))
            points = np.delete(points,self.points_to_drop,axis=0)
            remissions = np.delete(remissions,self.points_to_drop)

        self.set_points(points, remissions)

    def set_points(self, points, remissions=None):
        """ Set scan attributes (instead of opening from file)
        """
        # reset just in case there was an open structure
        self.reset()

        # check scan makes sense
        if not isinstance(points, np.ndarray):
            raise TypeError("Scan should be numpy array")

        # check remission makes sense
        if remissions is not None and not isinstance(remissions, np.ndarray):
            raise TypeError("Remissions should be numpy array")

        # put in attribute
        self.points = points  # get
        if self.flip_sign:
            self.points[:, 1] = -self.points[:, 1]
        if self.DA:
            jitter_x = random.uniform(-5,5)
            jitter_y = random.uniform(-3, 3)
            jitter_z = random.uniform(-1, 0)
            self.points[:, 0] += jitter_x
            self.points[:, 1] += jitter_y
            self.points[:, 2] += jitter_z
        if self.rot:
            self.points = self.points @ R.random(random_state=1234).as_dcm().T
        if remissions is not None:
            self.remissions = remissions  # get remission
            #if self.DA:
            #    self.remissions = self.remissions[::-1].copy()
        else:
            self.remissions = np.zeros((points.shape[0]), dtype=np.float32)

        # if projection is wanted, then do it and fill in the structure
        if self.project:
            self.do_range_projection()

    def do_range_projection(self):
        """ Project a pointcloud into a spherical projection image.projection.
            Function takes no arguments because it can be also called externally
            if the value of the constructor was not set (in case you change your
            mind about wanting the projection)
        """
        # laser parameters
        fov_up = self.proj_fov_up / 180.0 * np.pi  # field of view up in rad
        fov_down = self.proj_fov_down / 180.0 * np.pi  # field of view down in rad
        fov = abs(fov_down) + abs(fov_up)  # get field of view total in rad

        # get depth of all points
        depth = np.linalg.norm(self.points, 2, axis=1)

        # get scan components
        scan_x = self.points[:, 0]
        scan_y = self.points[:, 1]
        scan_z = self.points[:, 2]

        # get angles of all points
        yaw = -np.arctan2(scan_y, scan_x)
        pitch = np.arcsin(scan_z / depth)

        # get projections in image coords
        proj_x = 0.5 * (yaw / np.pi + 1.0)  # in [0.0, 1.0]
        proj_y = 1.0 - (pitch + abs(fov_down)) / fov  # in [0.0, 1.0]

        # scale to image size using angular resolution
        proj_x *= self.proj_W  # in [0.0, W]
        proj_y *= self.proj_H  # in [0.0, H]

        # round and clamp for use as index
        proj_x = np.floor(proj_x)
        proj_x = np.minimum(self.proj_W - 1, proj_x)
        proj_x = np.maximum(0, proj_x).astype(np.int32)  # in [0,W-1]
        self.proj_x = np.copy(proj_x)  # store a copy in orig order

        proj_y = np.floor(proj_y)
        proj_y = np.minimum(self.proj_H - 1, proj_y)
        proj_y = np.maximum(0, proj_y).astype(np.int32)  # in [0,H-1]
        self.proj_y = np.copy(proj_y)  # stope a copy in original order

        # copy of depth in original order
        self.unproj_range = np.copy(depth)

        # order in decreasing depth
        indices = np.arange(depth.shape[0])
        order = np.argsort(depth)[::-1]
        depth = depth[order]
        indices = indices[order]
        points = self.points[order]
        remission = self.remissions[order]
        proj_y = proj_y[order]
        proj_x = proj_x[order]

        # assing to images
        self.proj_range[proj_y, proj_x] = depth
        self.proj_xyz[proj_y, proj_x] = points
        self.proj_remission[proj_y, proj_x] = remission
        self.proj_idx[proj_y, proj_x] = indices
        self.proj_mask = (self.proj_idx > 0).astype(np.int32)

        if self.use_normal:
            # Plan A:
            # scan_x = scan_x[order]
            # scan_y = scan_y[order]
            # scan_z = scan_z[order]
            # self.proj_vertex[proj_y, proj_x] = np.array([scan_x, scan_y, scan_z]).T
            # self.normal_map = self.gen_normal_map(self.proj_range, self.proj_vertex, self.proj_H, self.proj_W)

            # Plan B:
            normals = self.gen_normal_map_open3d(points=points)
            self.normal_map = np.full((self.proj_H, self.proj_W, 3), -1, dtype=np.float32)
            self.normal_map[proj_y, proj_x] = normals

    def gen_normal_map_open3d(self, points):
        # If numpy is 1.18.1 (open3d 0.9.0 py3.7), estimate_normals will leak memory, 
        #  and then I updated it to 1.19.4 according to the issue below
        #  https://github.com/isl-org/Open3D/issues/1787
        #  https://github.com/isl-org/Open3D/issues/2107
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        # Calculate normal, search radius 20cm, only consider 15 points in the neighborhood
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.2, max_nn=15))
        ## Calculate normals, considering only 20 points in the neighborhood
        # pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=10))
        normals = np.asarray(pcd.normals)
        del pcd
        return normals #np.asarray(pcd.normals)

    # This function refers to https://github.com/PRBonn/range-mcl ,
    #  but it runs very slowly
    def gen_normal_map(self, current_range, current_vertex, proj_H=64, proj_W=900):
        """ Generate a normal image given the range projection of a point cloud.
            Args:
            current_range:  range projection of a point cloud, each pixel contains the corresponding depth
            current_vertex: range projection of a point cloud,
                            each pixel contains the corresponding point (x, y, z, 1)
            Returns:
            normal_data: each pixel contains the corresponding normal
        """
        def wrap(x, dim):
            """ Wrap the boarder of the range image.
            """
            value = x
            if value >= dim:
                value = (value - dim)
            if value < 0:
                value = (value + dim)
            return value

        normal_data = np.full((proj_H, proj_W, 3), -1, dtype=np.float32)

        # iterate over all pixels in the range image
        for x in range(proj_W):
            for y in range(proj_H - 1):
                p = current_vertex[y, x][:3]
                depth = current_range[y, x]

                if depth > 0:
                    wrap_x = wrap(x + 1, proj_W)
                    u = current_vertex[y, wrap_x][:3]
                    u_depth = current_range[y, wrap_x]
                    if u_depth <= 0:
                        continue

                    v = current_vertex[y + 1, x][:3]
                    v_depth = current_range[y + 1, x]
                    if v_depth <= 0:
                        continue

                    u_norm = (u - p) / np.linalg.norm(u - p)
                    v_norm = (v - p) / np.linalg.norm(v - p)

                    w = np.cross(v_norm, u_norm)
                    norm = np.linalg.norm(w)
                    if norm > 0:
                        normal = w / norm
                        normal_data[y, x] = normal

        return normal_data

class SemLaserScan(LaserScan):
    """Class that contains LaserScan with x,y,z,r,sem_label,sem_color_label,inst_label,inst_color_label"""
    EXTENSIONS_LABEL = ['.label']

    def __init__(self, sem_color_dict=None, project=False,
                 H=64, W=1024, fov_up=3.0, fov_down=-25.0, max_classes=300,
                 DA=False, flip_sign=False, drop_points=False, use_normal=False):
        super(SemLaserScan, self).__init__(project, H, W, fov_up, fov_down, DA=DA,
                                           flip_sign=flip_sign, drop_points=drop_points, use_normal=use_normal)
        self.reset()

        # make semantic colors
        if sem_color_dict:
            # if I have a dict, make it
            max_sem_key = 0
            for key, data in sem_color_dict.items():
                if key + 1 > max_sem_key:
                    max_sem_key = key + 1
            self.sem_color_lut = np.zeros((max_sem_key + 100, 3), dtype=np.float32)
            for key, value in sem_color_dict.items():
                self.sem_color_lut[key] = np.array(value, np.float32) / 255.0
        else:
            # otherwise make random
            max_sem_key = max_classes
            self.sem_color_lut = np.random.uniform(low=0.0, high=1.0, size=(max_sem_key, 3))
            # force zero to a gray-ish color
            self.sem_color_lut[0] = np.full((3), 0.1)

        # make instance colors
        max_inst_id = 100000
        self.inst_color_lut = np.random.uniform(low=0.0, high=1.0, size=(max_inst_id, 3))
        # force zero to a gray-ish color
        self.inst_color_lut[0] = np.full((3), 0.1)

    def reset(self):
        """ Reset scan members. """
        super(SemLaserScan, self).reset()

        # semantic labels
        self.sem_label = np.zeros((0, 1), dtype=np.int32)  # [m, 1]: label
        self.sem_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        # instance labels
        self.inst_label = np.zeros((0, 1), dtype=np.int32)  # [m, 1]: label
        self.inst_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        # projection color with semantic labels
        self.proj_sem_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W]  label
        self.proj_sem_movable_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W]  label
        self.proj_sem_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float)  # [H,W,3] color

        # projection color with instance labels
        self.proj_inst_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W]  label
        self.proj_inst_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float)  # [H,W,3] color

    def open_label(self, filename):
        """ Open raw scan and fill in attributes """
        # check filename is string
        if not isinstance(filename, str):
            raise TypeError("Filename should be string type, "
                            "but was {type}".format(type=str(type(filename))))

        # check extension is a laserscan
        if not any(filename.endswith(ext) for ext in self.EXTENSIONS_LABEL):
            raise RuntimeError("Filename extension is not valid label file.")

        # if all goes well, open label
        label = np.fromfile(filename, dtype=np.int32)
        label = label.reshape((-1))

        if self.drop_points is not False:
            label = np.delete(label,self.points_to_drop)
        # set it
        self.set_label(label)

    def set_label(self, label):
        """ Set points for label not from file but from np """
        # check label makes sense
        if not isinstance(label, np.ndarray):
            raise TypeError("Label should be numpy array")

        # only fill in attribute if the right size
        if label.shape[0] == self.points.shape[0]:
            self.sem_label = label & 0xFFFF  # semantic label in lower half
            self.inst_label = label >> 16  # instance id in upper half
        else:
            print("Points shape: ", self.points.shape)
            print("Label shape: ", label.shape)
            raise ValueError("Scan and Label don't contain same number of points")

        # sanity check
        assert ((self.sem_label + (self.inst_label << 16) == label).all())

        if self.project:
            self.do_label_projection()

    def colorize(self):
        """ Colorize pointcloud with the color of each semantic label """
        self.sem_label_color = self.sem_color_lut[self.sem_label]
        self.sem_label_color = self.sem_label_color.reshape((-1, 3))

        self.inst_label_color = self.inst_color_lut[self.inst_label]
        self.inst_label_color = self.inst_label_color.reshape((-1, 3))

    def do_label_projection(self):
        # only map colors to labels that exist
        mask = self.proj_idx >= 0

        # semantics
        self.proj_sem_label[mask] = self.sem_label[self.proj_idx[mask]]
        self.proj_sem_color[mask] = self.sem_color_lut[self.sem_label[self.proj_idx[mask]]]

        # instances
        self.proj_inst_label[mask] = self.inst_label[self.proj_idx[mask]]
        self.proj_inst_color[mask] = self.inst_color_lut[self.inst_label[self.proj_idx[mask]]]

common/laserscan.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
from utils.utils import *
from modules.user import *
from modules.user_refine import *

if __name__ == '__main__':
    
    parser = get_args(flags="infer")
    FLAGS, unparsed = parser.parse_known_args()

    print("----------")
    print("INTERFACE:")
    print("  dataset", FLAGS.dataset)
    print("  log", FLAGS.log)
    print("  model", FLAGS.model)
    print("  infering", FLAGS.split)
    print("  pointrefine", FLAGS.pointrefine)
    print("  save movable", FLAGS.movable)
    print("----------\n")

    # open arch / data config file
    ARCH = load_yaml(FLAGS.model + "/arch_cfg.yaml")
    DATA = load_yaml(FLAGS.model + "/data_cfg.yaml")

    make_predictions_dir(FLAGS, DATA, save_movable=FLAGS.movable) # create predictions file folder
    check_model_dir(FLAGS.model)      # does model folder exist?

    # create user and infer dataset
    if not FLAGS.pointrefine:
        user = User(ARCH, DATA, datadir=FLAGS.dataset, outputdir=FLAGS.log,
                    modeldir=FLAGS.model, split=FLAGS.split, save_movable=FLAGS.movable)
    else:
        user = UserRefine(ARCH, DATA, datadir=FLAGS.dataset, outputdir=FLAGS.log,
                          modeldir=FLAGS.model, split=FLAGS.split, save_movable=FLAGS.movable)
    user.infer()

infer.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import torch
import torch.nn as nn
import torch.nn.functional as F

# import softpool_cuda
# from SoftPool import soft_pool2d, SoftPool2d
from typing import Tuple

@torch.jit.script
def pixelshuffle(x: torch.Tensor, factor_hw: Tuple[int, int]):
    pH = factor_hw[0]
    pW = factor_hw[1]
    y = x
    B, iC, iH, iW = y.shape
    oC, oH, oW = iC//(pH*pW), iH*pH, iW*pW
    y = y.reshape(B, oC, pH, pW, iH, iW)
    y = y.permute(0, 1, 4, 2, 5, 3)     # B, oC, iH, pH, iW, pW
    y = y.reshape(B, oC, oH, oW)
    return y

class ResContextBlock(nn.Module):
    def __init__(self, in_filters, out_filters):
        super(ResContextBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_filters, out_filters, kernel_size=(1, 1), stride=1)
        self.act1 = nn.LeakyReLU()

        self.conv2 = nn.Conv2d(out_filters, out_filters, (3, 3), padding=1)
        self.act2 = nn.LeakyReLU()
        self.bn1 = nn.BatchNorm2d(out_filters)

        self.conv3 = nn.Conv2d(out_filters, out_filters, (3, 3), dilation=2, padding=2)
        self.act3 = nn.LeakyReLU()
        self.bn2 = nn.BatchNorm2d(out_filters)

    def forward(self, x):
        shortcut = self.conv1(x)
        shortcut = self.act1(shortcut)

        resA = self.conv2(shortcut)
        resA = self.act2(resA)
        resA1 = self.bn1(resA)

        resA = self.conv3(resA1)
        resA = self.act3(resA)
        resA2 = self.bn2(resA)

        output = shortcut + resA2
        return output

class ResBlock(nn.Module):
    def __init__(self, in_filters, out_filters, dropout_rate, 
                 kernel_size=(3, 3), stride=1, pooling=True, drop_out=True, use_softpool=False):
        super(ResBlock, self).__init__()
        self.pooling = pooling
        self.drop_out = drop_out
        self.conv1 = nn.Conv2d(in_filters, out_filters, kernel_size=(1, 1), stride=stride)
        self.act1 = nn.LeakyReLU()

        self.conv2 = nn.Conv2d(in_filters, out_filters, kernel_size=(3,3), padding=1)
        self.act2 = nn.LeakyReLU()
        self.bn1 = nn.BatchNorm2d(out_filters)

        self.conv3 = nn.Conv2d(out_filters, out_filters, kernel_size=(3,3),dilation=2, padding=2)
        self.act3 = nn.LeakyReLU()
        self.bn2 = nn.BatchNorm2d(out_filters)

        self.conv4 = nn.Conv2d(out_filters, out_filters, kernel_size=(2, 2), dilation=2, padding=1)
        self.act4 = nn.LeakyReLU()
        self.bn3 = nn.BatchNorm2d(out_filters)

        self.conv5 = nn.Conv2d(out_filters*3, out_filters, kernel_size=(1, 1))
        self.act5 = nn.LeakyReLU()
        self.bn4 = nn.BatchNorm2d(out_filters)

        if pooling:
            self.dropout = nn.Dropout2d(p=dropout_rate)
            if use_softpool:
                self.pool = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
            else:
                # self.pool = nn.AvgPool2d(kernel_size=kernel_size, stride=2, padding=1)
                self.pool = nn.AvgPool2d(kernel_size=kernel_size)
        else:
            self.dropout = nn.Dropout2d(p=dropout_rate)

    def forward(self, x):
        shortcut = self.conv1(x)    # (3, 64, 64, 2048)
        shortcut = self.act1(shortcut)

        resA = self.conv2(x)
        resA = self.act2(resA)
        resA1 = self.bn1(resA)

        resA = self.conv3(resA1)    # (3, 64, 64, 2048)
        resA = self.act3(resA)
        resA2 = self.bn2(resA)

        resA = self.conv4(resA2)    # (3, 64, 64, 2048)
        resA = self.act4(resA)
        resA3 = self.bn3(resA)

        concat = torch.cat((resA1, resA2, resA3), dim=1)    # (3, 192, 64, 2048)
        resA = self.conv5(concat)   # (3, 64, 64, 2048)
        resA = self.act5(resA)
        resA = self.bn4(resA)
        resA = shortcut + resA  # (3, 64, 64, 2048)

        if self.pooling:
            resB = self.dropout(resA) if self.drop_out else resA
            resB = self.pool(resB)

            # H, W = resB.shape[2:]
            # resB = resB.flatten(2).transpose(1, 2)
            # return resB, H, W
            return resB, resA
        else:
            resB = self.dropout(resA) if self.drop_out else resA
            return resB


class UpBlock(nn.Module):
    def __init__(self, in_filters, out_filters, dropout_rate, drop_out=True):
        super(UpBlock, self).__init__()
        self.drop_out = drop_out
        self.in_filters = in_filters
        self.out_filters = out_filters

        self.dropout1 = nn.Dropout2d(p=dropout_rate)
        self.dropout2 = nn.Dropout2d(p=dropout_rate)
        self.dropout3 = nn.Dropout2d(p=dropout_rate)

        # self.conv1 = nn.Conv2d(in_filters // 4 + 2*out_filters, out_filters, (3, 3), padding=1)
        self.conv1 = nn.Conv2d(in_filters // 8 + 2 * out_filters, out_filters, (3, 3), padding=1)
        self.act1 = nn.LeakyReLU()
        self.bn1 = nn.BatchNorm2d(out_filters)

        self.conv2 = nn.Conv2d(out_filters, out_filters, (3, 3), dilation=2, padding=2)
        self.act2 = nn.LeakyReLU()
        self.bn2 = nn.BatchNorm2d(out_filters)

        self.conv3 = nn.Conv2d(out_filters, out_filters, (2, 2), dilation=2, padding=1)
        self.act3 = nn.LeakyReLU()
        self.bn3 = nn.BatchNorm2d(out_filters)

        self.conv4 = nn.Conv2d(out_filters*3, out_filters, kernel_size=(1, 1))
        self.act4 = nn.LeakyReLU()
        self.bn4 = nn.BatchNorm2d(out_filters)

    def forward(self, x, skip):
        # upA = nn.PixelShuffle(2)(x)
        upA = pixelshuffle(x, (2, 4))
        if self.drop_out:
            upA = self.dropout1(upA)

        upB = torch.cat((upA, skip), dim=1)
        if self.drop_out:
            upB = self.dropout2(upB)

        upE = self.conv1(upB)
        upE = self.act1(upE)
        upE1 = self.bn1(upE)

        upE = self.conv2(upE1)
        upE = self.act2(upE)
        upE2 = self.bn2(upE)

        upE = self.conv3(upE2)
        upE = self.act3(upE)
        upE3 = self.bn3(upE)

        concat = torch.cat((upE1, upE2, upE3), dim=1)
        upE = self.conv4(concat)
        upE = self.act4(upE)
        upE = self.bn4(upE)
        if self.drop_out:
            upE = self.dropout3(upE)

        return upE


class MetaKernel(nn.Module):
    def __init__(self, num_batch, feat_height, feat_width, coord_channels, fp16=False, num_frame=1):
        super(MetaKernel, self).__init__()
        self.num_batch = num_batch
        self.H = feat_height
        self.W = feat_width
        self.fp16 = fp16
        self.num_frame = num_frame
        self.channel_list = [16, 32]
        self.use_norm = False
        self.coord_channels = coord_channels
        print("meta num_batch is : ", self.num_batch)
        # TODO: Best to put this part of the definition in __init__ function ?
        self.conv1 = nn.Conv2d(self.coord_channels,
                               self.channel_list[0],
                               kernel_size=(1, 1),
                               stride=(1, 1),
                               padding=0,
                               dilation=1,
                               bias=True)
        if self.use_norm:
            self.bn1 = nn.BatchNorm2d(self.channel_list[0])
        # self.act1 = nn.ReLU()
        self.act1 = nn.LeakyReLU()

        self.conv2 = nn.Conv2d(self.channel_list[0],
                               self.channel_list[1],
                               kernel_size=(1, 1),
                               stride=(1, 1),
                               padding=0,
                               dilation=1,
                               bias=True)
        self.conv1x1 = nn.Conv2d(in_channels=288, out_channels=32, kernel_size=1, stride=1, padding=0)

    def update_num_batch(self, num_batch):
        self.num_batch = num_batch

    @staticmethod
    def sampler_im2col(data, name=None, kernel=1, stride=1, pad=None, dilate=1):
        """ please refer to mx.symbol.im2col """
        if isinstance(kernel, int):
            kernel = (kernel, kernel)
        if isinstance(stride, int):
            stride = (stride, stride)
        if isinstance(dilate, int):
            dilate = (dilate, dilate)
        if pad is None:
            assert kernel[0] % 2 == 1, "Specify pad for an even kernel size in function sampler_im2col"
            pad = ((kernel[0] - 1) * dilate[0] + 1) // 2
        if isinstance(pad, int):
            pad = (pad, pad)

        # https://mxnet.apache.org/versions/1.8.0/api/python/docs/api/symbol/symbol.html#mxnet.symbol.im2col
        # output = mx.symbol.im2col(
        #     name=name + "sampler",
        #     data=data,
        #     kernel=kernel,
        #     stride=stride,
        #     dilate=dilate,
        #     pad=pad)

        # https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html
        # torch._C._nn.im2col(input, _pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride))
        output = F.unfold(data, kernel_size=kernel, 
                                dilation=dilate,
                                stride=stride, 
                                padding=pad)

        return output

    def sample_data(self, data, kernel_size):
        """
        data sample
        :param data: num_batch, num_channel_in, H, W
        :param kernel_size: int default=3
        :return: sample_output: num_batch, num_channel_in * kernel_size * kernel_size, H, W
        """
        sample_output = self.sampler_im2col(
            data=data,
            kernel=kernel_size,
            stride=1,
            pad=1,
            dilate=1)
        return sample_output

    def sample_coord(self, coord, kernel_size):
        """
        coord sample
        :param coord: num_batch, num_channel_in, H, W
        :param kernel_size: int default=3
        :return: coord_sample_data: num_batch, num_channel_in * kernel_size * kernel_size, H, W
        """
        coord_sample_data = self.sampler_im2col(
            data=coord,
            kernel=kernel_size,
            stride=1,
            pad=1,
            dilate=1)

        return coord_sample_data

    def relative_coord(self, sample_coord, center_coord, num_channel_in, kernel_size):
        """
        :param sample_coord: num_batch, num_channel_in * kernel_size * kernel_size, H, W
        :param center_coord: num_batch, num_channel_in, H, W
        :param num_channel_in: int
        :param kernel_size: int
        :return: rel_coord: num_batch, num_channel_in, kernel_size * kernel_size, H, W
        """
        sample_reshape = torch.reshape(
            sample_coord,
            shape=(
                self.num_batch,
                num_channel_in,
                kernel_size * kernel_size,
                self.H,
                self.W
            )
        )
        # ic(sample_reshape.size())

        center_coord_expand = torch.unsqueeze(center_coord, dim=2)  # expand_dims
        # ic(center_coord_expand.size())

        rel_coord = torch.subtract(sample_reshape, center_coord_expand)
        # ic(rel_coord.size())

        return rel_coord

    def mlp(self, data, in_channels, channel_list=None, b_mul=1, use_norm=False):
        """
        :param data: num_batch, num_channel_in * kernel_size * kernel_size, H, W
        :param in_channels: int
        :param norm: normalizer
        :param channel_list: List[int]
        :param b_mul: int default=1
        :param use_norm: bool default=False
        :return: mlp_output_reshape: num_batch, out_channels, kernel_size * kernel_size, H, W
        """
        assert isinstance(channel_list, list)
        # NOTE: Currently only supports two-layer Conv structure
        assert len(channel_list) == 2

        x = torch.reshape(
            data,
            shape=(self.num_batch * b_mul,
                in_channels,
                -1,
                self.W
            )
        )

        y = self.conv1(x)
        # ic(y.size())

        if use_norm:
            y = self.bn1(y)
        y = self.act1(y)
        y = self.conv2(y)
        # ic(y.size())

        mlp_output_reshape = torch.reshape(
            y,
            shape=(
                self.num_batch * b_mul,
                channel_list[-1],
                -1,
                self.H,
                self.W
            )
        )
        # ic(mlp_output_reshape.size())
        return mlp_output_reshape

    def forward(self, data, coord_data, data_channels, coord_channels, kernel_size=3):
        """
        # Without data mlp;
        # MLP: fc + norm + relu + fc;
        # Using normalized coordinates
        :param data: num_batch, num_channel_in, H, W
        :param coord_data: num_batch, 3, H, W
        :param data_channels: num_channel_in
        :param coord_channels: 3
        :param norm: normalizer
        :param conv1_filter: int
        :param kernel_size: int default=3
        :param kwargs:
        :return: conv1: num_batch, conv1_filter, H, W
        """

        coord_sample_data = self.sample_coord(coord_data, kernel_size)  # (1, 45, 131072)
        # ic(coord_sample_data.size())
        # ic(coord_data.size())

        rel_coord = self.relative_coord(    # (1, 5, 9, 64, 2048)
            coord_sample_data,
            coord_data,
            coord_channels,
            kernel_size)
        # ic(rel_coord.size())

        weights = self.mlp(     # (1, 32, 9, 64, 2048)
            rel_coord,
            in_channels=coord_channels,
            channel_list=self.channel_list)
        # ic(weights.size())

        data_sample = self.sample_data(data, kernel_size)   # (1, 288, 131072)
        # ic(data_sample.size())

        data_sample_reshape = torch.reshape(        # (1, 32, 9, 64, 2048)
            data_sample,
            shape=(
                self.num_batch,
                data_channels,
                kernel_size * kernel_size,
                self.H,
                self.W)
        )
        # ic(data_sample_reshape.size())

        output = data_sample_reshape * weights  # (1, 32, 9, 64, 2048)
        # ic(output.size())

        output_reshape = torch.reshape(     # (1, 288, 64, 2048)
            output,
            shape=(
                self.num_batch,
                -1,
                self.H,
                self.W)
        )
        # ic(output_reshape.size())

        output = self.conv1x1(output_reshape)   # (1, 32, 64, 2048)
        return output

if __name__ == '__main__':
    # from icecream import ic
    #
    # input = torch.randn(1, 8, 64, 2048)
    # METAConv = MetaKernel(
    #     num_batch=1,
    #     feat_height=64,
    #     feat_width=2048,
    #     coord_channels=8)
    #
    # downCntx = torch.randn(1, 32, 64, 2048)
    # output = METAConv(data=downCntx,
    #                   coord_data=input,
    #                   data_channels=downCntx.size()[1],
    #                   coord_channels=input.size()[1],
    #                   kernel_size=3)
    #
    # ic(output.size())

    input = torch.randn(1, 32, 64, 2048)
    block = ResBlock(32, 2 * 32, 0.2, pooling=True, drop_out=False, kernel_size=(2, 4))
    output1, output2 = block(input)

    print(output1.shape)
    print(output2.shape)

modules/BaseBlocks.py
SCNU-RISLAB/MF-MOS
# !/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import __init__ as booger
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.utils import *

from modules.BaseBlocks import MetaKernel, ResContextBlock, ResBlock, UpBlock

class MFMOS(nn.Module):
    def __init__(self, nclasses, movable_nclasses, params, num_batch=None, point_refine=None):
        super(MFMOS, self).__init__()
        self.nclasses = nclasses
        self.use_attention = "MGA"
        self.point_refine = point_refine

        self.range_channel = 5
        print("Channel of range image input = ", self.range_channel)
        print("Number of residual images input = ", params['train']['n_input_scans'])
        
        self.downCntx = ResContextBlock(self.range_channel, 32)
        self.downCntx2 = ResContextBlock(32, 32)
        self.downCntx3 = ResContextBlock(32, 32)
        print("params['train']['batch_size']", params['train']['batch_size'])
        self.metaConv = MetaKernel(num_batch=int(params['train']['batch_size']) if num_batch is None else num_batch,
                                   feat_height=params['dataset']['sensor']['img_prop']['height'],
                                   feat_width=params['dataset']['sensor']['img_prop']['width'],
                                   coord_channels=self.range_channel)

        self.resBlock1 = ResBlock(32, 2 * 32, 0.2, pooling=True, drop_out=False, kernel_size=(2, 4))
        self.resBlock2 = ResBlock(2 * 32, 2 * 2 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.resBlock3 = ResBlock(2 * 2 * 32, 2 * 4 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.resBlock4 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.resBlock5 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=False, kernel_size=(2, 4))

        self.upBlock1 = UpBlock(2 * 4 * 32, 4 * 32, 0.2)
        self.upBlock2 = UpBlock(4 * 32, 4 * 32, 0.2)
        self.upBlock3 = UpBlock(4 * 32, 2 * 32, 0.2)
        self.upBlock4 = UpBlock(2 * 32, 32, 0.2, drop_out=False)

        self.range_upBlock1 = UpBlock(256, 128, 0.2)
        self.range_upBlock2 = UpBlock(128, 64, 0.2)
        self.range_upBlock3 = UpBlock(64, 32, 0.2)
        # self.range_upBlock4 = UpBlock(64, 32, 0.2, drop_out=False)

        # Context Block for residual image 
        self.RI_downCntx = ResContextBlock(params['train']['n_input_scans'], 32)
        # self.RI_downCntx2 = ResContextBlock(32, 32)
        # self.RI_downCntx3 = ResContextBlock(32, 32)

        self.RI_resBlock1 = ResBlock(32, 2 * 32, 0.2, pooling=True, drop_out=False, kernel_size=(2, 4))
        self.RI_resBlock2 = ResBlock(2 * 32, 2 * 2 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.RI_resBlock3 = ResBlock(2 * 2 * 32, 2 * 4 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.RI_resBlock4 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=True, kernel_size=(2, 4))
        self.RI_resBlock5 = ResBlock(2 * 4 * 32, 4 * 4 * 32, 0.2, pooling=False, kernel_size=(2, 4))
        
        self.logits3 = nn.Conv2d(32, nclasses, kernel_size=(1, 1))
        self.movable_logits = nn.Conv2d(32, movable_nclasses, kernel_size=(1, 1))

        if self.use_attention == "MGA":
            self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))

            self.conv1x1_conv1_channel_wise = nn.Conv2d(32, 32, 1, bias=True)
            self.conv1x1_conv1_spatial = nn.Conv2d(32, 1, 1, bias=True)

            self.conv1x1_layer0_channel_wise = nn.Conv2d(64, 64, 1, bias=True)
            self.conv1x1_layer0_spatial = nn.Conv2d(64, 1, 1, bias=True)

            self.conv1x1_layer1_channel_wise  = nn.Conv2d(128, 128, 1, bias=True)
            self.conv1x1_layer1_spatial = nn.Conv2d(128, 1, 1, bias=True)

            self.conv1x1_layer2_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer2_spatial = nn.Conv2d(256, 1, 1, bias=True)

            self.conv1x1_layer3_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer3_spatial = nn.Conv2d(256, 1, 1, bias=True)

            self.conv1x1_layer4_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer4_spatial = nn.Conv2d(256, 1, 1, bias=True)
        else:
            pass # raise NotImplementedError

    def encoder_attention_module_MGA_tmc(self, img_feat, flow_feat, conv1x1_channel_wise, conv1x1_spatial):
        """
            flow_feat_map:  [bsize, 1, h, w]
            feat_vec:       [bsize, channel, 1, 1]
            channel_attentioned_img_feat:  [bsize, channel, h, w]
        """
        # spatial attention
        flow_feat_map = conv1x1_spatial(flow_feat)  
        flow_feat_map = nn.Sigmoid()(flow_feat_map)
        spatial_attentioned_img_feat = flow_feat_map * img_feat

        # channel-wise attention
        feat_vec = self.avg_pool(spatial_attentioned_img_feat)
        feat_vec = conv1x1_channel_wise(feat_vec)
        feat_vec = nn.Softmax(dim=1)(feat_vec) * feat_vec.shape[1]
        channel_attentioned_img_feat = spatial_attentioned_img_feat * feat_vec

        final_feat = channel_attentioned_img_feat + img_feat
        return final_feat

    def forward(self, x):
        """
            x: shape [bs, c, h, w],  c = range image channel + num of residual images
            *_downCntx:[bs, .., h, w]
            RI_down0c: [bs, c', h/2, w/2]       RI_down0b:  [bs, c', h, w] 
            RI_down1c: [bs, c'', h/4, w/4]      RI_down1b:  [bs, c'', h/2, w/2] 
            RI_down2c: [bs, c'', h/8, w/8]      RI_down2b:  [bs, c'', h/4, w/4] 
            RI_down3c: [bs, c'', h/16, w/16]    RI_down3b:  [bs, c'', h/8, w/8] 
            up4e: [bs, .., h/8, w/8] 
            up3e: [bs, .., h/4, w/4]
            up2e: [bs, .., h/2, w/2]
            up1e: [bs, .., h, w]
            logits: [bs, num_class, h, w]
        """

        # print("x shape is {}".format(x.shape))
        # split the input data to range image (5 channel) and residual images
        current_range_image = x[:, :self.range_channel, : ,:]
        residual_images = x[:, self.range_channel:, : ,:]

        # print("residual_images {}".format(residual_images.shape))
        ###### the Encoder for residual image ######
        RI_downCntx = self.RI_downCntx(residual_images)
        # print("RI_downCntx {}".format(RI_downCntx.shape))

        ###### the Encoder for range image ######       # range (3, 5, 64, 2048)
        downCntx = self.downCntx(current_range_image)   # (3, 32, 64, 2048)
        # Use MetaKernel to capture more spatial information
        downCntx = self.metaConv(data=downCntx,
                                 coord_data=current_range_image,
                                 data_channels=downCntx.size()[1],
                                 coord_channels=current_range_image.size()[1],
                                 kernel_size=3)
        downCntx = self.downCntx2(downCntx)
        downCntx = self.downCntx3(downCntx)

        Range_down0c, Range_down0b = self.RI_resBlock1(downCntx)        # (64, 32, 1024), (64, 64, 2048)
        Range_down1c, Range_down1b = self.RI_resBlock2(Range_down0c)    # (128, 16, 512), (128, 32, 1024)
        Range_down2c, Range_down2b = self.RI_resBlock3(Range_down1c)    # (256, 8, 256) , (256, 16, 512)
        Range_down3c, Range_down3b = self.RI_resBlock4(Range_down2c)    # (256, 4, 128) , (256, 8, 256)
        # Range_down4c = self.RI_resBlock5(Range_down3c)                  # (512, 4, 128)

        ###### Bridging two specific branches using MotionGuidedAttention ######
        if self.use_attention == "MGA":
            downCntx = self.encoder_attention_module_MGA_tmc(RI_downCntx, downCntx, self.conv1x1_conv1_channel_wise, self.conv1x1_conv1_spatial)
        elif self.use_attention == "Add":
            downCntx += RI_downCntx
        down0c, down0b = self.resBlock1(downCntx)

        if self.use_attention == "MGA":
            down0c = self.encoder_attention_module_MGA_tmc(down0c, Range_down0c, self.conv1x1_layer0_channel_wise, self.conv1x1_layer0_spatial)
        elif self.use_attention == "Add":
            down0c += Range_down0c
        down1c, down1b = self.resBlock2(down0c)

        if self.use_attention == "MGA":
            down1c = self.encoder_attention_module_MGA_tmc(down1c, Range_down1c, self.conv1x1_layer1_channel_wise, self.conv1x1_layer1_spatial)
        elif self.use_attention == "Add":
            down1c += Range_down1c
        down2c, down2b = self.resBlock3(down1c)

        if self.use_attention == "MGA":
            down2c = self.encoder_attention_module_MGA_tmc(down2c, Range_down2c, self.conv1x1_layer2_channel_wise, self.conv1x1_layer2_spatial)
        elif self.use_attention == "Add":
            down2c += Range_down2c
        down3c, down3b = self.resBlock4(down2c) # (1, 256, 4, 128) (1, 256, 8, 256)

        if self.use_attention == "MGA":
            down3c = self.encoder_attention_module_MGA_tmc(down3c, Range_down3c, self.conv1x1_layer3_channel_wise, self.conv1x1_layer3_spatial)
        elif self.use_attention == "Add":
            down3c += Range_down3c
        down5c = self.resBlock5(down3c)     # (1, 256, 4, 128)

        ###### the Decoder, same as SalsaNext ######
        up4e = self.upBlock1(down5c, down3b)
        up3e = self.upBlock2(up4e, down2b)
        up2e = self.upBlock3(up3e, down1b)
        up1e = self.upBlock4(up2e, down0b)

        logits = self.logits3(up1e)
        logits = F.softmax(logits, dim=1)

        range_up4e = self.range_upBlock1(Range_down3b, Range_down2b)
        range_up3e = self.range_upBlock2(range_up4e, Range_down1b)
        range_up2e = self.range_upBlock3(range_up3e, Range_down0b)
        # range_up1e = self.range_upBlock1(range_up2e, Range_down0b)

        movable_logits = self.movable_logits(range_up2e)
        movable_logits = F.softmax(movable_logits, dim=1)

        return logits, up1e, movable_logits, range_up2e

if __name__ == '__main__':
    from torchinfo import summary

    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    input = torch.randn(1, 13, 64, 2048)
    arch_cfg = load_yaml('../train_yaml/mos_coarse_stage.yml')
    # print(arch_cfg)
    model = MFMOS(nclasses=3, movable_nclasses=3, params=arch_cfg, num_batch=1)
    # print(model)
    output, movable_output = model(input)    # (1, 3, 64, 2048)

    summary(model, (1, 13, 64, 2048))

    print(output.shape)
    print(movable_output.shape)
modules/MFMOS.py
SCNU-RISLAB/MF-MOS
# A simple MLP network structure for point clouds,
# 
# Added by Jiadai Sun

import torch
import torch.nn as nn
import torch.nn.functional as F


class PointRefine(nn.Module):

    def __init__(self, n_class=3,
                 in_fea_dim=35,
                 out_point_fea_dim=64):
        super(PointRefine, self).__init__()

        self.n_class = n_class
        self.PPmodel = nn.Sequential(
            nn.BatchNorm1d(in_fea_dim),

            nn.Linear(in_fea_dim, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),

            nn.Linear(64, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),

            nn.Linear(128, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),

            nn.Linear(256, out_point_fea_dim)
        )

        self.logits = nn.Sequential(
            nn.Linear(out_point_fea_dim, self.n_class)
        )

    def forward(self, point_fea):
        # the point_fea need with size (b, N, c) e.g.  torch.Size([1, 121722, 35])
        # process feature
        # torch.Size([124668, 9]) --> torch.Size([124668, 256])
        processed_point_fea = self.PPmodel(point_fea)
        logits = self.logits(processed_point_fea)
        point_predict = F.softmax(logits, dim=1)
        return point_predict


if __name__ == '__main__':

    import time
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = PointRefine()
    model.train()

    # t0 = time.time()
    # pred = model(cloud)
    # t1 = time.time()
    # print(t1-t0)

    total = sum([param.nelement() for param in model.parameters()])
    print("Number of PointRefine parameter: %.2fM" % (total/1e6))
    # Number of PointRefine parameter: 0.04M

modules/PointRefine/PointMLP.py
SCNU-RISLAB/MF-MOS
import torch
import torchsparse
import torchsparse.nn as spnn
import torchsparse.nn.functional as F

from torch import nn
from torchsparse import PointTensor, SparseTensor
from torchsparse.nn.utils import get_kernel_offsets


# z: PointTensor
# return: SparseTensor
def initial_voxelize(z, init_res, after_res):
    new_float_coord = torch.cat(
        [(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)

    pc_hash = F.sphash(torch.floor(new_float_coord).int())
    sparse_hash = torch.unique(pc_hash)
    idx_query = F.sphashquery(pc_hash, sparse_hash)
    counts = F.spcount(idx_query.int(), len(sparse_hash))

    inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query,
                                   counts)
    inserted_coords = torch.round(inserted_coords).int()
    inserted_feat = F.spvoxelize(z.F, idx_query, counts)

    new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
    new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
    z.additional_features['idx_query'][1] = idx_query
    z.additional_features['counts'][1] = counts
    z.C = new_float_coord

    return new_tensor


# x: SparseTensor, z: PointTensor
# return: SparseTensor
def point_to_voxel(x, z):
    if z.additional_features is None or z.additional_features.get(
            'idx_query') is None or z.additional_features['idx_query'].get(
                x.s) is None:
        pc_hash = F.sphash(
            torch.cat([
                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
                z.C[:, -1].int().view(-1, 1)
            ], 1))
        sparse_hash = F.sphash(x.C)
        idx_query = F.sphashquery(pc_hash, sparse_hash)
        counts = F.spcount(idx_query.int(), x.C.shape[0])
        z.additional_features['idx_query'][x.s] = idx_query
        z.additional_features['counts'][x.s] = counts
    else:
        idx_query = z.additional_features['idx_query'][x.s]
        counts = z.additional_features['counts'][x.s]

    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
    new_tensor = SparseTensor(inserted_feat, x.C, x.s)
    new_tensor.cmaps = x.cmaps
    new_tensor.kmaps = x.kmaps

    return new_tensor


# x: SparseTensor, z: PointTensor
# return: PointTensor
def voxel_to_point(x, z, nearest=False):
    if z.idx_query is None or z.weights is None or z.idx_query.get(
            x.s) is None or z.weights.get(x.s) is None:
        off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
        old_hash = F.sphash(
            torch.cat([
                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
                z.C[:, -1].int().view(-1, 1)
            ], 1), off)
        pc_hash = F.sphash(x.C.to(z.F.device))
        idx_query = F.sphashquery(old_hash, pc_hash)
        weights = F.calc_ti_weights(z.C, idx_query,
                                    scale=x.s[0]).transpose(0, 1).contiguous()
        idx_query = idx_query.transpose(0, 1).contiguous()
        if nearest:
            weights[:, 1:] = 0.
            idx_query[:, 1:] = -1
        new_feat = F.spdevoxelize(x.F, idx_query, weights)
        new_tensor = PointTensor(new_feat,
                                 z.C,
                                 idx_query=z.idx_query,
                                 weights=z.weights)
        new_tensor.additional_features = z.additional_features
        new_tensor.idx_query[x.s] = idx_query
        new_tensor.weights[x.s] = weights
        z.idx_query[x.s] = idx_query
        z.weights[x.s] = weights

    else:
        new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s), z.weights.get(x.s))
        new_tensor = PointTensor(new_feat,
                                 z.C,
                                 idx_query=z.idx_query,
                                 weights=z.weights)
        new_tensor.additional_features = z.additional_features

    return new_tensor


class BasicConvolutionBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        dilation=dilation,
                        stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
        )

    def forward(self, x):
        out = self.net(x)
        return out


class BasicDeconvolutionBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        stride=stride,
                        transposed=True),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
        )

    def forward(self, x):
        return self.net(x)

class ResidualBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        dilation=dilation,
                        stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
            spnn.Conv3d(outc, outc, kernel_size=ks, dilation=dilation,
                        stride=1),
            spnn.BatchNorm(outc),
        )

        if inc == outc and stride == 1:
            self.downsample = nn.Identity()
        else:
            self.downsample = nn.Sequential(
                spnn.Conv3d(inc, outc, kernel_size=1, dilation=1,
                            stride=stride),
                spnn.BatchNorm(outc),
            )

        self.relu = spnn.ReLU(True)

    def forward(self, x):
        out = self.relu(self.net(x) + self.downsample(x))
        return out

class FeatureExtractBlock(nn.Module):
    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
        super().__init__()

        self.conv113 = nn.Sequential(
            spnn.Conv3d(inc, outc, kernel_size=(1, 1, ks), dilation=dilation, stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True)
        )
        self.conv131 = nn.Sequential(
            spnn.Conv3d(inc, outc, kernel_size=(1, ks, 1), dilation=dilation, stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True)
        )
        self.conv311 = nn.Sequential(
            spnn.Conv3d(inc, outc, kernel_size=(ks, 1, 1), dilation=dilation, stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True)
        )
        self.conv_cat = nn.Sequential(
            spnn.Conv3d(outc * 3, outc, kernel_size=1, dilation=dilation, stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
            spnn.Conv3d(outc, outc, kernel_size=3, dilation=dilation, stride=stride),
            spnn.BatchNorm(outc),
        )

        if inc == outc and stride == 1:
            self.downsample = nn.Identity()
        else:
            self.downsample = nn.Sequential(
                spnn.Conv3d(inc, outc, kernel_size=1, dilation=1,
                            stride=stride),
                spnn.BatchNorm(outc),
            )

        self.relu = spnn.ReLU(True)

    def forward(self, x):
        out = self.relu(self.conv_cat(torchsparse.cat([self.conv113(x), self.conv131(x), self.conv311(x)])) + self.downsample(x))
        return out


class SPVCNN(nn.Module):

    def __init__(self, **kwargs):
        super().__init__()

        cr = kwargs.get('cr', 1.0)
        cs = [32, 32, 64, 128, 256, 256, 128, 96, 96]
        cs = [int(cr * x) for x in cs]

        if 'pres' in kwargs and 'vres' in kwargs:
            self.pres = kwargs['pres']
            self.vres = kwargs['vres']

        self.stem = nn.Sequential(
            # spnn.Conv3d(4, cs[0], kernel_size=3, stride=1),
            # spnn.BatchNorm(cs[0]), spnn.ReLU(True),
            spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
            spnn.BatchNorm(cs[0]),
            spnn.ReLU(True))

        self.stage1 = nn.Sequential(
            BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
            FeatureExtractBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
        )

        self.stage2 = nn.Sequential(
            BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
            FeatureExtractBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
        )

        self.stage3 = nn.Sequential(
            BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
            FeatureExtractBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
        )

        self.stage4 = nn.Sequential(
            BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
            FeatureExtractBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
        )

        self.up1 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
            nn.Sequential(
                FeatureExtractBlock(cs[5] + cs[3], cs[5], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
            )
        ])

        self.up2 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
            nn.Sequential(
                FeatureExtractBlock(cs[6] + cs[2], cs[6], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
            )
        ])

        self.up3 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
            nn.Sequential(
                FeatureExtractBlock(cs[7] + cs[1], cs[7], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
            )
        ])

        self.up4 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
            nn.Sequential(
                FeatureExtractBlock(cs[8] + cs[0], cs[8], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
            )
        ])

        self.classifier = nn.Sequential(nn.Linear(cs[8], kwargs['num_classes']))

        self.point_transforms = nn.ModuleList([
            nn.Sequential(
                nn.Linear(cs[0], cs[4]),
                nn.BatchNorm1d(cs[4]),
                nn.ReLU(True),
            ),
            nn.Sequential(
                nn.Linear(cs[4], cs[6]),
                nn.BatchNorm1d(cs[6]),
                nn.ReLU(True),
            ),
            nn.Sequential(
                nn.Linear(cs[6], cs[8]),
                nn.BatchNorm1d(cs[8]),
                nn.ReLU(True),
            )
        ])

        self.weight_initialization()
        self.dropout = nn.Dropout(0.3, True)

    def weight_initialization(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        # x: SparseTensor z: PointTensor
        z = PointTensor(x.F, x.C.float())

        x0 = initial_voxelize(z, self.pres, self.vres)

        x0 = self.stem(x0)
        z0 = voxel_to_point(x0, z, nearest=False)
        z0.F = z0.F

        x1 = point_to_voxel(x0, z0)
        x1 = self.stage1(x1)    # x1 (95150, 32)
        x2 = self.stage2(x1)    # x2 (69207, 64)
        x3 = self.stage3(x2)    # x3 (16218, 128)
        x4 = self.stage4(x3)    # x4 (2197, 256)
        z1 = voxel_to_point(x4, z0)     # z1 (99349, 256)
        z1.F = z1.F + self.point_transforms[0](z0.F)

        y1 = point_to_voxel(x4, z1)     # y1 (2197, 256)
        y1.F = self.dropout(y1.F)
        y1 = self.up1[0](y1)            # y1 (2197, 256)  ->  (16205, 256)
        y1 = torchsparse.cat([y1, x3])
        y1 = self.up1[1](y1)

        y2 = self.up2[0](y1)        # y1 (16205, 256)  ->  (69181, 128)
        y2 = torchsparse.cat([y2, x2])
        y2 = self.up2[1](y2)
        z2 = voxel_to_point(y2, z1) # y2 (19181, 128)  ->  z2 (99348, 128)
        z2.F = z2.F + self.point_transforms[1](z1.F)

        y3 = point_to_voxel(y2, z2)
        y3.F = self.dropout(y3.F)
        y3 = self.up3[0](y3)
        y3 = torchsparse.cat([y3, x1])
        y3 = self.up3[1](y3)

        y4 = self.up4[0](y3)
        y4 = torchsparse.cat([y4, x0])
        y4 = self.up4[1](y4)
        z3 = voxel_to_point(y4, z2)
        z3.F = z3.F + self.point_transforms[2](z2.F)

        out = self.classifier(z3.F)

        scores = nn.functional.softmax(out, dim=1)
        return scores


if __name__ == "__main__":
    import numpy as np
    from torchsparse.utils.quantize import sparse_quantize
    from torchsparse.utils.collate import sparse_collate
    from torchsparse import SparseTensor

    N = 99972
    C = 32
    points_xyz = (np.random.random([N, 3]) - 0.5) * 10
    points_features = torch.tensor((np.random.random([C, N]) - 0.5) * 100, dtype=torch.float, device='cuda')

    coords = np.round(points_xyz / 0.05)
    coords -= coords.min(0, keepdims=1)

    sparse_coords, indices, inverse = sparse_quantize(coords, return_index=True, return_inverse=True)
    sparse_coords = torch.tensor(sparse_coords, dtype=torch.int, device='cuda')

    feats = points_features.permute(1, 0)[indices]

    print(points_features.permute(1, 0)[indices[0]])

    inputs = SparseTensor(coords=sparse_coords, feats=feats)
    sparse_inputs = sparse_collate([inputs]).cuda()

    # module = nn.Sequential(spnn.Conv3d(32, 32, kernel_size=(3, 1, 3), stride=1),
    #                        spnn.BatchNorm(32),
    #                        spnn.ReLU(True)).cuda()
    # output = module(sparse_inputs)

    PointHead = SPVCNN(num_classes=3, cr=1.0, pres=0.05, vres=0.05).cuda()
    output = PointHead(sparse_inputs)

    print()
modules/PointRefine/spvcnn.py
SCNU-RISLAB/MF-MOS
import torch
import torchsparse
import torchsparse.nn as spnn
import torchsparse.nn.functional as F

from torch import nn
from torchsparse import PointTensor, SparseTensor
from torchsparse.nn.utils import get_kernel_offsets


# z: PointTensor
# return: SparseTensor
def initial_voxelize(z, init_res, after_res):
    new_float_coord = torch.cat(
        [(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)

    pc_hash = F.sphash(torch.floor(new_float_coord).int())
    sparse_hash = torch.unique(pc_hash)
    idx_query = F.sphashquery(pc_hash, sparse_hash)
    counts = F.spcount(idx_query.int(), len(sparse_hash))

    inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query,
                                   counts)
    inserted_coords = torch.round(inserted_coords).int()
    inserted_feat = F.spvoxelize(z.F, idx_query, counts)

    new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
    new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
    z.additional_features['idx_query'][1] = idx_query
    z.additional_features['counts'][1] = counts
    z.C = new_float_coord

    return new_tensor


# x: SparseTensor, z: PointTensor
# return: SparseTensor
def point_to_voxel(x, z):
    if z.additional_features is None or z.additional_features.get(
            'idx_query') is None or z.additional_features['idx_query'].get(
                x.s) is None:
        pc_hash = F.sphash(
            torch.cat([
                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
                z.C[:, -1].int().view(-1, 1)
            ], 1))
        sparse_hash = F.sphash(x.C)
        idx_query = F.sphashquery(pc_hash, sparse_hash)
        counts = F.spcount(idx_query.int(), x.C.shape[0])
        z.additional_features['idx_query'][x.s] = idx_query
        z.additional_features['counts'][x.s] = counts
    else:
        idx_query = z.additional_features['idx_query'][x.s]
        counts = z.additional_features['counts'][x.s]

    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
    new_tensor = SparseTensor(inserted_feat, x.C, x.s)
    new_tensor.cmaps = x.cmaps
    new_tensor.kmaps = x.kmaps

    return new_tensor


# x: SparseTensor, z: PointTensor
# return: PointTensor
def voxel_to_point(x, z, nearest=False):
    if z.idx_query is None or z.weights is None or z.idx_query.get(
            x.s) is None or z.weights.get(x.s) is None:
        off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
        old_hash = F.sphash(
            torch.cat([
                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
                z.C[:, -1].int().view(-1, 1)
            ], 1), off)
        pc_hash = F.sphash(x.C.to(z.F.device))
        idx_query = F.sphashquery(old_hash, pc_hash)
        weights = F.calc_ti_weights(z.C, idx_query,
                                    scale=x.s[0]).transpose(0, 1).contiguous()
        idx_query = idx_query.transpose(0, 1).contiguous()
        if nearest:
            weights[:, 1:] = 0.
            idx_query[:, 1:] = -1
        new_feat = F.spdevoxelize(x.F, idx_query, weights)
        new_tensor = PointTensor(new_feat,
                                 z.C,
                                 idx_query=z.idx_query,
                                 weights=z.weights)
        new_tensor.additional_features = z.additional_features
        new_tensor.idx_query[x.s] = idx_query
        new_tensor.weights[x.s] = weights
        z.idx_query[x.s] = idx_query
        z.weights[x.s] = weights

    else:
        new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s), z.weights.get(x.s))
        new_tensor = PointTensor(new_feat,
                                 z.C,
                                 idx_query=z.idx_query,
                                 weights=z.weights)
        new_tensor.additional_features = z.additional_features

    return new_tensor


class BasicConvolutionBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        dilation=dilation,
                        stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
        )

    def forward(self, x):
        out = self.net(x)
        return out


class BasicDeconvolutionBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        stride=stride,
                        transposed=True),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
        )

    def forward(self, x):
        return self.net(x)


class ResidualBlock(nn.Module):

    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
        super().__init__()
        self.net = nn.Sequential(
            spnn.Conv3d(inc,
                        outc,
                        kernel_size=ks,
                        dilation=dilation,
                        stride=stride),
            spnn.BatchNorm(outc),
            spnn.ReLU(True),
            spnn.Conv3d(outc, outc, kernel_size=ks, dilation=dilation,
                        stride=1),
            spnn.BatchNorm(outc),
        )

        if inc == outc and stride == 1:
            self.downsample = nn.Identity()
        else:
            self.downsample = nn.Sequential(
                spnn.Conv3d(inc, outc, kernel_size=1, dilation=1,
                            stride=stride),
                spnn.BatchNorm(outc),
            )

        self.relu = spnn.ReLU(True)

    def forward(self, x):
        out = self.relu(self.net(x) + self.downsample(x))
        return out


class SPVCNN(nn.Module):

    def __init__(self, **kwargs):
        super().__init__()

        # cr = kwargs.get('cr', 1.0)
        # cs = [32, 32, 64, 128, 256, 256, 128, 96, 96]
        cs = [32, 64, 128, 96, 96]
        # cs = [int(cr * x) for x in cs]

        if 'pres' in kwargs and 'vres' in kwargs:
            self.pres = kwargs['pres']
            self.vres = kwargs['vres']

        self.stem = nn.Sequential(
            # spnn.Conv3d(4, cs[0], kernel_size=3, stride=1),
            # spnn.BatchNorm(cs[0]), spnn.ReLU(True),
            spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
            spnn.BatchNorm(cs[0]),
            spnn.ReLU(True))

        self.stage1 = nn.Sequential(
            BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
            ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
        )

        self.stage2 = nn.Sequential(
            BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
            ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
            ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
        )

        self.up1 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[2], cs[3], ks=2, stride=2),
            nn.Sequential(
                ResidualBlock(cs[3] + cs[1], cs[3], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
            )
        ])

        self.up2 = nn.ModuleList([
            BasicDeconvolutionBlock(cs[3], cs[4], ks=2, stride=2),
            nn.Sequential(
                ResidualBlock(cs[4] + cs[0], cs[4], ks=3, stride=1, dilation=1),
                ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
            )
        ])

        self.classifier = nn.Sequential(nn.Linear(cs[4], kwargs['num_classes']))

        self.point_transforms = nn.ModuleList([
            nn.Sequential(
                nn.Linear(cs[0], cs[2]),
                nn.BatchNorm1d(cs[2]),
                nn.ReLU(True),
            ),
            nn.Sequential(
                nn.Linear(cs[2], cs[4]),
                nn.BatchNorm1d(cs[4]),
                nn.ReLU(True),
            )
        ])

        self.weight_initialization()
        self.dropout = nn.Dropout(0.3, True)

    def weight_initialization(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


    def forward(self, x):
        # x: SparseTensor z: PointTensor
        z = PointTensor(x.F, x.C.float())

        x0 = initial_voxelize(z, self.pres, self.vres)

        x0 = self.stem(x0)
        z0 = voxel_to_point(x0, z, nearest=False)

        x1 = self.stage1(x0)
        x2 = self.stage2(x1)
        z1 = voxel_to_point(x2, z0)
        z1.F = z1.F + self.point_transforms[0](z0.F)

        y3 = point_to_voxel(x2, z1)
        # y3.F = self.dropout(y3.F)
        y3 = self.up1[0](y3)
        y3 = torchsparse.cat([y3, x1])
        y3 = self.up1[1](y3)

        y4 = self.up2[0](y3)
        y4 = torchsparse.cat([y4, x0])
        y4 = self.up2[1](y4)
        z3 = voxel_to_point(y4, z1)
        z3.F = z3.F + self.point_transforms[1](z1.F)

        out = self.classifier(z3.F)

        scores = nn.functional.softmax(out, dim=1)
        return scores
modules/PointRefine/spvcnn_lite.py
SCNU-RISLAB/MF-MOS
# !/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import __init__ as booger
import torch
import torch.nn as nn
import torch.nn.functional as F
# from icecream import ic
from utils.utils import *

from modules.BaseBlocks import MetaKernel, ResContextBlock, ResBlock, UpBlock


class SalsaNextWithMotionAttention(nn.Module):
    def __init__(self, nclasses, params, num_batch=None, point_refine=None):
        super(SalsaNextWithMotionAttention, self).__init__()
        self.nclasses = nclasses
        self.use_attention = "MGA"
        self.point_refine = point_refine

        self.range_channel = 5
        print("Channel of range image input = ", self.range_channel)
        print("Number of residual images input = ", params['train']['n_input_scans'])
        
        self.downCntx = ResContextBlock(self.range_channel, 32)
        self.downCntx2 = ResContextBlock(32, 32)
        self.downCntx3 = ResContextBlock(32, 32)
        print("params['train']['batch_size']", params['train']['batch_size'])
        self.metaConv = MetaKernel(num_batch=int(params['train']['batch_size'] / torch.cuda.device_count()) if num_batch is None else num_batch,
                                   feat_height=params['dataset']['sensor']['img_prop']['height'],
                                   feat_width=params['dataset']['sensor']['img_prop']['width'],
                                   coord_channels=self.range_channel)

        self.resBlock1 = ResBlock(32, 2 * 32, 0.2, pooling=True, drop_out=False)
        self.resBlock2 = ResBlock(2 * 32, 2 * 2 * 32, 0.2, pooling=True)
        self.resBlock3 = ResBlock(2 * 2 * 32, 2 * 4 * 32, 0.2, pooling=True)
        self.resBlock4 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=True)
        self.resBlock5 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=False)

        self.upBlock1 = UpBlock(2 * 4 * 32, 4 * 32, 0.2)
        self.upBlock2 = UpBlock(4 * 32, 4 * 32, 0.2)
        self.upBlock3 = UpBlock(4 * 32, 2 * 32, 0.2)
        self.upBlock4 = UpBlock(2 * 32, 32, 0.2, drop_out=False)

        # Context Block for residual image 
        self.RI_downCntx = ResContextBlock(params['train']['n_input_scans'], 32)
        # self.RI_downCntx2 = ResContextBlock(32, 32)
        # self.RI_downCntx3 = ResContextBlock(32, 32)

        self.RI_resBlock1 = ResBlock(32, 2 * 32, 0.2, pooling=True, drop_out=False)
        self.RI_resBlock2 = ResBlock(2 * 32, 2 * 2 * 32, 0.2, pooling=True)
        self.RI_resBlock3 = ResBlock(2 * 2 * 32, 2 * 4 * 32, 0.2, pooling=True)
        self.RI_resBlock4 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=True)
        self.RI_resBlock5 = ResBlock(2 * 4 * 32, 2 * 4 * 32, 0.2, pooling=False)

        self.RI_upBlock1 = UpBlock(2 * 4 * 32, 4 * 32, 0.2)
        self.RI_upBlock2 = UpBlock(4 * 32, 4 * 32, 0.2)
        self.RI_upBlock3 = UpBlock(4 * 32, 2 * 32, 0.2)
        self.RI_upBlock4 = UpBlock(2 * 32, 32, 0.2, drop_out=False)
        
        self.logits3 = nn.Conv2d(32, nclasses, kernel_size=(1, 1))

        if self.use_attention == "MGA":
            self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))

            self.conv1x1_conv1_channel_wise = nn.Conv2d(32, 32, 1, bias=True)
            self.conv1x1_conv1_spatial = nn.Conv2d(32, 1, 1, bias=True)

            self.conv1x1_layer0_channel_wise = nn.Conv2d(64, 64, 1, bias=True)
            self.conv1x1_layer0_spatial = nn.Conv2d(64, 1, 1, bias=True)

            self.conv1x1_layer1_channel_wise  = nn.Conv2d(128, 128, 1, bias=True)
            self.conv1x1_layer1_spatial = nn.Conv2d(128, 1, 1, bias=True)

            self.conv1x1_layer2_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer2_spatial = nn.Conv2d(256, 1, 1, bias=True)

            self.conv1x1_layer3_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer3_spatial = nn.Conv2d(256, 1, 1, bias=True)

            self.conv1x1_layer4_channel_wise = nn.Conv2d(256, 256, 1, bias=True)
            self.conv1x1_layer4_spatial = nn.Conv2d(256, 1, 1, bias=True)
        else:
            pass # raise NotImplementedError

    def encoder_attention_module_MGA_tmc(self, img_feat, flow_feat, conv1x1_channel_wise, conv1x1_spatial):
        """
            flow_feat_map:  [bsize, 1, h, w]
            feat_vec:       [bsize, channel, 1, 1]
            channel_attentioned_img_feat:  [bsize, channel, h, w]
        """
        # spatial attention
        flow_feat_map = conv1x1_spatial(flow_feat)  
        flow_feat_map = nn.Sigmoid()(flow_feat_map)
        spatial_attentioned_img_feat = flow_feat_map * img_feat

        # channel-wise attention
        feat_vec = self.avg_pool(spatial_attentioned_img_feat)
        feat_vec = conv1x1_channel_wise(feat_vec)
        feat_vec = nn.Softmax(dim=1)(feat_vec) * feat_vec.shape[1]
        channel_attentioned_img_feat = spatial_attentioned_img_feat * feat_vec

        final_feat = channel_attentioned_img_feat + img_feat
        return final_feat

    def forward(self, x):
        """
            x: shape [bs, c, h, w],  c = range image channel + num of residual images
            *_downCntx:[bs, .., h, w]
            RI_down0c: [bs, c', h/2, w/2]       RI_down0b:  [bs, c', h, w] 
            RI_down1c: [bs, c'', h/4, w/4]      RI_down1b:  [bs, c'', h/2, w/2] 
            RI_down2c: [bs, c'', h/8, w/8]      RI_down2b:  [bs, c'', h/4, w/4] 
            RI_down3c: [bs, c'', h/16, w/16]    RI_down3b:  [bs, c'', h/8, w/8] 
            up4e: [bs, .., h/8, w/8] 
            up3e: [bs, .., h/4, w/4]
            up2e: [bs, .., h/2, w/2]
            up1e: [bs, .., h, w]
            logits: [bs, num_class, h, w]
        """

        # print("x shape is {}".format(x.shape))
        # split the input data to range image (5 channel) and residual images
        current_range_image = x[:, :self.range_channel, : ,:]
        residual_images = x[:, self.range_channel:, : ,:]

        # print("residual_images {}".format(residual_images.shape))
        ###### the Encoder for residual image ######
        RI_downCntx = self.RI_downCntx(residual_images)
        # print("RI_downCntx {}".format(RI_downCntx.shape))

        RI_down0c, RI_down0b = self.RI_resBlock1(RI_downCntx)
        RI_down1c, RI_down1b = self.RI_resBlock2(RI_down0c)
        RI_down2c, RI_down2b = self.RI_resBlock3(RI_down1c)
        RI_down3c, RI_down3b = self.RI_resBlock4(RI_down2c)
        # RI_down5c = self.RI_resBlock5(RI_down3c)

        ###### the Encoder for range image ######       # range (3, 5, 64, 2048)
        downCntx = self.downCntx(current_range_image)   # (3, 32, 64, 2048)
        # print("downCntx1 {}".format(downCntx.shape))
        # Use MetaKernel to capture more spatial information
        downCntx = self.metaConv(data=downCntx,
                                 coord_data=current_range_image,
                                 data_channels=downCntx.size()[1],
                                 coord_channels=current_range_image.size()[1],
                                 kernel_size=3)
        # print("downCntx2 {}".format(downCntx.shape))
        downCntx = self.downCntx2(downCntx)
        # print("downCntx3 {}".format(downCntx.shape))
        downCntx = self.downCntx3(downCntx)
        # print("downCntx4 {}".format(downCntx.shape))

        ###### Bridging two specific branches using MotionGuidedAttention ######
        if self.use_attention == "MGA":
            downCntx = self.encoder_attention_module_MGA_tmc(downCntx, RI_downCntx, self.conv1x1_conv1_channel_wise, self.conv1x1_conv1_spatial)
        elif self.use_attention == "Add":
            downCntx += RI_downCntx
        down0c, down0b = self.resBlock1(downCntx)

        if self.use_attention == "MGA":
            down0c = self.encoder_attention_module_MGA_tmc(down0c, RI_down0c, self.conv1x1_layer0_channel_wise, self.conv1x1_layer0_spatial)
        elif self.use_attention == "Add":
            down0c += RI_down0c
        down1c, down1b = self.resBlock2(down0c)

        if self.use_attention == "MGA":
            down1c = self.encoder_attention_module_MGA_tmc(down1c, RI_down1c, self.conv1x1_layer1_channel_wise, self.conv1x1_layer1_spatial)
        elif self.use_attention == "Add":
            down1c += RI_down1c
        down2c, down2b = self.resBlock3(down1c)

        if self.use_attention == "MGA":
            down2c = self.encoder_attention_module_MGA_tmc(down2c, RI_down2c, self.conv1x1_layer2_channel_wise, self.conv1x1_layer2_spatial)
        elif self.use_attention == "Add":
            down2c += RI_down2c
        down3c, down3b = self.resBlock4(down2c) # (1, 256, 4, 128) (1, 256, 8, 256)

        if self.use_attention == "MGA":
            down3c = self.encoder_attention_module_MGA_tmc(down3c, RI_down3c, self.conv1x1_layer3_channel_wise, self.conv1x1_layer3_spatial)
        elif self.use_attention == "Add":
            down3c += RI_down3c
        down5c = self.resBlock5(down3c)     # (1, 256, 4, 128)

        ###### the Decoder, same as SalsaNext ######
        up4e = self.upBlock1(down5c, down3b)
        up3e = self.upBlock2(up4e, down2b)
        up2e = self.upBlock3(up3e, down1b)
        up1e = self.upBlock4(up2e, down0b)

        logits = self.logits3(up1e)
        logits = F.softmax(logits, dim=1)

        return logits, up1e

if __name__ == '__main__':
    from torchinfo import summary
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    input = torch.randn(1, 13, 64, 2048)
    arch_cfg = load_yaml('../train_yaml/mos_coarse_stage.yml')
    # print(arch_cfg)
    model = SalsaNextWithMotionAttention(nclasses=3, params=arch_cfg, num_batch=1)
    # print(model)
    output, _ = model(input)    # (1, 3, 64, 2048)

    summary(model, (1, 13, 64, 2048))
modules/SalsaNextWithMotionAttention.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import numpy
import torch
import torch.nn as nn
import torch.nn.functional as F

#PyTorch
# class DiceLoss(nn.Module):
#     def __init__(self, weight=None, size_average=True):
#         super(DiceLoss, self).__init__()

#     def forward(self, inputs, targets, smooth=1):
        
#         #comment out if your model contains a sigmoid or equivalent activation layer
#         inputs = F.sigmoid(inputs)       
        
#         #flatten label and prediction tensors
#         inputs = inputs.view(-1)
#         targets = targets.view(-1)
        
#         intersection = (inputs * targets).sum()                            
#         dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  
        
#         return 1 - dice

# https://smp.readthedocs.io/en/latest/losses.html
# https://github.com/pytorch/pytorch/issues/1249
# https://www.kaggle.com/bigironsphere/loss-function-library-keras-pytorch#Dice-Loss
# https://kornia.readthedocs.io/en/v0.1.2/_modules/torchgeometry/losses/dice.html


# based on:
# https://github.com/kevinzakka/pytorch-goodies/blob/master/losses.py

class DiceLoss(nn.Module):
    r"""Criterion that computes Sørensen-Dice Coefficient loss.

    According to [1], we compute the Sørensen-Dice Coefficient as follows:

    .. math::

        \text{Dice}(x, class) = \frac{2 |X| \cap |Y|}{|X| + |Y|}

    where:
       - :math:`X` expects to be the scores of each class.
       - :math:`Y` expects to be the one-hot tensor with the class labels.

    the loss, is finally computed as:

    .. math::

        \text{loss}(x, class) = 1 - \text{Dice}(x, class)

    [1] https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient

    Shape:
        - Input: :math:`(N, C, H, W)` where C = number of classes.
        - Target: :math:`(N, H, W)` where each value is
          :math:`0 ≤ targets[i] ≤ C−1`.

    Examples:
        >>> N = 5  # num_classes
        >>> loss = tgm.losses.DiceLoss()
        >>> input = torch.randn(1, N, 3, 5, requires_grad=True)
        >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
        >>> output = loss(input, target)
        >>> output.backward()
    """

    def __init__(self) -> None:
        super(DiceLoss, self).__init__()
        self.eps: float = 1e-6

    def forward(self, input: torch.Tensor,
                target: torch.Tensor) -> torch.Tensor:
        if not torch.is_tensor(input):
            raise TypeError("Input type is not a torch.Tensor. Got {}"
                            .format(type(input)))
        if not len(input.shape) == 4:
            raise ValueError("Invalid input shape, we expect BxNxHxW. Got: {}"
                             .format(input.shape))
        if not input.shape[-2:] == target.shape[-2:]:
            raise ValueError("input and target shapes must be the same. Got: {}"
                             .format(input.shape, input.shape))
        if not input.device == target.device:
            raise ValueError(
                "input and target must be in the same device. Got: {}" .format(
                    input.device, target.device))
        # compute softmax over the classes axis
        # input_soft = F.softmax(input, dim=1) # have done is network last layer

        # create the labels one hot tensor
        # target_one_hot = one_hot(target, num_classes=input.shape[1],
        #                          device=input.device, dtype=input.dtype)
        target_one_hot = F.one_hot(target, num_classes=input.shape[1]).permute(0, 3, 1, 2)

        # compute the actual dice score
        dims = (1, 2, 3)
        # intersection = torch.sum(input_soft * target_one_hot, dims)
        # cardinality = torch.sum(input_soft + target_one_hot, dims)

        ## if we need to ignore the class=0
        input_filter = input[:, 1:, :, :]
        target_one_hot_filter = input[:, 1:, :, :]
        intersection = torch.sum(input_filter * target_one_hot_filter, dims)
        cardinality = torch.sum(input_filter + target_one_hot_filter, dims)

        dice_score = 2. * intersection / (cardinality + self.eps)
        return torch.mean(1. - dice_score)



######################
# functional interface
######################


def dice_loss(input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    r"""Function that computes Sørensen-Dice Coefficient loss.

    See :class:`~torchgeometry.losses.DiceLoss` for details.
    """
    return DiceLoss()(input, target)


modules/loss/DiceLoss.py
SCNU-RISLAB/MF-MOS
"""

MIT License

Copyright (c) 2018 Maxim Berman
Copyright (c) 2020 Tiago Cortinhal, George Tzelepis and Eren Erdal Aksoy


Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

"""
import torch
import torch.nn as nn
from torch.autograd import Variable


try:
    from itertools import ifilterfalse
except ImportError:
    from itertools import filterfalse as ifilterfalse


def isnan(x):
    return x != x


def mean(l, ignore_nan=False, empty=0):
    """
    nanmean compatible with generators.
    """
    l = iter(l)
    if ignore_nan:
        l = ifilterfalse(isnan, l)
    try:
        n = 1
        acc = next(l)
    except StopIteration:
        if empty == 'raise':
            raise ValueError('Empty mean')
        return empty
    for n, v in enumerate(l, 2):
        acc += v
    if n == 1:
        return acc
    return acc / n


def lovasz_grad(gt_sorted):
    """
    Computes gradient of the Lovasz extension w.r.t sorted errors
    See Alg. 1 in paper
    """
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts - gt_sorted.float().cumsum(0)
    union = gts + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1. - intersection / union
    if p > 1:  # cover 1-pixel case
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard


def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None):
    """
    Multi-class Lovasz-Softmax loss
      probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
              Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
      labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
      per_image: compute the loss per image instead of per batch
      ignore: void class labels
    """
    if per_image:
        loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes)
                    for prob, lab in zip(probas, labels))
    else:
        loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes)
    return loss


def lovasz_softmax_flat(probas, labels, classes='present'):
    """
    Multi-class Lovasz-Softmax loss
      probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
      labels: [P] Tensor, ground truth labels (between 0 and C - 1)
      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    """
    if probas.numel() == 0:
        # only void pixels, the gradients should be 0
        return probas * 0.
    C = probas.size(1)
    losses = []
    class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes
    for c in class_to_sum:
        fg = (labels == c).float()  # foreground for class c
        if (classes == 'present' and fg.sum() == 0):
            continue
        if C == 1:
            if len(classes) > 1:
                raise ValueError('Sigmoid output possible only with 1 class')
            class_pred = probas[:, 0]
        else:
            class_pred = probas[:, c]
        errors = (Variable(fg) - class_pred).abs()
        errors_sorted, perm = torch.sort(errors, 0, descending=True)
        perm = perm.data
        fg_sorted = fg[perm]
        losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
    return mean(losses)


def flatten_probas(probas, labels, ignore=None):
    """
    Flattens predictions in the batch
    """
    if probas.dim() == 3:
        # assumes output of a sigmoid layer
        B, H, W = probas.size()
        probas = probas.view(B, 1, H, W)
    B, C, H, W = probas.size()
    probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
    labels = labels.view(-1)
    if ignore is None:
        return probas, labels
    valid = (labels != ignore)
    vprobas = probas[valid.nonzero(as_tuple=False).squeeze()]
    vlabels = labels[valid]
    return vprobas, vlabels


class Lovasz_softmax(nn.Module):
    def __init__(self, classes='present', per_image=False, ignore=None):
        super(Lovasz_softmax, self).__init__()
        self.classes = classes
        self.per_image = per_image
        self.ignore = ignore

    def forward(self, probas, labels):
        return lovasz_softmax(probas, labels, self.classes, self.per_image, self.ignore)


# Used to calculate Lovasz Loss with point cloud as input
# Add by Jiadai Sun
class Lovasz_softmax_PointCloud(nn.Module):
    def __init__(self, classes='present', ignore=None):
        super(Lovasz_softmax_PointCloud, self).__init__()
        self.classes = classes
        self.ignore = ignore

    def forward(self, probas, labels):

        B, C, N = probas.size()
        probas = probas.permute(0, 2, 1).contiguous().view(-1, C)
        labels = labels.view(-1)
        if self.ignore is not None:
            valid = (labels != self.ignore)
            vprobas = probas[valid.nonzero(as_tuple=False).squeeze()]
            vlabels = labels[valid]
            return lovasz_softmax_flat(vprobas, vlabels, classes=self.classes)
        else:
            return lovasz_softmax_flat(probas, labels, classes=self.classes)

modules/loss/Lovasz_Softmax.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import torch
import numpy as np
import cv2
from matplotlib import pyplot as plt


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


# def one_hot_pred_from_label(y_pred, labels):
#     y_true = torch.zeros_like(y_pred)
#     ones = torch.ones_like(y_pred)
#     indexes = [l for l in labels]
#     y_true[torch.arange(labels.size(0)), indexes] = ones[torch.arange(
#         labels.size(0)), indexes]
#     return y_true


# def keep_variance_fn(x):
#     return x + 1e-3


# class SoftmaxHeteroscedasticLoss(torch.nn.Module):
#     def __init__(self):
#         super(SoftmaxHeteroscedasticLoss, self).__init__()
#         self.adf_softmax = adf.Softmax(
#             dim=1, keep_variance_fn=keep_variance_fn)

#     def forward(self, outputs, targets, eps=1e-5):
#         mean, var = self.adf_softmax(*outputs)
#         targets = torch.nn.functional.one_hot(
#             targets, num_classes=20).permute(0, 3, 1, 2).float()

#         precision = 1 / (var + eps)
#         return torch.mean(0.5 * precision * (targets - mean) ** 2 + 0.5 * torch.log(var + eps))


def save_to_txtlog(logdir, logfile, message):
    f = open(logdir + '/' + logfile, "a")
    f.write(message + '\n')
    f.close()
    return


def save_checkpoint(to_save, logdir, suffix=""):
    # Save the weights
    torch.save(to_save, logdir +
               "/MFMOS" + suffix)


def get_mpl_colormap(cmap_name):
    cmap = plt.get_cmap(cmap_name)
    # Initialize the matplotlib color map
    sm = plt.cm.ScalarMappable(cmap=cmap)
    # Obtain linear color range
    color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1]
    return color_range.reshape(256, 1, 3)


def make_log_img(depth, mask, pred, gt, color_fn, movable=False):
    # input should be [depth, pred, gt]
    # make range image (normalized to 0,1 for saving)
    depth = (cv2.normalize(depth, None, alpha=0, beta=1,
                           norm_type=cv2.NORM_MINMAX,
                           dtype=cv2.CV_32F) * 255.0).astype(np.uint8)
    out_img = cv2.applyColorMap(
        depth, get_mpl_colormap('viridis')) * mask[..., None]
    # make label prediction
    pred_color = color_fn((pred * mask).astype(np.int32), movable=movable)
    out_img = np.concatenate([out_img, pred_color], axis=0)
    # make label gt
    gt_color = color_fn(gt, movable=movable)
    out_img = np.concatenate([out_img, gt_color], axis=0)
    return (out_img).astype(np.uint8)

def show_scans_in_training(proj_mask, in_vol, argmax, proj_labels, color_fn, movable=False):
    # get the first scan in batch and project points
    mask_np = proj_mask[0].cpu().numpy()
    depth_np = in_vol[0][0].cpu().numpy()
    pred_np = argmax[0].cpu().numpy()
    gt_np = proj_labels[0].cpu().numpy()
    out = make_log_img(depth_np, mask_np, pred_np, gt_np, color_fn, movable=movable)

    mask_np = proj_mask[1].cpu().numpy()
    depth_np = in_vol[1][0].cpu().numpy()
    pred_np = argmax[1].cpu().numpy()
    gt_np = proj_labels[1].cpu().numpy()
    out2 = make_log_img(depth_np, mask_np, pred_np, gt_np, color_fn, movable=movable)

    out = np.concatenate([out, out2], axis=0)

    cv2.imshow("sample_training", out)
    cv2.waitKey(1)


class iouEval:
    def __init__(self, n_classes, device, ignore=None):
        self.n_classes = n_classes
        self.device = device
        # if ignore is larger than n_classes, consider no ignoreIndex
        self.ignore = torch.tensor(ignore).long()
        self.include = torch.tensor(
            [n for n in range(self.n_classes) if n not in self.ignore]).long()
        print("[IOU EVAL] IGNORE: ", self.ignore)
        print("[IOU EVAL] INCLUDE: ", self.include)
        self.reset()

    def num_classes(self):
        return self.n_classes

    def reset(self):
        self.conf_matrix = torch.zeros(
            (self.n_classes, self.n_classes), device=self.device).long()
        self.ones = None
        self.last_scan_size = None  # for when variable scan size is used

    def addBatch(self, x, y):  # x=preds, y=targets
        # if numpy, pass to pytorch
        # to tensor
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(np.array(x)).long().to(self.device)
        if isinstance(y, np.ndarray):
            y = torch.from_numpy(np.array(y)).long().to(self.device)

        # sizes should be "batch_size x H x W"
        x_row = x.reshape(-1)  # de-batchify
        y_row = y.reshape(-1)  # de-batchify

        # idxs are labels and predictions
        idxs = torch.stack([x_row, y_row], dim=0)

        # ones is what I want to add to conf when I
        if self.ones is None or self.last_scan_size != idxs.shape[-1]:
            self.ones = torch.ones((idxs.shape[-1]), device=self.device).long()
            self.last_scan_size = idxs.shape[-1]

        # make confusion matrix (cols = gt, rows = pred)
        self.conf_matrix = self.conf_matrix.index_put_(
            tuple(idxs), self.ones, accumulate=True)

    def getStats(self):
        # remove fp and fn from confusion on the ignore classes cols and rows
        conf = self.conf_matrix.clone().double()
        conf[self.ignore] = 0
        conf[:, self.ignore] = 0

        # get the clean stats
        tp = conf.diag()
        fp = conf.sum(dim=1) - tp
        fn = conf.sum(dim=0) - tp
        return tp, fp, fn

    def getIoU(self):
        tp, fp, fn = self.getStats()
        intersection = tp
        union = tp + fp + fn + 1e-15
        iou = intersection / union
        iou_mean = (intersection[self.include] / union[self.include]).mean()
        return iou_mean, iou  # returns "iou mean", "iou per class" ALL CLASSES

    def getacc(self):
        tp, fp, fn = self.getStats()
        total_tp = tp.sum()
        total = tp[self.include].sum() + fp[self.include].sum() + 1e-15
        acc_mean = total_tp / total
        return acc_mean  # returns "acc mean"

modules/tools.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
import datetime
import os
import time
import imp
import cv2
import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
import __init__ as booger


import torch.optim as optim
from tensorboardX import SummaryWriter as Logger
from common.sync_batchnorm.batchnorm import convert_model
from common.warmupLR import warmupLR

# from modules.SalsaNextWithMotionAttention import SalsaNextWithMotionAttention
from modules.MFMOS import MFMOS

from modules.loss.Lovasz_Softmax import Lovasz_softmax, Lovasz_softmax_PointCloud
from modules.tools import AverageMeter, iouEval, save_checkpoint, show_scans_in_training, save_to_txtlog, make_log_img

from torch import distributed as dist

class Trainer():
    def __init__(self, ARCH, DATA, datadir, logdir, path=None, point_refine=False, local_rank=0):
        # parameters
        self.ARCH = ARCH
        self.DATA = DATA
        self.datadir = datadir
        self.logdir = logdir
        self.path = path
        self.epoch = 0
        self.point_refine = point_refine
        self.local_rank = local_rank

        self.batch_time_t = AverageMeter()
        self.data_time_t = AverageMeter()
        self.batch_time_e = AverageMeter()

        # put logger where it belongs
        self.tb_logger = Logger(self.logdir + "/tb")
        self.info = {"train_update": 0,
                     "train_loss": 0, "train_acc": 0, "train_iou": 0,
                     "valid_loss": 0, "valid_acc": 0, "valid_iou": 0,
                     "best_train_iou": 0, "best_val_iou": 0}

        # get the data
        parserModule = imp.load_source("parserModule",
                                       f"{booger.TRAIN_PATH}/common/dataset/{self.DATA['name']}/parser.py")
        self.parser = parserModule.Parser(root=self.datadir,
                                          train_sequences=self.DATA["split"]["train"],
                                          valid_sequences=self.DATA["split"]["valid"],
                                          test_sequences=None,
                                          split='train',
                                          labels=self.DATA["labels"],
                                          residual_aug=self.ARCH["train"]["residual_aug"],
                                          valid_residual_delta_t=self.ARCH["train"]["valid_residual_delta_t"],
                                          color_map=self.DATA["color_map"],
                                          learning_map=self.DATA["moving_learning_map"],
                                          movable_learning_map=self.DATA["movable_learning_map"],
                                          learning_map_inv=self.DATA["moving_learning_map_inv"],
                                          movable_learning_map_inv=self.DATA["movable_learning_map_inv"],
                                          sensor=self.ARCH["dataset"]["sensor"],
                                          max_points=self.ARCH["dataset"]["max_points"],
                                          batch_size=self.ARCH["train"]["batch_size"],
                                          workers=self.ARCH["train"]["workers"],
                                          gt=True,
                                          shuffle_train=True)

        self.set_loss_weight()

        with torch.no_grad():
            self.model = MFMOS(nclasses=self.parser.get_n_classes(),
                               movable_nclasses=self.parser.get_n_classes(movable=True), params=self.ARCH)

        self.set_gpu_cuda()
        self.set_loss_function(point_refine)
        self.set_optim_scheduler()

        # if need load the pre-trained model from checkpoint
        if self.path is not None:
            self.load_pretrained_model()

    def set_loss_weight(self):
        """
            Used to calculate the weights for each class
            weights for loss (and bias)
        """
        epsilon_w = self.ARCH["train"]["epsilon_w"]
        content = torch.zeros(self.parser.get_n_classes(), dtype=torch.float)
        movable_content = torch.zeros(self.parser.get_n_classes(movable=True), dtype=torch.float)
        for cl, freq in self.DATA["content"].items():
            x_cl = self.parser.to_xentropy(cl)   # map actual class to xentropy class
            content[x_cl] += freq

            movable_x_cl = self.parser.to_xentropy(cl, movable=True)
            movable_content[movable_x_cl] += freq
        self.loss_w = 1 / (content + epsilon_w)  # get weights
        self.movable_loss_w = 1 / (movable_content + epsilon_w)  # get weights
        for x_cl, w in enumerate(self.loss_w):   # ignore the ones necessary to ignore
            if self.DATA["learning_ignore"][x_cl]:    # don't weigh
                self.loss_w[x_cl] = 0

        for x_cl, w in enumerate(self.movable_loss_w):   # ignore the ones necessary to ignore
            if self.DATA["learning_ignore"][x_cl]:    # don't weigh
                self.movable_loss_w[x_cl] = 0
        print("Loss weights from content: ", self.loss_w.data)
        print("Movable Loss weights from content: ", self.movable_loss_w.data)

    def set_loss_function(self, point_refine):
        """
            Used to define the loss function, multiple gpus need to be parallel
            # self.dice = DiceLoss().to(self.device)
            # self.dice = nn.DataParallel(self.dice).cuda()
        """
        self.criterion = nn.NLLLoss(weight=self.loss_w.double()).cuda()
        self.movable_criterion = nn.NLLLoss(weight=self.movable_loss_w.double()).cuda()
        if not point_refine:
            # self.ls = Lovasz_softmax(ignore=0).to(self.device)
            self.ls = Lovasz_softmax(ignore=0).cuda()
            self.movable_ls = Lovasz_softmax(ignore=0).cuda()
        else:
            self.ls = Lovasz_softmax_PointCloud(ignore=0).to(self.device)

        # loss as dataparallel too (more images in batch)
        # if self.n_gpus > 1:
            # self.criterion = nn.DataParallel(self.criterion).cuda()  # spread in gpus
            # self.movable_criterion = nn.DataParallel(self.movable_criterion).cuda()  # spread in gpus
            # self.ls = nn.DataParallel(self.ls).cuda()
            # print()

    def set_gpu_cuda(self):
        """
            Used to set gpus and cuda information
        """
        self.gpu = False
        self.multi_gpu = False
        self.n_gpus = 0
        self.model_single = self.model
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("Training in device: ", self.device)

        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            # cudnn.benchmark = True
            # cudnn.fastest = True
            self.gpu = True
            self.n_gpus = 1
            self.model.cuda()

        if torch.cuda.is_available() and torch.cuda.device_count() > 1:
            print(f"Let's use {torch.cuda.device_count()} GPUs!")
            # self.model = nn.DataParallel(self.model)      # spread in gpus
            self.model = convert_model(self.model).cuda() # sync batchnorm
            self.model = torch.nn.parallel.DistributedDataParallel(self.model,
                                                                   device_ids=[self.local_rank],
                                                                   output_device=self.local_rank,
                                                                   find_unused_parameters=True,
                                                                   broadcast_buffers=False)

            self.model_single = self.model.module         # single model to get weight names
            self.multi_gpu = True
            self.n_gpus = torch.cuda.device_count()


    def set_optim_scheduler(self):
        """
            Used to set the optimizer and scheduler
        """
        self.optimizer = optim.SGD([{'params': self.model.parameters()}],
                                   lr=self.ARCH["train"]["lr"],
                                   momentum=self.ARCH["train"]["momentum"],
                                   weight_decay=self.ARCH["train"]["w_decay"])

        # Use warmup learning rate
        # post decay and step sizes come in epochs and we want it in steps
        steps_per_epoch = self.parser.get_train_size()
        up_steps = int(self.ARCH["train"]["wup_epochs"] * steps_per_epoch)
        final_decay = self.ARCH["train"]["lr_decay"] ** (1 / steps_per_epoch)
        self.scheduler = warmupLR(optimizer=self.optimizer,
                                  lr=self.ARCH["train"]["lr"],
                                  warmup_steps=up_steps,
                                  momentum=self.ARCH["train"]["momentum"],
                                  decay=final_decay)

    def load_pretrained_model(self):
        """
            If you want to resume training, reload the model
        """
        torch.nn.Module.dump_patches = True
        if not self.point_refine:
            checkpoint = "MFMOS"
            w_dict = torch.load(f"{self.path}/{checkpoint}", map_location=lambda storage, loc: storage)
            self.model.load_state_dict(w_dict['state_dict'], strict=True)
            self.optimizer.load_state_dict(w_dict['optimizer'])
            self.epoch = w_dict['epoch'] + 1
            self.scheduler.load_state_dict(w_dict['scheduler'])
            print("dict epoch:", w_dict['epoch'])
            self.info = w_dict['info']
            print("info", w_dict['info'])
            print("load the pretrained model of MFMOS")
        else:
            checkpoint = "MFMOS_valid_best"
            w_dict = torch.load(f"{self.path}/{checkpoint}", map_location=lambda storage, loc: storage)
            # self.model.load_state_dict(w_dict['state_dict'], strict=True)
            self.model.load_state_dict({k.replace('module.',''):v for k,v in w_dict['state_dict'].items()})
            self.optimizer.load_state_dict(w_dict['optimizer'])
            print("load the coarse model of MFMOS_valid_best")


    def calculate_estimate(self, epoch, iter):
        estimate = int((self.data_time_t.avg + self.batch_time_t.avg) *
                       (self.parser.get_train_size() * self.ARCH['train']['max_epochs'] - (
                           iter + 1 + epoch * self.parser.get_train_size()))) + \
            int(self.batch_time_e.avg * self.parser.get_valid_size() * (
                self.ARCH['train']['max_epochs'] - (epoch)))
        return str(datetime.timedelta(seconds=estimate))

    @staticmethod
    def save_to_tensorboard(logdir, logger, info, epoch, w_summary=False, model=None, img_summary=False, imgs=[]):
        # save scalars
        for tag, value in info.items():
            if 'valid_classes' in tag:
                continue # solve the bug of saving tensor type of value
            logger.add_scalar(tag, value, epoch)

        # save summaries of weights and biases
        if w_summary and model:
            for tag, value in model.named_parameters():
                tag = tag.replace('.', '/')
                logger.add_histogram(tag, value.data.cpu().numpy(), epoch)
                if value.grad is not None:
                    logger.add_histogram(tag + '/grad', value.grad.data.cpu().numpy(), epoch)

        if img_summary and len(imgs) > 0:
            directory = os.path.join(logdir, "predictions")
            if not os.path.isdir(directory):
                os.makedirs(directory)
            for i, img in enumerate(imgs):
                name = os.path.join(directory, str(i) + ".png")
                cv2.imwrite(name, img)

    def init_evaluator(self):
        self.ignore_class = []
        for i, w in enumerate(self.loss_w):
            if w < 1e-10:
                self.ignore_class.append(i)
                print("Ignoring class ", i, " in IoU evaluation")
        # self.evaluator = iouEval(self.parser.get_n_classes(),
        #                          self.device, self.ignore_class)
        # self.movable_evaluator = iouEval(self.parser.get_n_classes(),
        #                                  self.device, self.ignore_class)
        self.evaluator = iouEval(self.parser.get_n_classes(),
                                 self.local_rank, self.ignore_class)
        self.movable_evaluator = iouEval(self.parser.get_n_classes(movable=True),
                                         self.local_rank, self.ignore_class)

    def train(self):

        self.init_evaluator()

        # train for n epochs
        for epoch in range(self.epoch, self.ARCH["train"]["max_epochs"]):
            self.parser.train_sampler.set_epoch(epoch)
            # train for 1 epoch
            acc, iou, loss, update_mean, hetero_l = self.train_epoch(train_loader=self.parser.get_train_set(),
                                                                     model=self.model,
                                                                     all_criterion=(self.criterion, self.movable_criterion),
                                                                     optimizer=self.optimizer,
                                                                     epoch=epoch,
                                                                     all_evaluator=(self.evaluator, self.movable_evaluator),
                                                                     scheduler=self.scheduler,
                                                                     color_fn=self.parser.to_color,
                                                                     report=self.ARCH["train"]["report_batch"],
                                                                     show_scans=self.ARCH["train"]["show_scans"])

            if self.local_rank == 0:
                # update the info dict and save the training checkpoint
                self.update_training_info(epoch, acc, iou, loss, update_mean, hetero_l)

            # evaluate on validation set
            if epoch % self.ARCH["train"]["report_epoch"] == 0 and self.local_rank == 0:
                acc, iou, loss, rand_img, hetero_l = self.validate(val_loader=self.parser.get_valid_set(),
                                                                   model=self.model,
                                                                   all_criterion=(self.criterion, self.movable_criterion),
                                                                   all_evaluator=(self.evaluator, self.movable_evaluator),
                                                                   class_func=self.parser.get_xentropy_class_string,
                                                                   color_fn=self.parser.to_color,
                                                                   save_scans=self.ARCH["train"]["save_scans"])

                self.update_validation_info(epoch, acc, iou, loss, hetero_l)

            if self.local_rank == 0:
                # save to tensorboard log
                Trainer.save_to_tensorboard(logdir=self.logdir,
                                            logger=self.tb_logger,
                                            info=self.info,
                                            epoch=epoch,
                                            w_summary=self.ARCH["train"]["save_summary"],
                                            model=self.model_single,
                                            img_summary=self.ARCH["train"]["save_scans"],
                                            imgs=rand_img)

            dist.barrier()

        print('Finished Training')

        return

    def train_epoch(self, train_loader, model, all_criterion, optimizer,
                    epoch, all_evaluator, scheduler, color_fn, report=10,
                    show_scans=False):

        losses = AverageMeter()

        moving_losses = AverageMeter()
        movable_losses = AverageMeter()

        acc = AverageMeter()
        iou = AverageMeter()

        movable_acc = AverageMeter()
        movable_iou = AverageMeter()

        hetero_l = AverageMeter()
        update_ratio_meter = AverageMeter()

        evaluator, movable_evaluator = all_evaluator
        criterion, movable_criterion = all_criterion

        # empty the cache to train now
        # if self.gpu:
        #     torch.cuda.empty_cache()

        # switch to train mode
        model.train()

        end = time.time()
        for i, (in_vol, proj_mask, all_proj_labels, _, path_seq, path_name,
                _, _, _, _, _, _, _, _, _) in enumerate(train_loader):
            # measure data loading time
            self.data_time_t.update(time.time() - end)

            proj_labels, movable_proj_labels = all_proj_labels

            if not self.multi_gpu and self.gpu:
                in_vol = in_vol.cuda()
            if self.gpu:
                proj_labels = proj_labels.cuda().long()
                movable_proj_labels = movable_proj_labels.cuda().long()

            output, _, movable_output, _ = model(in_vol)

            # loss_m = criterion(torch.log(output.clamp(min=1e-8)), proj_labels) + self.ls(output, proj_labels.long())
            moving_loss_m = criterion(torch.log(output.clamp(min=1e-8)).double(), proj_labels).float() + self.ls(output, proj_labels.long())
            movable_loss_m = movable_criterion(torch.log(movable_output.clamp(min=1e-8)).double(), movable_proj_labels).float() \
                             + self.movable_ls(movable_output, movable_proj_labels.long())

            loss_m = moving_loss_m + movable_loss_m

            optimizer.zero_grad()
            # if self.n_gpus > 1:
            #     idx = torch.ones(self.n_gpus).cuda()
            #     loss_m.backward(idx)
            # else:
            #     loss_m.backward()
            loss_m.backward()
            optimizer.step()

            # measure accuracy and record loss
            loss = loss_m.mean()
            with torch.no_grad():
                evaluator.reset()
                argmax = output.argmax(dim=1)
                evaluator.addBatch(argmax, proj_labels)
                accuracy = evaluator.getacc()
                jaccard, class_jaccard = evaluator.getIoU()

                movable_evaluator.reset()
                movable_argmax = movable_output.argmax(dim=1)
                movable_evaluator.addBatch(movable_argmax, movable_proj_labels)
                movable_accuracy = movable_evaluator.getacc()
                movable_jaccard, movable_class_jaccard = movable_evaluator.getIoU()

            losses.update(loss.item(), in_vol.size(0))
            moving_losses.update(moving_loss_m.mean().item(), in_vol.size(0))
            movable_losses.update(movable_loss_m.mean().item(), in_vol.size(0))

            acc.update(accuracy.item(), in_vol.size(0))
            iou.update(jaccard.item(), in_vol.size(0))

            movable_acc.update(movable_accuracy.item(), in_vol.size(0))
            movable_iou.update(movable_jaccard.item(), in_vol.size(0))

            # measure elapsed time
            self.batch_time_t.update(time.time() - end)
            end = time.time()

            # get gradient updates and weights, so I can print the relationship of
            # their norms
            update_ratios = []
            for g in self.optimizer.param_groups:
                lr = g["lr"]
                for value in g["params"]:
                    if value.grad is not None:
                        w = np.linalg.norm(value.data.cpu().numpy().reshape((-1)))
                        update = np.linalg.norm(-max(lr, 1e-10) * value.grad.cpu().numpy().reshape((-1)))
                        update_ratios.append(update / max(w, 1e-10))
            update_ratios = np.array(update_ratios)
            update_mean = update_ratios.mean()
            update_std = update_ratios.std()
            update_ratio_meter.update(update_mean)  # over the epoch

            if show_scans:
                show_scans_in_training(
                    proj_mask, in_vol, argmax, proj_labels, color_fn)
                show_scans_in_training(
                    proj_mask, in_vol, movable_argmax, movable_proj_labels, color_fn, movable=True)

            if i % self.ARCH["train"]["report_batch"] == 0 and self.local_rank == 0:
                str_line = ('Lr: {lr:.3e} | '
                            'Update: {umean:.3e} mean,{ustd:.3e} std | '
                            'Epoch: [{0}][{1}/{2}] | '
                            'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) | '
                            'Data {data_time.val:.3f} ({data_time.avg:.3f}) | '
                            'Loss {loss.val:.4f} ({loss.avg:.4f}) | '
                            'MovingLoss {moving_losses.val:.4f} ({moving_losses.avg:.4f}) | '
                            'MovableLoss {movable_losses.val:.4f} ({movable_losses.avg:.4f}) | '
                            'MovingAcc {moving_acc.val:.3f} ({moving_acc.avg:.3f}) | '
                            'MovingIoU {moving_iou.val:.3f} ({moving_iou.avg:.3f}) | '
                            'movable_acc {movable_acc.val:.3f} ({movable_acc.avg:.3f}) | '
                            'movable_IoU {movable_iou.val:.3f} ({movable_iou.avg:.3f}) | '
                            '[{estim}]').format(
                    epoch, i, len(train_loader), batch_time=self.batch_time_t, data_time=self.data_time_t, loss=losses,
                    moving_losses=moving_losses, movable_losses=movable_losses,
                    moving_acc=acc, moving_iou=iou,
                    movable_acc=movable_acc, movable_iou=movable_iou,
                    lr=lr, umean=update_mean, ustd=update_std, estim=self.calculate_estimate(epoch, i))
                print(str_line)
                save_to_txtlog(self.logdir, 'log.txt', str_line)

            # step scheduler
            scheduler.step()

        return acc.avg, iou.avg, losses.avg, update_ratio_meter.avg, hetero_l.avg

    def validate(self, val_loader, model, all_criterion, all_evaluator, class_func, color_fn, save_scans=False):
        evaluator, movable_evaluator = all_evaluator
        criterion, movable_criterion = all_criterion

        losses = AverageMeter()
        moving_losses = AverageMeter()
        movable_losses = AverageMeter()

        jaccs = AverageMeter()
        wces = AverageMeter()
        acc = AverageMeter()
        iou = AverageMeter()

        movable_jaccs = AverageMeter()
        movable_wces = AverageMeter()
        movable_acc = AverageMeter()
        movable_iou = AverageMeter()

        hetero_l = AverageMeter()
        rand_imgs = []

        # switch to evaluate mode
        model.eval()
        evaluator.reset()
        movable_evaluator.reset()

        # empty the cache to infer in high res
        # if self.gpu:
        #     torch.cuda.empty_cache()

        with torch.no_grad():
            end = time.time()
            for i, (in_vol, proj_mask, all_proj_labels, _, path_seq, path_name,
                    _, _, _, _, _, _, _, _, _)\
                    in enumerate(tqdm(val_loader, desc="Validation:", ncols=80)):
                proj_labels, movable_proj_labels = all_proj_labels

                if not self.multi_gpu and self.gpu:
                    in_vol = in_vol.cuda()
                    proj_mask = proj_mask.cuda()
                if self.gpu:
                    proj_labels = proj_labels.cuda(non_blocking=True).long()
                    movable_proj_labels = movable_proj_labels.cuda(non_blocking=True).long()

                # compute output
                output, _, movable_output, _ = model(in_vol)

                log_out = torch.log(output.clamp(min=1e-8))
                movable_log_out = torch.log(movable_output.clamp(min=1e-8))

                # wce = criterion(log_out, proj_labels)
                jacc = self.ls(output, proj_labels)
                wce = criterion(log_out.double(), proj_labels).float()
                moving_loss = wce + jacc

                movable_jacc = self.ls(movable_output, movable_proj_labels)
                movable_wce = movable_criterion(movable_log_out.double(), movable_proj_labels).float()
                movable_loss = movable_wce + jacc

                loss = moving_loss + movable_loss

                # measure accuracy and record loss
                argmax = output.argmax(dim=1)
                evaluator.addBatch(argmax, proj_labels)

                movable_argmax = movable_output.argmax(dim=1)
                movable_evaluator.addBatch(movable_argmax, movable_proj_labels)

                losses.update(loss.mean().item(), in_vol.size(0))

                moving_losses.update(moving_loss.mean().item(), in_vol.size(0))
                jaccs.update(jacc.mean().item(), in_vol.size(0))
                wces.update(wce.mean().item(), in_vol.size(0))

                movable_losses.update(movable_loss.mean().item(), in_vol.size(0))
                movable_jaccs.update(movable_jacc.mean().item(), in_vol.size(0))
                movable_wces.update(movable_wce.mean().item(), in_vol.size(0))


                if save_scans:
                    # get the first scan in batch and project points
                    mask_np = proj_mask[0].cpu().numpy()
                    depth_np = in_vol[0][0].cpu().numpy()
                    pred_np = argmax[0].cpu().numpy()
                    gt_np = proj_labels[0].cpu().numpy()
                    out = make_log_img(depth_np, mask_np, pred_np, gt_np, color_fn)
                    rand_imgs.append(out)

                # measure elapsed time
                self.batch_time_e.update(time.time() - end)
                end = time.time()

            accuracy = evaluator.getacc()
            jaccard, class_jaccard = evaluator.getIoU()
            acc.update(accuracy.item(), in_vol.size(0))
            iou.update(jaccard.item(), in_vol.size(0))

            movable_accuracy = movable_evaluator.getacc()
            movable_jaccard, movable_class_jaccard = movable_evaluator.getIoU()
            movable_acc.update(movable_accuracy.item(), in_vol.size(0))
            movable_iou.update(movable_jaccard.item(), in_vol.size(0))

            str_line = ("*" * 80 + '\n'
                        'Validation set:\n'
                        'Time avg per batch {batch_time.avg:.3f}\n'
                        'Loss avg {loss.avg:.4f}\n'
                        'MovingLoss avg {moving_loss.avg:.4f}\n'
                        'MovableLoss avg {movable_loss.avg:.4f}\n'
                        'MovingJaccard avg {moving_jac.avg:.4f}\n'
                        'MovingWCE avg {moving_wces.avg:.4f}\n'
                        'MovingAcc avg {moving_acc.avg:.6f}\n'
                        'MovingIoU avg {moving_iou.avg:.6f}\n'
                        'MovableJaccard avg {movable_jac.avg:.4f}\n'
                        'MovableWCE avg {movable_wce.avg:.4f}\n'
                        'MovableAcc avg {movable_acc.avg:.6f}\n'
                        'MovableIoU avg {movable_iou.avg:.6f}'
                        ).format(
                            batch_time=self.batch_time_e, loss=losses,
                            moving_loss=moving_losses, movable_loss=movable_losses,
                            moving_jac=jaccs, moving_wces=wces, moving_acc=acc, moving_iou=iou,
                            movable_jac=movable_jaccs, movable_wce=movable_wces,
                            movable_acc=movable_acc, movable_iou=movable_iou)
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)

            print("-" * 80)
            # print also classwise
            for i, jacc in enumerate(class_jaccard):
                self.info["valid_classes/" + class_func(i)] = jacc
                str_line = 'IoU class {i:} [{class_str:}] = {jacc:.6f}'.format(
                    i=i, class_str=class_func(i), jacc=jacc)
                print(str_line)
                save_to_txtlog(self.logdir, 'log.txt', str_line)

            print("-" * 80)
            for i, jacc in enumerate(movable_class_jaccard):
                self.info["valid_classes/" + class_func(i, movable=True)] = jacc
                str_line = 'IoU class {i:} [{class_str:}] = {jacc:.6f}'.format(
                    i=i, class_str=class_func(i, movable=True), jacc=jacc)
                print(str_line)
                save_to_txtlog(self.logdir, 'log.txt', str_line)

            str_line = '*' * 80
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)

        return acc.avg, iou.avg, losses.avg, rand_imgs, hetero_l.avg

    def update_training_info(self, epoch, acc, iou, loss, update_mean, hetero_l):
        # update info
        self.info["train_update"] = update_mean
        self.info["train_loss"] = loss
        self.info["train_acc"] = acc
        self.info["train_iou"] = iou
        self.info["train_hetero"] = hetero_l

        # remember best iou and save checkpoint
        state = {'epoch': epoch,
                 'state_dict': self.model.state_dict(),
                 'optimizer': self.optimizer.state_dict(),
                 'info': self.info,
                 'scheduler': self.scheduler.state_dict()}
        save_checkpoint(state, self.logdir, suffix="")

        if self.info['train_iou'] > self.info['best_train_iou']:
            print("Best mean iou in training set so far, save model!")
            self.info['best_train_iou'] = self.info['train_iou']
            state = {'epoch': epoch,
                     'state_dict': self.model.state_dict(),
                     'optimizer': self.optimizer.state_dict(),
                     'info': self.info,
                     'scheduler': self.scheduler.state_dict()}
            save_checkpoint(state, self.logdir, suffix="_train_best")

    def update_validation_info(self, epoch, acc, iou, loss, hetero_l):
        # update info
        self.info["valid_loss"] = loss
        self.info["valid_acc"] = acc
        self.info["valid_iou"] = iou
        self.info['valid_heteros'] = hetero_l

        # remember best iou and save checkpoint
        if self.info['valid_iou'] > self.info['best_val_iou']:
            str_line = ("Best mean iou in validation so far, save model!\n" + "*" * 80)
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)
            self.info['best_val_iou'] = self.info['valid_iou']

            # save the weights!
            state = {'epoch': epoch,
                     'state_dict': self.model.state_dict(),
                     'optimizer': self.optimizer.state_dict(),
                     'info': self.info,
                     'scheduler': self.scheduler.state_dict()}
            save_checkpoint(state, self.logdir, suffix="_valid_best")
            save_checkpoint(state, self.logdir, suffix=f"_valid_best_{epoch}")


modules/trainer.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import datetime
import os
import time
import imp
import cv2
import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
from modules.trainer import Trainer
import __init__ as booger


import torch.optim as optim
from tensorboardX import SummaryWriter as Logger
from common.sync_batchnorm.batchnorm import convert_model
from common.warmupLR import warmupLR

# from modules.SalsaNextWithMotionAttention import SalsaNextWithMotionAttention
# from modules.loss.Lovasz_Softmax import Lovasz_softmax, Lovasz_softmax_PointCloud
from modules.tools import AverageMeter, iouEval, save_checkpoint, show_scans_in_training, save_to_txtlog, make_log_img

from modules.PointRefine.spvcnn import SPVCNN
# from modules.PointRefine.spvcnn_lite import SPVCNN
from torchsparse.utils.quantize import sparse_quantize
from torchsparse.utils.collate import sparse_collate
from torchsparse import SparseTensor


class TrainerRefine(Trainer):
    def __init__(self, ARCH, DATA, datadir, logdir, path=None):

        super(TrainerRefine, self).__init__(ARCH, DATA, datadir, logdir, path, point_refine=True)

        """ New variables for the PointRefine module """
        net_config = {'num_classes': self.parser.get_n_classes(),
                      'cr': 1.0, 'pres': 0.05, 'vres': 0.05}
        self.refine_module = SPVCNN(num_classes=net_config['num_classes'],
                                    cr=net_config['cr'],
                                    pres=net_config['pres'],
                                    vres=net_config['vres'])

        if torch.cuda.is_available():
            if torch.cuda.device_count() > 0:
               self.refine_module.cuda()
            if torch.cuda.device_count() > 1:
                self.refine_module = nn.DataParallel(self.refine_module)

        self.set_refine_optim_scheduler()
        """"""""""""""""""""""""""""""""""""""""""""""""


    def set_refine_optim_scheduler(self):
        """
            Used to set the optimizer and scheduler for PointRefine module
        """
        self.refine_optimizer = optim.Adam(
            [{'params': self.refine_module.parameters()}], self.ARCH["train"]["lr"])

        self.refine_scheduler = torch.optim.lr_scheduler.ExponentialLR(
            self.refine_optimizer, 0.95)


    def train(self):

        self.init_evaluator()

        # train for n epochs
        for epoch in range(self.epoch, self.ARCH["train"]["max_epochs"]):

            # train for 1 epoch
            acc, iou, loss, update_mean, hetero_l = self.train_epoch(train_loader=self.parser.get_train_set(),
                                                                     model=self.model,
                                                                     criterion=self.criterion,
                                                                     optimizer=self.optimizer,
                                                                     epoch=epoch,
                                                                     evaluator=self.evaluator,
                                                                     scheduler=self.scheduler,
                                                                     color_fn=self.parser.to_color,
                                                                     report=self.ARCH["train"]["report_batch"],
                                                                     show_scans=self.ARCH["train"]["show_scans"])

            # update the info dict and save the training checkpoint
            self.update_training_info(epoch, acc, iou, loss, update_mean, hetero_l)

            # evaluate on validation set
            if epoch % self.ARCH["train"]["report_epoch"] == 0:
                acc, iou, loss, rand_img, hetero_l = self.validate(val_loader=self.parser.get_valid_set(),
                                                                   model=self.model,
                                                                   criterion=self.criterion,
                                                                   evaluator=self.evaluator,
                                                                   class_func=self.parser.get_xentropy_class_string,
                                                                   color_fn=self.parser.to_color,
                                                                   save_scans=self.ARCH["train"]["save_scans"])
                
                self.update_validation_info(epoch, acc, iou, loss, hetero_l)

            # save to tensorboard log
            Trainer.save_to_tensorboard(logdir=self.logdir,
                                        logger=self.tb_logger,
                                        info=self.info,
                                        epoch=epoch,
                                        w_summary=self.ARCH["train"]["save_summary"],
                                        model=self.model_single,
                                        img_summary=self.ARCH["train"]["save_scans"],
                                        imgs=rand_img)

        print('Finished Training')

        return

    def train_epoch(self, train_loader, model, criterion, optimizer,
                    epoch, evaluator, scheduler, color_fn, report=10,
                    show_scans=False):

        losses = AverageMeter()
        acc = AverageMeter()
        iou = AverageMeter()
        hetero_l = AverageMeter()
        update_ratio_meter = AverageMeter()

        # empty the cache to train now
        # if self.gpu:
        #     torch.cuda.empty_cache()

        # switch to train mode
        start_epoch = 5
        # switch to train mode
        if epoch <= start_epoch:
            model.eval()
        else:
            model.train()
        
        self.refine_module.train()

        end = time.time()

        for i, (in_vol, proj_mask, all_proj_labels, unproj_labels, path_seq, path_name,
                p_x, p_y, proj_range, unproj_range, _, unproj_xyz, _, unproj_remissions, n_points)\
                in enumerate(train_loader):

            proj_labels, movable_proj_labels = all_proj_labels

            # measure data loading time
            self.data_time_t.update(time.time() - end)

            if not self.multi_gpu and self.gpu:
                in_vol = in_vol.cuda()
                #proj_mask = proj_mask.cuda()
            if self.gpu:
                proj_labels = proj_labels.cuda().long()
                unproj_labels = unproj_labels.cuda().long()
                unproj_xyz = unproj_xyz.cuda()

            # compute output
            # output = model(in_vol)
            if epoch <= start_epoch:
                with torch.no_grad():
                    output, last_feature, movable_output, movalbe_last_feature = model(in_vol)
            else:
                output, last_feature, movable_output, movalbe_last_feature = model(in_vol)


            tmp_pred = []
            tmp_labels = []
            for j in range(len(n_points)):
                _npoints = n_points[j]
                _px = p_x[j, :_npoints]
                _py = p_y[j, :_npoints]
                _unproj_labels = unproj_labels[j, :_npoints]
                _points_xyz = unproj_xyz[j, :_npoints]

                # put in original pointcloud using indexes
                _points_feature = last_feature[j, :, _py, _px]

                # Filter out duplicate points
                coords = np.round(_points_xyz[:, :3].cpu().numpy() / 0.05)
                coords -= coords.min(0, keepdims=1)
                coords, indices, inverse = sparse_quantize(coords, return_index=True, return_inverse=True)
                coords = torch.tensor(coords, dtype=torch.int, device='cuda')

                feats = _points_feature.permute(1,0)[indices]

                inputs = SparseTensor(coords=coords, feats=feats)
                inputs = sparse_collate([inputs]).cuda()

                _predict = self.refine_module(inputs)
                _predict = _predict[inverse].permute(1,0)

                tmp_pred.append(_predict)
                tmp_labels.append(_unproj_labels)

            predict = torch.cat(tmp_pred, -1).unsqueeze(0)
            unproj_labels = torch.cat(tmp_labels).unsqueeze(0)

            loss_m = criterion(torch.log(predict.clamp(min=1e-8)).double(), unproj_labels).float() + self.ls(predict, unproj_labels)

            if epoch <= start_epoch:
                self.refine_optimizer.zero_grad()
            else:
                self.refine_optimizer.zero_grad()
                optimizer.zero_grad()
            if self.n_gpus > 1:
                idx = torch.ones(self.n_gpus).cuda()
                loss_m.backward(idx)
            else:
                loss_m.backward()

            if epoch <= start_epoch:
                self.refine_optimizer.step()
            else:
                self.refine_optimizer.step()
                optimizer.step()

            # measure accuracy and record loss
            loss = loss_m.mean()
            with torch.no_grad():
                evaluator.reset()
                # argmax = output.argmax(dim=1)
                # evaluator.addBatch(argmax, proj_labels)
                argmax = predict.argmax(dim=1)
                evaluator.addBatch(argmax, unproj_labels)
                accuracy = evaluator.getacc()
                jaccard, class_jaccard = evaluator.getIoU()

            losses.update(loss.item(), in_vol.size(0))
            acc.update(accuracy.item(), in_vol.size(0))
            iou.update(jaccard.item(), in_vol.size(0))

            # measure elapsed time
            self.batch_time_t.update(time.time() - end)
            end = time.time()

            # get gradient updates and weights, so I can print the relationship of
            # their norms
            # update_ratios = []
            for g in self.optimizer.param_groups:
                lr = g["lr"]

            if show_scans:
                show_scans_in_training(
                    proj_mask, in_vol, argmax, proj_labels, color_fn)

            if i % self.ARCH["train"]["report_batch"] == 0:
                str_line = ('Lr: {lr:.3e} | '
                            # 'Update: {umean:.3e} mean,{ustd:.3e} std | '
                            'Epoch: [{0}][{1}/{2}] | '
                            'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) | '
                            'Data {data_time.val:.3f} ({data_time.avg:.3f}) | '
                            'MovingLoss {moving_loss.val:.4f} ({moving_loss.avg:.4f}) | '
                            'MovingAcc {moving_acc.val:.3f} ({moving_acc.avg:.3f}) | '
                            'MovingIoU {moving_iou.val:.3f} ({moving_iou.avg:.3f}) | [{estim}]').format(
                    epoch, i, len(train_loader), batch_time=self.batch_time_t,
                    data_time=self.data_time_t, moving_loss=losses, moving_acc=acc, moving_iou=iou, lr=lr,
                    # umean=update_mean, ustd=update_std, 
                    estim=self.calculate_estimate(epoch, i))
                print(str_line)
                save_to_txtlog(self.logdir, 'log.txt', str_line)

            # step scheduler
            # scheduler.step()
            if i != 0 and i % 1000 == 0:
                self.refine_scheduler.step()

            if i != 0 and i % 1000 == 0: 
                # evaluate on validation set
                tmp_evaluator = iouEval(self.parser.get_n_classes(), self.device, self.ignore_class)

                _acc, _iou, _loss, rand_img, _hetero_l = self.validate(val_loader=self.parser.get_valid_set(),
                                                                       model=self.model,
                                                                       criterion=self.criterion,
                                                                       evaluator=tmp_evaluator,
                                                                       class_func=self.parser.get_xentropy_class_string,
                                                                       color_fn=self.parser.to_color,
                                                                       save_scans=self.ARCH["train"]["save_scans"])
                
                self.update_validation_info(epoch, _acc, _iou, _loss, _hetero_l, i)

            # step scheduler
            scheduler.step()

        return acc.avg, iou.avg, losses.avg, update_ratio_meter.avg, hetero_l.avg

    def validate(self, val_loader, model, criterion, evaluator, class_func, color_fn, save_scans=False):
        losses = AverageMeter()
        jaccs = AverageMeter()
        wces = AverageMeter()
        acc = AverageMeter()
        iou = AverageMeter()
        hetero_l = AverageMeter()
        rand_imgs = []

        # switch to evaluate mode
        model.eval()
        evaluator.reset()
        self.refine_module.eval()

        # empty the cache to infer in high res
        # if self.gpu:
        #     torch.cuda.empty_cache()

        with torch.no_grad():
            end = time.time()
            for i, (in_vol, proj_mask, all_proj_labels, unproj_labels, path_seq, path_name,
                    p_x, p_y, proj_range, unproj_range, _, unproj_xyz, _, unproj_remissions, n_points)\
                    in enumerate(tqdm(val_loader, desc="Validation:", ncols=80)):

                proj_labels, movable_proj_labels = all_proj_labels

                if not self.multi_gpu and self.gpu:
                    in_vol = in_vol.cuda()
                    proj_mask = proj_mask.cuda()
                if self.gpu:
                    proj_labels = proj_labels.cuda(non_blocking=True).long()
                    unproj_labels = unproj_labels.cuda().long()
                    unproj_xyz = unproj_xyz.cuda()

                # compute output
                # output = model(in_vol)
                with torch.no_grad():
                    output, last_feature, movable_output, movable_last_feature = model(in_vol)
                
                """ Reproject 2D features to 3D based on indices and form sparse Tensor"""
                tmp_pred = []
                tmp_labels = []
                for j in range(len(n_points)):
                    _npoints = n_points[j]
                    _px = p_x[j, :_npoints]
                    _py = p_y[j, :_npoints]
                    _unproj_labels = unproj_labels[j, :_npoints]
                    _points_xyz = unproj_xyz[j, :_npoints]

                    # put in original pointcloud using indexes
                    _points_feature = last_feature[j, :, _py, _px]
                
                # Filter out duplicate points
                coords = np.round(_points_xyz[:, :3].cpu().numpy() / 0.05)
                coords -= coords.min(0, keepdims=1)
                coords, indices, inverse = sparse_quantize(coords, return_index=True, return_inverse=True)
                coords = torch.tensor(coords, dtype=torch.int, device='cuda')
                feats = _points_feature.permute(1,0)[indices]
                # feats = torch.cat((_points_xyz[indices] , _points_feature.permute(1, 0)[indices]), dim=1)
                # feats = torch.cat((_points_xyz[indices], _range_pred.permute(1, 0)[indices], _points_feature.permute(1, 0)[indices]), dim=1)

                inputs = SparseTensor(coords=coords, feats=feats)
                inputs = sparse_collate([inputs]).cuda()

                # _concate_feat = torch.cat((_points_xyz, _range_pred.permute(1, 0)), dim=1).unsqueeze(0) 
                _predict = self.refine_module(inputs)
                _predict = _predict[inverse].permute(1,0)

                tmp_pred.append(_predict)
                tmp_labels.append(_unproj_labels)

                predict = torch.cat(tmp_pred, -1).unsqueeze(0)
                unproj_labels = torch.cat(tmp_labels).unsqueeze(0)
                """""""""""""""""""""""" 

                log_out = torch.log(predict.clamp(min=1e-8))
                jacc = self.ls(predict, unproj_labels)
                # wce = criterion(log_out, proj_labels)
                wce = criterion(log_out.double(), unproj_labels).float()
                loss = wce + jacc

                # measure accuracy and record loss
                # argmax = output.argmax(dim=1)
                # evaluator.addBatch(argmax, proj_labels)
                argmax = predict.argmax(dim=1)
                evaluator.addBatch(argmax, unproj_labels)
                losses.update(loss.mean().item(), in_vol.size(0))
                jaccs.update(jacc.mean().item(),in_vol.size(0))
                wces.update(wce.mean().item(),in_vol.size(0))

                if save_scans:
                    # get the first scan in batch and project points
                    mask_np = proj_mask[0].cpu().numpy()
                    depth_np = in_vol[0][0].cpu().numpy()
                    pred_np = argmax[0].cpu().numpy()
                    gt_np = proj_labels[0].cpu().numpy()
                    out = make_log_img(depth_np, mask_np, pred_np, gt_np, color_fn)
                    rand_imgs.append(out)

                # measure elapsed time
                self.batch_time_e.update(time.time() - end)
                end = time.time()

            accuracy = evaluator.getacc()
            jaccard, class_jaccard = evaluator.getIoU()
            acc.update(accuracy.item(), in_vol.size(0))
            iou.update(jaccard.item(), in_vol.size(0))

            str_line = ("*" * 80 + '\n'
                        'Validation set:\n'
                        'Time avg per batch {batch_time.avg:.3f}\n'
                        'MovingLoss avg {moving_loss.avg:.4f}\n'
                        'Jaccard avg {jac.avg:.4f}\n'
                        'WCE avg {wces.avg:.4f}\n'
                        'MovingAcc avg {moving_acc.avg:.6f}\n'
                        'MovingIoU avg {moving_iou.avg:.6f}').format(
                            batch_time=self.batch_time_e, moving_loss=losses,
                            jac=jaccs, wces=wces, moving_acc=acc, moving_iou=iou)
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)

            # print also classwise
            for i, jacc in enumerate(class_jaccard):
                self.info["valid_classes/" + class_func(i)] = jacc
                str_line = 'IoU class {i:} [{class_str:}] = {jacc:.6f}'.format(
                    i=i, class_str=class_func(i), jacc=jacc)
                print(str_line)
                save_to_txtlog(self.logdir, 'log.txt', str_line)
            str_line = '*' * 80
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)

        return acc.avg, iou.avg, losses.avg, rand_imgs, hetero_l.avg

    def update_training_info(self, epoch, acc, iou, loss, update_mean, hetero_l):
        # update info
        self.info["train_update"] = update_mean
        self.info["train_loss"] = loss
        self.info["train_acc"] = acc
        self.info["train_iou"] = iou
        self.info["train_hetero"] = hetero_l

        # remember best iou and save checkpoint
        state = {'epoch': epoch,
                 'main_state_dict': self.model.state_dict(),
                 'refine_state_dict': self.refine_module.state_dict(),
                 'main_optimizer': self.optimizer.state_dict(),
                 'refine_optimizer': self.refine_optimizer.state_dict(),
                 'info': self.info,
                 'scheduler': self.scheduler.state_dict(),
                 'refine_scheduler': self.refine_scheduler.state_dict()}
        save_checkpoint(state, self.logdir, suffix="")

        if self.info['train_iou'] > self.info['best_train_iou']:
            print("Best mean iou in training set so far, save model!")
            self.info['best_train_iou'] = self.info['train_iou']
            state = {'epoch': epoch,
                     'main_state_dict': self.model.state_dict(),
                     'refine_state_dict': self.refine_module.state_dict(),
                     'main_optimizer': self.optimizer.state_dict(),
                     'refine_optimizer': self.refine_optimizer.state_dict(),
                     'info': self.info,
                     'scheduler': self.scheduler.state_dict(),
                     'refine_scheduler': self.refine_scheduler.state_dict()}
            save_checkpoint(state, self.logdir, suffix="_train_best")

    def update_validation_info(self, epoch, acc, iou, loss, hetero_l, iter=None):
        # update info
        self.info["valid_loss"] = loss
        self.info["valid_acc"] = acc
        self.info["valid_iou"] = iou
        self.info['valid_heteros'] = hetero_l

        # remember best iou and save checkpoint
        if self.info['valid_iou'] > self.info['best_val_iou']:
            str_line = ("Best mean iou in validation so far, save model!\n" + "*" * 80)
            print(str_line)
            save_to_txtlog(self.logdir, 'log.txt', str_line)

            self.info['best_val_iou'] = self.info['valid_iou']

            # save the weights!
            state = {'epoch': epoch, 
                     'main_state_dict': self.model.state_dict(),
                     'refine_state_dict': self.refine_module.state_dict(),
                     'main_optimizer': self.optimizer.state_dict(),
                     'refine_optimizer': self.refine_optimizer.state_dict(),
                     'info': self.info,
                     'scheduler': self.scheduler.state_dict(),
                     'refine_scheduler': self.refine_scheduler.state_dict()}
            save_checkpoint(state, self.logdir, suffix="_SIEM_valid_best")
            if iter is None:
                save_checkpoint(state, self.logdir, suffix=f"_SIEM_valid_best_{epoch}")
            else:
                save_checkpoint(state, self.logdir, suffix=f"_SIEM_valid_best_{epoch}_it{iter}")


modules/trainer_refine.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
import imp
import time
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import __init__ as booger

from tqdm import tqdm
from modules.KNN import KNN
# from modules.SalsaNextWithMotionAttention import *
from modules.MFMOS import *

from modules.PointRefine.spvcnn import SPVCNN
# from modules.PointRefine.spvcnn_lite import SPVCNN

class User():
    def __init__(self, ARCH, DATA, datadir, outputdir, modeldir, split, point_refine=False, save_movable=False):
        # parameters
        self.ARCH = ARCH
        self.DATA = DATA
        self.datadir = datadir
        self.outputdir = outputdir
        self.modeldir = modeldir
        self.split = split
        self.post = None
        self.infer_batch_size = 1
        self.point_refine = point_refine
        self.save_movable = save_movable
        # get the data
        parserModule = imp.load_source("parserModule",
                                       f"{booger.TRAIN_PATH}/common/dataset/{self.DATA['name']}/parser.py")
        self.parser = parserModule.Parser(root=self.datadir,
                                          train_sequences=self.DATA["split"]["train"],
                                          valid_sequences=self.DATA["split"]["valid"],
                                          test_sequences=self.DATA["split"]["test"],
                                          split=self.split,
                                          labels=self.DATA["labels"],
                                          residual_aug=self.ARCH["train"]["residual_aug"],
                                          valid_residual_delta_t=self.ARCH["train"]["valid_residual_delta_t"],
                                          color_map=self.DATA["color_map"],
                                          learning_map=self.DATA["moving_learning_map"],
                                          movable_learning_map=self.DATA["movable_learning_map"],
                                          learning_map_inv=self.DATA["moving_learning_map_inv"],
                                          movable_learning_map_inv=self.DATA["movable_learning_map_inv"],
                                          sensor=self.ARCH["dataset"]["sensor"],
                                          max_points=self.ARCH["dataset"]["max_points"],
                                          batch_size=self.infer_batch_size,
                                          workers=2, # self.ARCH["train"]["workers"],
                                          gt=True,
                                          shuffle_train=False)

        with torch.no_grad():
            torch.nn.Module.dump_patches = True
            if not point_refine:
                self.model = MFMOS(nclasses=self.parser.get_n_classes(),
                                   movable_nclasses=self.parser.get_n_classes(movable=True),
                                   params=ARCH, num_batch=self.infer_batch_size)
                self.model = nn.DataParallel(self.model)
                checkpoint = "MFMOS_valid_best"
                w_dict = torch.load(f"{self.modeldir}/{checkpoint}", map_location=lambda storage, loc: storage)
                self.model.load_state_dict(w_dict['state_dict'], strict=True)

                self.set_knn_post()
            else:
                self.model = MFMOS(nclasses=self.parser.get_n_classes(),
                                   movable_nclasses=self.parser.get_n_classes(movable=True),
                                   params=ARCH, num_batch=self.infer_batch_size)
                self.model = nn.DataParallel(self.model)
                checkpoint = "MFMOS_SIEM_valid_best"
                w_dict = torch.load(f"{self.modeldir}/{checkpoint}", map_location=lambda storage, loc: storage)
                self.model.load_state_dict({f"module.{k}":v for k,v in w_dict['main_state_dict'].items()}, strict=True)

                net_config = {'num_classes': self.parser.get_n_classes(), 'cr': 1.0, 'pres': 0.05, 'vres': 0.05}
                self.refine_module = SPVCNN(num_classes=net_config['num_classes'],
                                            cr=net_config['cr'],
                                            pres=net_config['pres'],
                                            vres=net_config['vres'])
                self.refine_module = nn.DataParallel(self.refine_module)
                w_dict = torch.load(f"{modeldir}/{checkpoint}", map_location=lambda storage, loc: storage)
                self.refine_module.load_state_dict({f"module.{k}":v for k,v in w_dict['refine_state_dict'].items()}, strict=True)

        self.set_gpu_cuda()

    def set_knn_post(self):
        # use knn post processing?
        if self.ARCH["post"]["KNN"]["use"]:
            self.post = KNN(self.ARCH["post"]["KNN"]["params"], self.parser.get_n_classes())

    def set_gpu_cuda(self):
        # GPU?
        self.gpu = False
        self.model_single = self.model
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("Infering in device: ", self.device)
        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            cudnn.benchmark = True
            cudnn.fastest = True
            self.gpu = True
            self.model.cuda()
            if self.point_refine:
                self.refine_module.cuda()


    def infer(self):
        cnn, knn = [], []

        if self.split == 'valid':
            self.infer_subset(loader=self.parser.get_valid_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
        elif self.split == 'train':
            self.infer_subset(loader=self.parser.get_train_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
        elif self.split == 'test':
            self.infer_subset(loader=self.parser.get_test_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
        elif self.split == None:
            self.infer_subset(loader=self.parser.get_train_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
            # do valid set
            self.infer_subset(loader=self.parser.get_valid_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
            # do test set
            self.infer_subset(loader=self.parser.get_test_set(),
                              to_orig_fn=self.parser.to_original,
                              cnn=cnn, knn=knn)
        else:
            raise NotImplementedError

        print(f"Mean CNN inference time:{'%.8f'%np.mean(cnn)}\t std:{'%.8f'%np.std(cnn)}")
        print(f"Mean KNN inference time:{'%.8f'%np.mean(knn)}\t std:{'%.8f'%np.std(knn)}")
        print(f"Total Frames: {len(cnn)}")
        print("Finished Infering")

        return

    def infer_subset(self, loader, to_orig_fn, cnn, knn):

        # switch to evaluate mode
        self.model.eval()

        # empty the cache to infer in high res
        if self.gpu:
            torch.cuda.empty_cache()

        with torch.no_grad():

            end = time.time()

            for i, (proj_in, proj_mask, _, _, path_seq, path_name,
                    p_x, p_y, proj_range, unproj_range, _, _, _, _, npoints)\
                    in enumerate(tqdm(loader, ncols=80)):
                # first cut to rela size (batch size one allows it)
                p_x = p_x[0, :npoints]
                p_y = p_y[0, :npoints]
                proj_range = proj_range[0, :npoints]
                unproj_range = unproj_range[0, :npoints]
                path_seq = path_seq[0]
                path_name = path_name[0]

                if self.gpu:
                    proj_in = proj_in.cuda()
                    p_x = p_x.cuda()
                    p_y = p_y.cuda()
                    if self.post:
                        proj_range = proj_range.cuda()
                        unproj_range = unproj_range.cuda()
                
                end = time.time()
                # compute output
                proj_output, _, movable_proj_output, _ = self.model(proj_in)
                res = time.time() - end
                cnn.append(res)

                proj_argmax = proj_output[0].argmax(dim=0)
                movable_proj_argmax = movable_proj_output[0].argmax(dim=0)

                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                end = time.time()
                # print(f"Network seq {path_seq} scan {path_name} in {res} sec")

                # if knn --> use knn to postprocess
                # 	else put in original pointcloud using indexes
                if self.post:
                    unproj_argmax = self.post(proj_range, unproj_range,
                                              proj_argmax, p_x, p_y)
                    if self.save_movable:
                        movable_unproj_argmax = self.post(proj_range, unproj_range,
                                                          movable_proj_argmax, p_x, p_y)
                else:
                    unproj_argmax = proj_argmax[p_y, p_x]
                    if self.save_movable:
                        movable_unproj_argmax = movable_proj_argmax[p_y, p_x]

                # measure elapsed time
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                res = time.time() - end
                knn.append(res)
                # print(f"KNN Infered seq {path_seq} scan {path_name} in {res} sec")

                # save scan # get the first scan in batch and project scan
                pred_np = unproj_argmax.cpu().numpy()
                pred_np = pred_np.reshape((-1)).astype(np.int32)

                # map to original label
                pred_np = to_orig_fn(pred_np)

                path = os.path.join(self.outputdir, "sequences", path_seq, "predictions", path_name)
                pred_np.tofile(path)

                if self.save_movable:
                    movable_pred_np = movable_unproj_argmax.cpu().numpy()
                    movable_pred_np = movable_pred_np.reshape((-1)).astype(np.int32)

                    # map to original label
                    movable_pred_np = to_orig_fn(movable_pred_np, movable=True)
                    path = os.path.join(self.outputdir, "sequences", path_seq, "predictions_movable", path_name)
                    movable_pred_np.tofile(path)

                    movable_pred_np[np.where(pred_np == 251)] = 251
                    path = os.path.join(self.outputdir, "sequences", path_seq, "predictions_fuse", path_name)
                    movable_pred_np.tofile(path)
modules/user.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
import imp
import time
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import __init__ as booger

from tqdm import tqdm
from modules.user import User
# from modules.SalsaNextWithMotionAttention import *

# from modules.PointRefine.spvcnn import SPVCNN
# from modules.PointRefine.spvcnn_lite import SPVCNN
from torchsparse.utils.quantize import sparse_quantize
from torchsparse.utils.collate import sparse_collate
from torchsparse import SparseTensor


class UserRefine(User):
    def __init__(self, ARCH, DATA, datadir, outputdir, modeldir, split, save_movable=False):

        super(UserRefine, self).__init__(ARCH, DATA, datadir, outputdir, modeldir, split,
                                         point_refine=True, save_movable=save_movable)

    def infer(self):
        coarse, reproj, refine = [], [], []

        if self.split == 'valid':
            self.infer_subset(loader=self.parser.get_valid_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
        elif self.split == 'train':
            self.infer_subset(loader=self.parser.get_train_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
        elif self.split == 'test':
            self.infer_subset(loader=self.parser.get_test_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
        elif self.split == None:
            self.infer_subset(loader=self.parser.get_train_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
            self.infer_subset(loader=self.parser.get_valid_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
            self.infer_subset(loader=self.parser.get_test_set(),
                              to_orig_fn=self.parser.to_original,
                              coarse=coarse, reproj=reproj, refine=refine)
        else:
            raise NotImplementedError

        print(f"Mean Coarse inference time:{'%.8f'%np.mean(coarse)}\t std:{'%.8f'%np.std(coarse)}")
        print(f"Mean Reproject inference time:{'%.8f'%np.mean(reproj)}\t std:{'%.8f'%np.std(reproj)}")
        print(f"Mean Refine inference time:{'%.8f'%np.mean(refine)}\t std:{'%.8f'%np.std(refine)}")
        print(f"Total Frames: {len(coarse)}")
        print("Finished Infering")

        return

    def infer_subset(self, loader, to_orig_fn, coarse, reproj, refine):

        # switch to evaluate mode
        self.model.eval()
        self.refine_module.eval()

        # empty the cache to infer in high res
        if self.gpu:
            torch.cuda.empty_cache()

        with torch.no_grad():

            end = time.time()

            for i, (proj_in, proj_mask, _, _, path_seq, path_name,
                    p_x, p_y, proj_range, unproj_range, _, unproj_xyz, _, _, npoints)\
                    in enumerate(tqdm(loader, ncols=80)):

                # first cut to rela size (batch size one allows it)
                p_x = p_x[0, :npoints]
                p_y = p_y[0, :npoints]
                proj_range = proj_range[0, :npoints]
                unproj_range = unproj_range[0, :npoints]
                path_seq = path_seq[0]
                path_name = path_name[0]
                points_xyz = unproj_xyz[0, :npoints]

                if self.gpu:
                    proj_in = proj_in.cuda()
                    p_x = p_x.cuda()
                    p_y = p_y.cuda()
                    if self.post:
                        proj_range = proj_range.cuda()
                        unproj_range = unproj_range.cuda()

                end = time.time()
                # compute output
                proj_output, last_feature, movable_proj_output, _ = self.model(proj_in)


                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                res = time.time() - end
                coarse.append(res)

                if self.save_movable:
                    movable_proj_argmax = movable_proj_output[0].argmax(dim=0)
                    if self.post:
                        movable_unproj_argmax = self.post(proj_range, unproj_range,
                                                          movable_proj_argmax, p_x, p_y)
                    else:
                        movable_unproj_argmax = movable_proj_argmax[p_y, p_x]

                end = time.time()
                # print(f"CoarseModule seq {path_seq} scan {path_name} in {res} sec")

                """ Reproject 2D features to 3D based on indices and form sparse Tensor"""
                points_feature = last_feature[0, :, p_y, p_x]
                coords = np.round(points_xyz[:, :3].cpu().numpy() / 0.05)
                coords -= coords.min(0, keepdims=1)
                coords, indices, inverse = sparse_quantize(coords, return_index=True, return_inverse=True)
                coords = torch.tensor(coords, dtype=torch.int, device='cuda')
                feats = points_feature.permute(1,0)[indices] #torch.tensor(, dtype=torch.float)
                inputs = SparseTensor(coords=coords, feats=feats)
                inputs = sparse_collate([inputs]).cuda()
                """"""""""""""""""""""""

                # measure elapsed time
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                res = time.time() - end
                reproj.append(res)
                end = time.time()
                # print(f"DataConvert seq {path_seq} scan {path_name} in {res} sec")

                """ Input to PointHead, refine prediction """
                predict = self.refine_module(inputs)

                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                res = time.time() - end
                refine.append(res)
                # print(f"RefineModule seq {path_seq} scan {path_name} in {res} sec")

                predict = predict[inverse] #.permute(1,0)
                unproj_argmax = predict.argmax(dim=1)

                # save scan # get the first scan in batch and project scan
                pred_np = unproj_argmax.cpu().numpy()
                pred_np = pred_np.reshape((-1)).astype(np.int32)

                # map to original label
                pred_np = to_orig_fn(pred_np)

                path = os.path.join(self.outputdir, "sequences", path_seq, "predictions", path_name)
                pred_np.tofile(path)

                if self.save_movable:
                    movable_pred_np = movable_unproj_argmax.cpu().numpy()
                    movable_pred_np = movable_pred_np.reshape((-1)).astype(np.int32)

                    # map to original label
                    movable_pred_np = to_orig_fn(movable_pred_np, movable=True)
                    path = os.path.join(self.outputdir, "sequences", path_seq, "predictions_movable", path_name)
                    movable_pred_np.tofile(path)

                    movable_pred_np[np.where(pred_np == 251)] = 251
                    path = os.path.join(self.outputdir, "sequences", path_seq, "predictions_fuse", path_name)
                    movable_pred_np.tofile(path)

modules/user_refine.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import torch
print("device count: ", torch.cuda.device_count())
from torch import distributed as dist
dist.init_process_group(backend="nccl")
print("world_size: ", dist.get_world_size())

import random
import numpy as np
import __init__ as booger

from modules.trainer import Trainer
# from modules.SalsaNextWithMotionAttention import *
from modules.MFMOS import *

def set_seed(seed=1024):
    random.seed(seed)
    # os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.

    # torch.backends.cudnn.benchmark = False
    # torch.backends.cudnn.deterministic = True
    # torch.backends.cudnn.enabled = False
    # If we need to reproduce the results, increase the training speed
    #    set benchmark = False
    # If we don’t need to reproduce the results, improve the network performance as much as possible
    #    set benchmark = True


if __name__ == '__main__':
    parser = get_args(flags="train")
    FLAGS, unparsed = parser.parse_known_args()
    local_rank = FLAGS.local_rank
    torch.cuda.set_device(local_rank)

    FLAGS.log = os.path.join(FLAGS.log, datetime.now().strftime("%Y-%-m-%d-%H:%M") + FLAGS.name)
    print(FLAGS.log)
    # open arch / data config file
    ARCH = load_yaml(FLAGS.arch_cfg)
    DATA = load_yaml(FLAGS.data_cfg)

    params = MFMOS(nclasses=3, params=ARCH, movable_nclasses=3)
    pytorch_total_params = sum(p.numel() for p in params.parameters() if p.requires_grad)
    del params

    if local_rank == 0:
        make_logdir(FLAGS=FLAGS, resume_train=False) # create log folder
        check_pretrained_dir(FLAGS.pretrained)       # does model folder exist?
        backup_to_logdir(FLAGS=FLAGS)                # backup code and config files to logdir

    set_seed()
    # create trainer and start the training
    trainer = Trainer(ARCH, DATA, FLAGS.dataset, FLAGS.log, FLAGS.pretrained, local_rank=local_rank)

    if local_rank == 0:
        print("----------")
        print("INTERFACE:")
        print("  dataset:", FLAGS.dataset)
        print("  arch_cfg:", FLAGS.arch_cfg)
        print("  data_cfg:", FLAGS.data_cfg)
        print("  Total of Trainable Parameters: {}".format(millify(pytorch_total_params, 2)))
        print("  log:", FLAGS.log)
        print("  pretrained:", FLAGS.pretrained)
        print("  Augmentation for residual: {}, interval in validation: {}".format(ARCH["train"]["residual_aug"],
                                                                                   ARCH["train"]["valid_residual_delta_t"]))
        print("----------\n")

    trainer.train()

train.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
import random
import numpy as np
import torch
import __init__ as booger

from datetime import datetime
from utils.utils import *
from modules.trainer_refine import TrainerRefine
# from modules.SalsaNextWithMotionAttention import *


def set_seed(seed=1024):
    random.seed(seed)
    # os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    # torch.backends.cudnn.benchmark = False
    # torch.backends.cudnn.deterministic = True
    # torch.backends.cudnn.enabled = False
    # If we need to reproduce the results, increase the training speed
    #    set benchmark = False
    # If we don’t need to reproduce the results, improve the network performance as much as possible
    #    set benchmark = True


if __name__ == '__main__':
    parser = get_args(flags="train")
    FLAGS, unparsed = parser.parse_known_args()
    FLAGS.log = os.path.join(FLAGS.log, datetime.now().strftime("%Y-%-m-%d-%H:%M") + FLAGS.name)
    
    # open arch / data config file
    ARCH = load_yaml(FLAGS.arch_cfg)
    DATA = load_yaml(FLAGS.data_cfg)
    
    # params = SalsaNextWithMotionAttention(nclasses=3, params=ARCH)
    # pytorch_total_params = sum(p.numel() for p in params.parameters() if p.requires_grad)
    # del params

    make_logdir(FLAGS=FLAGS, resume_train=False) # create log folder
    check_pretrained_dir(FLAGS.pretrained) # does model folder exist?
    backup_to_logdir(FLAGS=FLAGS, pretrain_model=True) # backup code and config files to logdir

    set_seed()
    # create trainer and start the training
    trainer = TrainerRefine(ARCH, DATA, FLAGS.dataset, FLAGS.log, FLAGS.pretrained)

    print("----------")
    print("INTERFACE:")
    print("  dataset:", FLAGS.dataset)
    print("  arch_cfg:", FLAGS.arch_cfg)
    print("  data_cfg:", FLAGS.data_cfg)
    print("  log:", FLAGS.log)
    print("  pretrained:", FLAGS.pretrained)
    print("  Augmentation for residual: {}, interval in validation: {}".format(ARCH["train"]["residual_aug"],
                                                                               ARCH["train"]["valid_residual_delta_t"]))
    print("----------\n")

    trainer.train()

train_2stage.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# Developed by Jiadai Sun
# 	and the main_funciton 'prosess_one_seq' refers to Xieyuanli Chen’s gen_residual_images.py
# This file is covered by the LICENSE file in the root of this project.
# Brief: This script generates residual images

import os
os.environ["OMP_NUM_THREADS"] = "4"
import yaml
import numpy as np
import matplotlib.pyplot as plt

from tqdm import tqdm
from icecream import ic
from kitti_utils import load_poses, load_calib, load_files, load_vertex

try:
    from c_gen_virtual_scan import gen_virtual_scan as range_projection
except:
    print("Using clib by $export PYTHONPATH=$PYTHONPATH:<path-to-library>")
    print("Currently using python-lib to generate range images.")
    from kitti_utils import range_projection


def check_and_makedirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


def load_yaml(path):
    if yaml.__version__ >= '5.1':
        config = yaml.load(open(path), Loader=yaml.FullLoader)
    else:
        config = yaml.load(open(path))
    return config


def process_one_seq(config):
    # specify parameters
    num_frames = config['num_frames']
    debug = config['debug']
    normalize = config['normalize']
    num_last_n = config['num_last_n']
    visualize = config['visualize']
    visualization_folder = config['visualization_folder']

    # specify the output folders
    residual_image_folder = config['residual_image_folder']
    check_and_makedirs(residual_image_folder)

    if visualize:
        check_and_makedirs(visualization_folder)

    # load poses
    pose_file = config['pose_file']
    poses = np.array(load_poses(pose_file))
    inv_frame0 = np.linalg.inv(poses[0])

    # load calibrations
    calib_file = config['calib_file']
    T_cam_velo = load_calib(calib_file)
    T_cam_velo = np.asarray(T_cam_velo).reshape((4, 4))
    T_velo_cam = np.linalg.inv(T_cam_velo)

    # convert kitti poses from camera coord to LiDAR coord
    new_poses = []
    for pose in poses:
        new_poses.append(T_velo_cam.dot(inv_frame0).dot(pose).dot(T_cam_velo))
    poses = np.array(new_poses)

    # load LiDAR scans
    scan_folder = config['scan_folder']
    scan_paths = load_files(scan_folder)

    # test for the first N scans
    if num_frames >= len(poses) or num_frames <= 0:
        print('generate training data for all frames with number of: ', len(poses))
    else:
        poses = poses[:num_frames]
        scan_paths = scan_paths[:num_frames]

    range_image_params = config['range_image']

    # generate residual images for the whole sequence
    for frame_idx in tqdm(range(len(scan_paths))):
        file_name = os.path.join(residual_image_folder, str(frame_idx).zfill(6))
        diff_image = np.full((range_image_params['height'], range_image_params['width']), 0, dtype=np.float32)  # [H,W] range (0 is no data)

        # for the first N frame we generate a dummy file
        if frame_idx < num_last_n:
            np.save(file_name, diff_image)
        else:
            # load current scan and generate current range image
            current_pose = poses[frame_idx]
            current_scan = load_vertex(scan_paths[frame_idx])
            current_range = range_projection(current_scan.astype(np.float32),
                                             range_image_params['height'], range_image_params['width'],
                                             range_image_params['fov_up'], range_image_params['fov_down'],
                                             range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]

            # load last scan, transform into the current coord and generate a transformed last range image
            last_pose = poses[frame_idx - num_last_n]
            last_scan = load_vertex(scan_paths[frame_idx - num_last_n])
            last_scan_transformed = np.linalg.inv(current_pose).dot(last_pose).dot(last_scan.T).T
            last_range_transformed = range_projection(last_scan_transformed.astype(np.float32),
                                             range_image_params['height'], range_image_params['width'],
                                             range_image_params['fov_up'], range_image_params['fov_down'],
                                             range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]

            # generate residual image
            valid_mask = (current_range > range_image_params['min_range']) & \
                            (current_range < range_image_params['max_range']) & \
                            (last_range_transformed > range_image_params['min_range']) & \
                            (last_range_transformed < range_image_params['max_range'])
            difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask])

            if normalize:
                difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask]) / current_range[valid_mask]

            diff_image[valid_mask] = difference

            if debug:
                fig, axs = plt.subplots(3)
                axs[0].imshow(last_range_transformed)
                axs[1].imshow(current_range)
                axs[2].imshow(diff_image, vmin=0, vmax=1)
                plt.show()

            if visualize:
                fig = plt.figure(frameon=False, figsize=(16, 10))
                fig.set_size_inches(20.48, 0.64)
                ax = plt.Axes(fig, [0., 0., 1., 1.])
                ax.set_axis_off()
                fig.add_axes(ax)
                ax.imshow(diff_image, vmin=0, vmax=1)
                image_name = os.path.join(visualization_folder, str(frame_idx).zfill(6))
                plt.savefig(image_name)
                plt.close()

            # save residual image
            np.save(file_name, diff_image)


if __name__ == '__main__':

    # load config file
    dataset = "DATAROOT"
    config_filename = 'config/data_preparing.yaml'
    config = load_yaml(config_filename)

    seq_list = [i for i in range(0, 11)]    # 0-10 for training
    if config['residual_aug']:
        residual_list = list(set([k * i for i in range(1, 9) for k in range(1, 4)]))
    else:
        residual_list = [1 * i for i in range(1, 9)]

    # used for kitti-raw and kitti-road
    for seq in seq_list: # sequences id

        for i in residual_list: # residual_image_i

            # Update the value in config to facilitate the iterative loop
            config['num_last_n'] = i
            config['scan_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/velodyne")
            config['pose_file'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/poses.txt")
            config['calib_file'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/calib.txt")
            config['residual_image_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/residual_images_{i}")
            config['visualization_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/visualization_{i}")

            ic(config)
            process_one_seq(config)

utils/auto_gen_residual_images.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# Developed by Jiadai Sun
# 	and the main_funciton 'prosess_one_seq' refers to Xieyuanli Chen’s gen_residual_images.py
# This file is covered by the LICENSE file in the root of this project.
# Brief: This script generates residual images

import os
import yaml
import numpy as np
import matplotlib.pyplot as plt

from tqdm import tqdm
from icecream import ic
from kitti_utils import load_poses, load_calib, load_files, load_vertex
from multiprocessing import Process

try:
    from c_gen_virtual_scan import gen_virtual_scan as range_projection
except:
    print("Using clib by $export PYTHONPATH=$PYTHONPATH:<path-to-library>")
    print("Currently using python-lib to generate range images.")
    from kitti_utils import range_projection


def check_and_makedirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


def load_yaml(path):
    if yaml.__version__ >= '5.1':
        config = yaml.load(open(path), Loader=yaml.FullLoader)
    else:
        config = yaml.load(open(path))
    return config


def process_one_seq(config, seq, i):
    # specify parameters
    num_frames = config['num_frames']
    debug = config['debug']
    normalize = config['normalize']
    num_last_n = config['num_last_n']
    visualize = config['visualize']
    visualization_folder = config['visualization_folder']

    # specify the output folders
    residual_image_folder = config['residual_image_folder']
    check_and_makedirs(residual_image_folder)

    if visualize:
        check_and_makedirs(visualization_folder)

    # load poses
    pose_file = config['pose_file']
    poses = np.array(load_poses(pose_file))
    inv_frame0 = np.linalg.inv(poses[0])

    # load calibrations
    calib_file = config['calib_file']
    T_cam_velo = load_calib(calib_file)
    T_cam_velo = np.asarray(T_cam_velo).reshape((4, 4))
    T_velo_cam = np.linalg.inv(T_cam_velo)

    # convert kitti poses from camera coord to LiDAR coord
    new_poses = []
    for pose in poses:
        new_poses.append(T_velo_cam.dot(inv_frame0).dot(pose).dot(T_cam_velo))
    poses = np.array(new_poses)

    # load LiDAR scans
    scan_folder = config['scan_folder']
    scan_paths = load_files(scan_folder)

    # test for the first N scans
    if num_frames >= len(poses) or num_frames <= 0:
        print('generate training data for all frames with number of: ', len(poses))
    else:
        poses = poses[:num_frames]
        scan_paths = scan_paths[:num_frames]

    range_image_params = config['range_image']

    # generate residual images for the whole sequence
    for frame_idx in tqdm(range(len(scan_paths))):
        file_name = os.path.join(residual_image_folder, str(frame_idx).zfill(6))
        diff_image = np.full((range_image_params['height'], range_image_params['width']), 0, dtype=np.float32)  # [H,W] range (0 is no data)

        # for the first N frame we generate a dummy file
        if frame_idx < num_last_n:
            np.save(file_name, diff_image)
        else:
            # load current scan and generate current range image
            current_pose = poses[frame_idx]
            current_scan = load_vertex(scan_paths[frame_idx])
            current_range = range_projection(current_scan.astype(np.float32),
                                             range_image_params['height'], range_image_params['width'],
                                             range_image_params['fov_up'], range_image_params['fov_down'],
                                             range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]

            # load last scan, transform into the current coord and generate a transformed last range image
            last_pose = poses[frame_idx - num_last_n]
            last_scan = load_vertex(scan_paths[frame_idx - num_last_n])
            last_scan_transformed = np.linalg.inv(current_pose).dot(last_pose).dot(last_scan.T).T
            last_range_transformed = range_projection(last_scan_transformed.astype(np.float32),
                                             range_image_params['height'], range_image_params['width'],
                                             range_image_params['fov_up'], range_image_params['fov_down'],
                                             range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]

            # generate residual image
            valid_mask = (current_range > range_image_params['min_range']) & \
                            (current_range < range_image_params['max_range']) & \
                            (last_range_transformed > range_image_params['min_range']) & \
                            (last_range_transformed < range_image_params['max_range'])
            difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask])

            if normalize:
                difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask]) / current_range[valid_mask]

            diff_image[valid_mask] = difference

            if debug:
                fig, axs = plt.subplots(3)
                axs[0].imshow(last_range_transformed)
                axs[1].imshow(current_range)
                axs[2].imshow(diff_image, vmin=0, vmax=1)
                plt.show()

            if visualize:
                fig = plt.figure(frameon=False, figsize=(16, 10))
                fig.set_size_inches(20.48, 0.64)
                ax = plt.Axes(fig, [0., 0., 1., 1.])
                ax.set_axis_off()
                fig.add_axes(ax)
                ax.imshow(diff_image, vmin=0, vmax=1)
                image_name = os.path.join(visualization_folder, str(frame_idx).zfill(6))
                plt.savefig(image_name)
                plt.close()

            # save residual image
            np.save(file_name, diff_image)

    print(f"seq {seq} residual {i} finish!")

def check_residual(path):
    folder = os.listdir(path)


if __name__ == '__main__':

    # load config file
    dataset = "DATAROOT"
    config_filename = '../config/data_preparing.yaml'
    config = load_yaml(config_filename)

    seq_split_part = 5
    residual_split_part = 5

    all_seq_list = [i for i in range(0, 11)]   # 0-10 for training
    if config['residual_aug']:
        residual_list = list(set([k * i for i in range(1, 9) for k in range(1, 4)]))
    else:
        residual_list = [1 * i for i in range(1, 9)]

    split_seq_list = [all_seq_list[i:i + (len(all_seq_list) // seq_split_part)] for i in range(0, len(all_seq_list), len(all_seq_list) // seq_split_part)]
    split_residual_list = [residual_list[i:i + (len(residual_list) // residual_split_part)] for i in range(0, len(residual_list), len(residual_list) // residual_split_part)]

    # process_num = [split_seq_list[i] * split_residual_list[i] for i in range(split_part+1)]

    print(split_seq_list)
    print(split_residual_list)

    for residual_id in range(len(split_residual_list)):
        for seq_id in range(len(split_seq_list)):
            processes = []
            for seq in split_seq_list[seq_id]:
                for i in split_residual_list[residual_id]:
                    config['num_last_n'] = i
                    config['scan_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/velodyne")
                    config['pose_file'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/poses.txt")
                    config['calib_file'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/calib.txt")
                    config['residual_image_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/residual_images_{i}")
                    config['visualization_folder'] = os.path.join(dataset, f"sequences/{'%02d' % seq}/visualization_{i}")

                    if os.path.exists(config['residual_image_folder']):
                        if len(os.listdir(config['residual_image_folder'])) == len(os.listdir(config['scan_folder'])):
                            print(f"seq {seq} residual {i} has been created!")
                            continue

                    process = Process(target=process_one_seq, args=(config, seq, i))  # 指定进程要执行的函数和参数
                    processes.append(process)
                    process.start()  # 启动进程

            for process in processes:
                process.join()  # 等待进程完成

utils/auto_gen_residual_images_mp.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
from matplotlib.pyplot import sca
import numpy as np


class LaserScan:
    """Class that contains LaserScan with x,y,z,r"""
    EXTENSIONS_SCAN = ['.bin']

    def __init__(self, project=False, H=64, W=1024, fov_up=3.0, fov_down=-25.0):
        self.project = project
        self.proj_H = H
        self.proj_W = W
        self.proj_fov_up = fov_up
        self.proj_fov_down = fov_down
        self.reset()

    def reset(self):
        """ Reset scan members. """
        self.points = np.zeros((0, 3), dtype=np.float32)        # [m, 3]: x, y, z
        self.remissions = np.zeros((0, 1), dtype=np.float32)    # [m ,1]: remission

        # projected range image - [H,W] range (-1 is no data)
        self.proj_range = np.full((self.proj_H, self.proj_W), -1, dtype=np.float32)

        # unprojected range (list of depths for each point)
        self.unproj_range = np.zeros((0, 1), dtype=np.float32)

        # projected point cloud xyz - [H,W,3] xyz coord (-1 is no data)
        self.proj_xyz = np.full((self.proj_H, self.proj_W, 3), -1, dtype=np.float32)

        # projected remission - [H,W] intensity (-1 is no data)
        self.proj_remission = np.full((self.proj_H, self.proj_W), -1, dtype=np.float32)

        # projected index (for each pixel, what I am in the pointcloud)
        # [H,W] index (-1 is no data)
        self.proj_idx = np.full((self.proj_H, self.proj_W), -1, dtype=np.int32)

        # for each point, where it is in the range image
        self.proj_x = np.zeros((0, 1), dtype=np.float32) # [m, 1]: x
        self.proj_y = np.zeros((0, 1), dtype=np.float32) # [m, 1]: y

        # mask containing for each pixel, if it contains a point or not
        self.proj_mask = np.zeros((self.proj_H, self.proj_W), dtype=np.int32) # [H,W] mask

    def size(self):
        """ Return the size of the point cloud. """
        return self.points.shape[0]

    def __len__(self):
        return self.size()

    def open_scan(self, filename):
        """ Open raw scan and fill in attributes"""
        # reset just in case there was an open structure
        self.reset()

        # check filename is string
        if not isinstance(filename, str):
            raise TypeError("Filename should be string type, "
                            "but was {type}".format(type=str(type(filename))))

        # check extension is a laserscan
        if not any(filename.endswith(ext) for ext in self.EXTENSIONS_SCAN):
            raise RuntimeError("Filename extension is not valid scan file.")

        # if all goes well, open pointcloud
        # import pdb; pdb.set_trace()
        scan = np.fromfile(filename, dtype=np.float32)
        scan = scan.reshape((-1, 4))
        # print(filename, scan.shape)

        # put in attribute
        points = scan[:, 0:3]    # get xyz
        remissions = scan[:, 3]  # get remission
        self.set_points(points, remissions)

    def set_points(self, points, remissions=None):
        """ Set scan attributes (instead of opening from file)
        """
        # reset just in case there was an open structure
        self.reset()

        # check scan makes sense
        if not isinstance(points, np.ndarray):
            raise TypeError("Scan should be numpy array")

        # check remission makes sense
        if remissions is not None and not isinstance(remissions, np.ndarray):
            raise TypeError("Remissions should be numpy array")

        # put in attribute
        self.points = points    # get xyz
        if remissions is not None:
            self.remissions = remissions  # get remission
        else:
            self.remissions = np.zeros((points.shape[0]), dtype=np.float32)

        # if projection is wanted, then do it and fill in the structure
        if self.project:
            self.do_range_projection()

    def do_range_projection(self):
        """ Project a pointcloud into a spherical projection image.projection.
                Function takes no arguments because it can be also called externally
                if the value of the constructor was not set (in case you change your
                mind about wanting the projection)
        """
        # laser parameters
        fov_up = self.proj_fov_up / 180.0 * np.pi      # field of view up in rad
        fov_down = self.proj_fov_down / 180.0 * np.pi  # field of view down in rad
        fov = abs(fov_down) + abs(fov_up)  # get field of view total in rad

        # get depth of all points
        depth = np.linalg.norm(self.points, 2, axis=1)

        # get scan components
        scan_x = self.points[:, 0]
        scan_y = self.points[:, 1]
        scan_z = self.points[:, 2]

        # get angles of all points
        yaw = -np.arctan2(scan_y, scan_x)
        pitch = np.arcsin(scan_z / depth)

        # get projections in image coords
        proj_x = 0.5 * (yaw / np.pi + 1.0)           # in [0.0, 1.0]
        proj_y = 1.0 - (pitch + abs(fov_down)) / fov # in [0.0, 1.0]

        # scale to image size using angular resolution
        proj_x *= self.proj_W   # in [0.0, W]
        proj_y *= self.proj_H   # in [0.0, H]

        # round and clamp for use as index
        proj_x = np.floor(proj_x)
        proj_x = np.minimum(self.proj_W - 1, proj_x)
        proj_x = np.maximum(0, proj_x).astype(np.int32)   # in [0,W-1]
        self.proj_x = np.copy(proj_x)  # store a copy in orig order

        proj_y = np.floor(proj_y)
        proj_y = np.minimum(self.proj_H - 1, proj_y)
        proj_y = np.maximum(0, proj_y).astype(np.int32)   # in [0,H-1]
        self.proj_y = np.copy(proj_y)  # stope a copy in original order

        # copy of depth in original order
        self.unproj_range = np.copy(depth)

        # order in decreasing depth
        indices = np.arange(depth.shape[0])
        order = np.argsort(depth)[::-1]
        depth = depth[order]
        indices = indices[order]
        points = self.points[order]
        remission = self.remissions[order]
        proj_y = proj_y[order]
        proj_x = proj_x[order]

        # assing to images
        self.proj_range[proj_y, proj_x] = depth
        self.proj_xyz[proj_y, proj_x] = points
        self.proj_remission[proj_y, proj_x] = remission
        self.proj_idx[proj_y, proj_x] = indices
        self.proj_mask = (self.proj_idx > 0).astype(np.float32)


class SemLaserScan(LaserScan):
    """Class that contains LaserScan with x,y,z,r,sem_label,sem_color_label,inst_label,inst_color_label"""
    EXTENSIONS_LABEL = ['.label']

    def __init__(self, nclasses, sem_color_dict=None, project=False, H=64, W=1024, fov_up=3.0, fov_down=-25.0):
        super(SemLaserScan, self).__init__(project, H, W, fov_up, fov_down)
        self.reset()
        self.nclasses = nclasses    # number of classes

        # make semantic colors
        max_sem_key = 0
        for key, data in sem_color_dict.items():
            if key + 1 > max_sem_key:
                max_sem_key = key + 1
        self.sem_color_lut = np.zeros((max_sem_key + 100, 3), dtype=np.float32)
        self.gt_sem_color_lut = np.zeros((max_sem_key + 100, 3), dtype=np.float32)
        for key, value in sem_color_dict.items():
            self.sem_color_lut[key] = np.array(value, np.float32) / 255.0
            self.gt_sem_color_lut[key] = np.array(value, np.float32) / 255.0

        # make instance colors
        max_inst_id = 100000
        self.inst_color_lut = np.random.uniform(low=0.0,
                                                high=1.0,
                                                size=(max_inst_id, 3))
        self.gt_inst_color_lut = np.random.uniform(low=0.0,
                                                   high=1.0,
                                                   size=(max_inst_id, 3))
        # force zero to a gray-ish color
        self.inst_color_lut[0] = np.full((3), 0.1)

    def reset(self):
        """ Reset scan members. """
        super(SemLaserScan, self).reset()

        # semantic labels
        self.sem_label = np.zeros((0, 1), dtype=np.uint32)         # [m, 1]: label
        self.sem_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        self.gt_sem_label = np.zeros((0, 1), dtype=np.uint32)  # [m, 1]: label
        self.gt_sem_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        # instance labels
        self.inst_label = np.zeros((0, 1), dtype=np.uint32)         # [m, 1]: label
        self.inst_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        self.gt_inst_label = np.zeros((0, 1), dtype=np.uint32)  # [m, 1]: label
        self.gt_inst_label_color = np.zeros((0, 3), dtype=np.float32)  # [m ,3]: color

        # projection color with semantic labels
        self.proj_sem_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)    # [H,W]  label
        self.proj_sem_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float)    # [H,W,3] color

        self.gt_proj_sem_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W]  label
        self.gt_proj_sem_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float)  # [H,W,3] color

        # projection color with instance labels
        self.proj_inst_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32) # [H,W]  label
        self.proj_inst_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float) # [H,W,3] color

        self.gt_proj_inst_label = np.zeros((self.proj_H, self.proj_W), dtype=np.int32)  # [H,W]  label
        self.gt_proj_inst_color = np.zeros((self.proj_H, self.proj_W, 3), dtype=np.float)  # [H,W,3] color

    def open_label(self, filename, gt=False):
        """ Open raw scan and fill in attributes
        """
        # check filename is string
        if not isinstance(filename, str):
            raise TypeError("Filename should be string type, "
                            "but was {type}".format(type=str(type(filename))))

        # check extension is a laserscan
        if not any(filename.endswith(ext) for ext in self.EXTENSIONS_LABEL):
            raise RuntimeError("Filename extension is not valid label file.")

        # if all goes well, open label
        label = np.fromfile(filename, dtype=np.uint32)
        label = label.reshape((-1))

        # set it
        self.set_label(label, gt=gt)

    def set_label(self, label, gt=False):
        """ Set points for label not from file but from np
        """
        # check label makes sense
        if not isinstance(label, np.ndarray):
            raise TypeError("Label should be numpy array")

        # only fill in attribute if the right size
        if label.shape[0] == self.points.shape[0]:
            if gt:
                self.gt_sem_label = label & 0xFFFF
                self.gt_inst_label = label >> 16
                # sanity check
                assert ((self.gt_sem_label + (self.gt_inst_label << 16) == label).all())
            else:
                self.sem_label = label & 0xFFFF  # semantic label in lower half
                self.inst_label = label >> 16    # instance id in upper half
                # sanity check
                assert ((self.sem_label + (self.inst_label << 16) == label).all())
        else:
            print("Points shape: ", self.points.shape)
            print("Label shape: ", label.shape)
            raise ValueError("Scan and Label don't contain same number of points")

        if self.project:
            self.do_label_projection(gt=gt)

    def colorize(self, gt=False):
        """ Colorize pointcloud with the color of each semantic label
        """
        if gt:
            self.gt_sem_label_color = self.gt_sem_color_lut[self.gt_sem_label]
            self.gt_sem_label_color = self.gt_sem_label_color.reshape((-1, 3))

            self.gt_inst_label_color = self.inst_color_lut[self.gt_inst_label]
            self.gt_inst_label_color = self.gt_inst_label_color.reshape((-1, 3))
        else:
            self.sem_label_color = self.sem_color_lut[self.sem_label]
            self.sem_label_color = self.sem_label_color.reshape((-1, 3))

            self.inst_label_color = self.inst_color_lut[self.inst_label]
            self.inst_label_color = self.inst_label_color.reshape((-1, 3))

    def do_label_projection(self, gt=False):
        # only map colors to labels that exist
        mask = self.proj_idx >= 0

        if gt:
            # semantics
            self.gt_proj_sem_label[mask] = self.gt_sem_label[self.proj_idx[mask]]
            self.gt_proj_sem_color[mask] = self.gt_sem_color_lut[self.gt_sem_label[self.proj_idx[mask]]]

            # instances
            self.gt_proj_inst_label[mask] = self.gt_inst_label[self.proj_idx[mask]]
            self.gt_proj_inst_color[mask] = self.gt_inst_color_lut[self.gt_inst_label[self.proj_idx[mask]]]
        else:
            # semantics
            self.proj_sem_label[mask] = self.sem_label[self.proj_idx[mask]]
            self.proj_sem_color[mask] = self.sem_color_lut[self.sem_label[self.proj_idx[mask]]]

            # instances
            self.proj_inst_label[mask] = self.inst_label[self.proj_idx[mask]]
            self.proj_inst_color[mask] = self.inst_color_lut[self.inst_label[self.proj_idx[mask]]]

utils/auxiliary/laserscan.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import vispy
from vispy.scene import visuals, SceneCanvas
import numpy as np
from matplotlib import pyplot as plt
from auxiliary.laserscan import LaserScan, SemLaserScan


class LaserScanVis:
    """Class that creates and handles a visualizer for a pointcloud"""

    def __init__(self, H, W, scan, scan_names, gt_label_names, pred_label_names, offset=0,
                 semantics=True, instances=False):
        self.H = H
        self.W = W
        self.scan = scan
        self.scan_names = scan_names
        self.gt_label_names = gt_label_names
        self.pred_label_names = pred_label_names
        self.offset = offset
        self.total = len(self.scan_names)
        self.semantics = semantics
        self.instances = instances
        # sanity check
        if not self.semantics and self.instances:
            print("Instances are only allowed in when semantics=True")
            raise ValueError

        self.reset()
        self.update_scan()

    def reset(self):
        """ Reset. """
        # last key press (it should have a mutex, but visualization is not
        # safety critical, so let's do things wrong)
        self.action = "no"  # no, next, back, quit are the possibilities

        # new canvas prepared for visualizing data
        self.canvas = SceneCanvas(keys='interactive', show=True)
        # interface (n next, b back, q quit, very simple)
        self.canvas.events.key_press.connect(self.key_press)
        self.canvas.events.draw.connect(self.draw)
        # grid
        self.grid = self.canvas.central_widget.add_grid()

        # add semantics
        if self.semantics:
            print("Using semantics in visualizer")
            self.gt_sem_view = vispy.scene.widgets.ViewBox(
                border_color='white', parent=self.canvas.scene)
            self.grid.add_widget(self.gt_sem_view, 0, 0)
            self.gt_sem_vis = visuals.Markers()
            self.gt_sem_view.camera = 'turntable'
            self.gt_sem_view.add(self.gt_sem_vis)
            visuals.XYZAxis(parent=self.gt_sem_view.scene)
            # self.sem_view.camera.link(self.scan_view.camera)

            if self.pred_label_names is not None:
                print("Using semantics in visualizer")
                self.sem_view = vispy.scene.widgets.ViewBox(
                    border_color='white', parent=self.canvas.scene)
                self.grid.add_widget(self.sem_view, 0, 1)
                self.sem_vis = visuals.Markers()
                self.sem_view.camera = 'turntable'
                self.sem_view.add(self.sem_vis)
                visuals.XYZAxis(parent=self.sem_view.scene)
                # self.sem_view.camera.link(self.scan_view.camera)

        if self.instances:
            print("Using instances in visualizer")
            self.gt_inst_view = vispy.scene.widgets.ViewBox(
                    border_color='white', parent=self.canvas.scene)
            self.grid.add_widget(self.gt_inst_view, 0, 2)
            self.gt_inst_vis = visuals.Markers()
            self.gt_inst_view.camera = 'turntable'
            self.gt_inst_view.add(self.gt_inst_vis)
            visuals.XYZAxis(parent=self.gt_inst_view.scene)
            # self.inst_view.camera.link(self.scan_view.camera)
            self.gt_inst_view.camera.link(self.sem_view.camera)

        # img canvas size
        self.multiplier = 1
        self.canvas_W = self.W
        self.canvas_H = self.H
        if self.semantics:
            self.multiplier += 1
            if self.pred_label_names is not None: self.multiplier += 1
        if self.instances:
            self.multiplier += 1

        # new canvas for img
        self.img_canvas = SceneCanvas(keys='interactive', show=True,
                                      size=(self.canvas_W, self.canvas_H * self.multiplier))
        # grid
        self.img_grid = self.img_canvas.central_widget.add_grid()
        # interface (n next, b back, q quit, very simple)
        self.img_canvas.events.key_press.connect(self.key_press)
        self.img_canvas.events.draw.connect(self.draw)

        # add a view for the depth
        self.img_view = vispy.scene.widgets.ViewBox(
            border_color='white', parent=self.img_canvas.scene)
        self.img_grid.add_widget(self.img_view, 0, 0)
        self.img_vis = visuals.Image(cmap='viridis')
        self.img_view.add(self.img_vis)

        # add semantics
        if self.semantics:
            self.gt_sem_img_view = vispy.scene.widgets.ViewBox(
                border_color='white', parent=self.img_canvas.scene)
            self.img_grid.add_widget(self.gt_sem_img_view, 1, 0)
            self.gt_sem_img_vis = visuals.Image(cmap='viridis')
            self.gt_sem_img_view.add(self.gt_sem_img_vis)

            if self.pred_label_names is not None:
                self.sem_img_view = vispy.scene.widgets.ViewBox(
                    border_color='white', parent=self.img_canvas.scene)
                self.img_grid.add_widget(self.sem_img_view, 2, 0)
                self.sem_img_vis = visuals.Image(cmap='viridis')
                self.sem_img_view.add(self.sem_img_vis)

        # add instances
        if self.instances:
            self.gt_inst_img_view = vispy.scene.widgets.ViewBox(
                border_color='white', parent=self.img_canvas.scene)
            self.img_grid.add_widget(self.gt_inst_img_view, 3, 0)
            self.gt_inst_img_vis = visuals.Image(cmap='viridis')
            self.gt_inst_img_view.add(self.gt_inst_img_vis)

    def get_mpl_colormap(self, cmap_name):
        cmap = plt.get_cmap(cmap_name)

        # Initialize the matplotlib color map
        sm = plt.cm.ScalarMappable(cmap=cmap)

        # Obtain linear color range
        color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1]

        return color_range.reshape(256, 3).astype(np.float32) / 255.0

    def update_scan(self):
        # first open data
        self.scan.open_scan(self.scan_names[self.offset])
        if self.semantics:
            self.scan.open_label(self.gt_label_names[self.offset], gt=True)
            self.scan.colorize(gt=True)

            if self.pred_label_names is not None:
                self.scan.open_label(self.pred_label_names[self.offset], gt=False)
                self.scan.colorize(gt=False)

        # then change names
        title = "scan " + str(self.offset)
        self.canvas.title = title
        self.img_canvas.title = title

        # then do all the point cloud stuff

        # plot scan
        power = 16
        # print()
        range_data = np.copy(self.scan.unproj_range)
        # print(range_data.max(), range_data.min())
        range_data = range_data**(1 / power)
        # print(range_data.max(), range_data.min())
        viridis_range = ((range_data - range_data.min()) /
                         (range_data.max() - range_data.min()) *
                         255).astype(np.uint8)
        viridis_map = self.get_mpl_colormap("viridis")
        viridis_colors = viridis_map[viridis_range]
        # self.scan_vis.set_data(self.scan.points,
        #                        face_color=viridis_colors[..., ::-1],
        #                        edge_color=viridis_colors[..., ::-1],
        #                        size=1)

        # plot semantics
        if self.semantics:
            self.gt_sem_vis.set_data(self.scan.points,
                                     face_color=self.scan.gt_sem_label_color[..., ::-1],
                                     edge_color=self.scan.gt_sem_label_color[..., ::-1],
                                     size=1)
            if self.pred_label_names is not None:
                self.sem_vis.set_data(self.scan.points,
                                      face_color=self.scan.sem_label_color[..., ::-1],
                                      edge_color=self.scan.sem_label_color[..., ::-1],
                                      size=1)


        # plot instances
        if self.instances:
            self.gt_inst_vis.set_data(self.scan.points,
                                      face_color=self.scan.gt_inst_label_color[..., ::-1],
                                      edge_color=self.scan.gt_inst_label_color[..., ::-1],
                                      size=1)

        # now do all the range image stuff
        # plot range image
        data = np.copy(self.scan.proj_range)
        # print(data[data > 0].max(), data[data > 0].min())
        data[data > 0] = data[data > 0]**(1 / power)
        data[data < 0] = data[data > 0].min()
        # print(data.max(), data.min())
        data = (data - data[data > 0].min()) / \
            (data.max() - data[data > 0].min())
        # print(data.max(), data.min())
        self.img_vis.set_data(data)
        self.img_vis.update()

        if self.semantics:
            self.gt_sem_img_vis.set_data(self.scan.gt_proj_sem_color[..., ::-1])
            self.gt_sem_img_vis.update()

            if self.pred_label_names is not None:
                self.sem_img_vis.set_data(self.scan.proj_sem_color[..., ::-1])
                self.sem_img_vis.update()

        if self.instances:
            self.gt_inst_img_vis.set_data(self.scan.gt_proj_inst_color[..., ::-1])
            self.gt_inst_img_vis.update()

    # interface
    def key_press(self, event):
        self.canvas.events.key_press.block()
        self.img_canvas.events.key_press.block()
        if event.key == 'N':
            self.offset += 1
            if self.offset >= self.total:
                self.offset = 0
            self.update_scan()
        elif event.key == 'B':
            self.offset -= 1
            if self.offset < 0:
                self.offset = self.total - 1
            self.update_scan()
        elif event.key == 'Q' or event.key == 'Escape':
            self.destroy()

    def draw(self, event):
        if self.canvas.events.key_press.blocked():
            self.canvas.events.key_press.unblock()
        if self.img_canvas.events.key_press.blocked():
            self.img_canvas.events.key_press.unblock()

    def destroy(self):
        # destroy the visualization
        self.canvas.close()
        self.img_canvas.close()
        vispy.app.quit()

    def run(self):
        vispy.app.run()

utils/auxiliary/laserscanvis.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import sys
import numpy as np


class iouEval:
    def __init__(self, n_classes, ignore=None):
        # classes
        self.n_classes = n_classes

        # What to include and ignore from the means
        self.ignore = np.array(ignore, dtype=np.int64)
        self.include = np.array([n for n in range(self.n_classes) if n not in self.ignore], dtype=np.int64)
        print("[IOU EVAL] IGNORE: ", self.ignore)
        print("[IOU EVAL] INCLUDE: ", self.include)

        # reset the class counters
        self.reset()

    def num_classes(self):
        return self.n_classes

    def reset(self):
        self.conf_matrix = np.zeros((self.n_classes, self.n_classes), dtype=np.int64)

    def addBatch(self, x, y):  # x=preds, y=targets
        # sizes should be matching
        x_row = x.reshape(-1)  # de-batchify
        y_row = y.reshape(-1)  # de-batchify

        # check
        assert(x_row.shape == x_row.shape)

        # create indexes
        idxs = tuple(np.stack((x_row, y_row), axis=0))

        # make confusion matrix (cols = gt, rows = pred)
        np.add.at(self.conf_matrix, idxs, 1)

    def getStats(self):
        # remove fp from confusion on the ignore classes cols
        conf = self.conf_matrix.copy()
        conf[:, self.ignore] = 0

        # get the clean stats
        tp = np.diag(conf)
        fp = conf.sum(axis=1) - tp
        fn = conf.sum(axis=0) - tp
        return tp, fp, fn

    def getIoU(self):
        tp, fp, fn = self.getStats()
        intersection = tp
        union = tp + fp + fn + 1e-15
        iou = intersection / union
        iou_mean = (intersection[self.include] / union[self.include]).mean()
        return iou_mean, iou  # returns "iou mean", "iou per class" ALL CLASSES

    def getacc(self):
        tp, fp, fn = self.getStats()
        total_tp = tp.sum()
        total = tp[self.include].sum() + fp[self.include].sum() + 1e-15
        acc_mean = total_tp / total
        return acc_mean  # returns "acc mean"
        
    def get_confusion(self):
        return self.conf_matrix.copy()



if __name__ == "__main__":
    # mock problem
    nclasses = 2
    ignore = []

    # test with 2 squares and a known IOU
    lbl = np.zeros((7, 7), dtype=np.int64)
    argmax = np.zeros((7, 7), dtype=np.int64)
    
    eval = iouEval(nclasses, ignore)

    eval.addBatch(argmax, lbl)

    # put squares
    lbl[2:4, 2:4] = 1
    argmax[3:5, 3:5] = 1

    # make evaluator

    # run
    eval.addBatch(argmax, lbl)

    m_iou, iou = eval.getIoU()
    print("IoU: ", m_iou)
    print("IoU class: ", iou)
    m_acc = eval.getacc()
    print("Acc: ", m_acc)

utils/auxiliary/np_ioueval.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# Developed by Xieyuanli Chen
# Brief: This script combines moving object segmentation with semantic information

import os
import sys
import yaml
import numpy as np
from tqdm import tqdm

from kitti_utils import load_files, load_labels


if __name__ == '__main__':
	# load config file
	config_filename = 'config/post-processing.yaml'
	if len(sys.argv) > 1:
		config_filename = sys.argv[1]

	if yaml.__version__ >= '5.1':
		config = yaml.load(open(config_filename), Loader=yaml.FullLoader)
	else:
		config = yaml.load(open(config_filename))

	# specify moving object segmentation results folder
	mos_pred_root = config['mos_pred_root']

	# specify semantic segmentation results folder
	semantic_pred_root = config['semantic_pred_root']

	# create a new folder for combined results
	combined_results_root = config['combined_results_root']

	# specify the split
	split = config['split']
	data_yaml = yaml.load(open('config/semantic-kitti-mos.yaml'))

	# create output folder
	seqs = []
	if not os.path.exists(os.path.join(combined_results_root, "sequences")):
		os.makedirs(os.path.join(combined_results_root, "sequences"))

	if split == 'train':
		print(data_yaml["split"]["train"])
		for seq in data_yaml["split"]["train"]:
			seq = '{0:02d}'.format(int(seq))
			print("train", seq)
			# if not os.path.exists(os.path.join(combined_results_root, "sequences", seq, "clean_scans")):
			#   os.makedirs(os.path.join(combined_results_root, "sequences", seq, "clean_scans"))
			if not os.path.exists(os.path.join(combined_results_root, "sequences", seq, "predictions")):
				os.makedirs(os.path.join(combined_results_root, "sequences", seq, "predictions"))
			seqs.append(seq)
	if split == 'valid':
		for seq in data_yaml["split"]["valid"]:
			seq = '{0:02d}'.format(int(seq))
			print("valid", seq)
			if not os.path.exists(os.path.join(combined_results_root, "sequences", seq, "clean_scans")):
				os.makedirs(os.path.join(combined_results_root, "sequences", seq, "clean_scans"))
			seqs.append(seq)
	if split == 'test':
		for seq in data_yaml["split"]["test"]:
			seq = '{0:02d}'.format(int(seq))
			print("test", seq)
			# if not os.path.exists(os.path.join(combined_results_root, "sequences", seq, "clean_scans")):
			#   os.makedirs(os.path.join(combined_results_root, "sequences", seq, "clean_scans"))
			if not os.path.exists(os.path.join(combined_results_root, "sequences", seq, "predictions")):
				os.makedirs(os.path.join(combined_results_root, "sequences", seq, "predictions"))
			seqs.append(seq)

	for seq in seqs:
		# load moving object segmentation files
		mos_pred_seq_path = os.path.join(mos_pred_root, "sequences", seq, "predictions")
		mos_pred_files = load_files(mos_pred_seq_path)

		# load semantic segmentation files
		semantic_pred_seq_path = os.path.join(semantic_pred_root, "sequences", seq, "predictions")
		semantic_pred_files = load_files(semantic_pred_seq_path)

		print('processing seq:', seq)

		for frame_idx in tqdm(range(len(mos_pred_files))):
			mos_pred, _ = load_labels(mos_pred_files[frame_idx])  # mos_pred should be 9/251 for static/dynamic
			semantic_pred, _ = load_labels(semantic_pred_files[frame_idx])  # mos_pred should be full classes
			semantic_pred_mapped = np.ones(len(mos_pred), dtype=np.uint32) * 9
			combine_pred = np.ones(len(mos_pred), dtype=np.uint32) * 9

			# mapping semantic into static and movable classes
			movable_mask = (semantic_pred > 0) & (semantic_pred < 40)
			semantic_pred_mapped[movable_mask] = 251

			# if consistent keep the same, otherwise labeled as static
			combined_mask = (semantic_pred_mapped == mos_pred)
			combine_pred[combined_mask] = mos_pred[combined_mask]

			file_name = os.path.join(combined_results_root, "sequences", seq, "predictions", str(frame_idx).zfill(6))
			combine_pred.reshape((-1)).astype(np.uint32)
			# print(file_name)
			combine_pred.tofile(file_name + '.label')
utils/combine_semantics.py
SCNU-RISLAB/MF-MOS
import numpy as np
import cv2
import os

if __name__ == "__main__":
    residual_root_path = "/data1/MotionSeg3d_Dataset/data_odometry_velodyne/dataset/sequences/08"
    residual_id = [1, 2, 3, 4, 5, 6, 7, 8]
    # residual_id = [1, 3, 5, 7, 9, 11, 13, 15]
    # residual_id = [1, 4, 7, 10, 13, 16, 19, 22]

    residual_image_name = "267".zfill(6)
    concat_image_save_path = os.path.join(residual_root_path, f"concat_{residual_image_name}.png")

    for per_residual_id in residual_id:
        residual_np = np.load(os.path.join(residual_root_path, f"residual_images_{per_residual_id}/{residual_image_name}.npy"))
        zero_index = np.where(residual_np == 0)
        residual_image = cv2.imread(os.path.join(residual_root_path, f"visualization_{per_residual_id}/{residual_image_name}.png"))
        if per_residual_id == residual_id[0]:
            all_resdual_image = residual_image
        else:
            all_resdual_image = np.vstack((all_resdual_image, np.zeros((10, residual_image.shape[1], 3)), residual_image))

    cv2.imwrite(concat_image_save_path, all_resdual_image)
utils/concat_residual_image.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# Developed by: Xieyuanli Chen
# Updated by Jiadai Sun with new features:
#   Allows to calculate the IoU of points within a specified radius R

import argparse
import os
import yaml
import sys
import numpy as np
from tqdm import tqdm

# possible splits
splits = ["train", "valid", "test"]

# possible backends
backends = ["numpy", "torch"]


def get_args():
    parser = argparse.ArgumentParser("./evaluate_mos.py")
    parser.add_argument(
        '--dataset', '-d',
        type=str,
        required=True,
        help='Dataset dir. No Default',
    )
    parser.add_argument(
        '--predictions', '-p',
        type=str,
        required=None,
        help='Prediction dir. Same organization as dataset, but predictions in'
        'each sequences "prediction" directory. No Default. If no option is set'
        ' we look for the labels in the same directory as dataset'
    )
    parser.add_argument(
        '--split', '-s',
        type=str,
        required=False,
        choices=["train", "valid", "test"],
        default="valid",
        help='Split to evaluate on. One of ' +
        str(splits) + '. Defaults to %(default)s',
    )
    parser.add_argument(
        '--backend', '-b',
        type=str,
        required=False,
        choices=["numpy", "torch"],
        default="numpy",
        help='Backend for evaluation. One of ' +
        str(backends) + ' Defaults to %(default)s',
    )
    parser.add_argument(
        '--datacfg', '-dc',
        type=str,
        required=False,
        default="config/semantic-kitti-mos.yaml",
        help='Dataset config file. Defaults to %(default)s',
    )
    parser.add_argument(
        '--limit', '-l',
        type=int,
        required=False,
        default=None,
        help='Limit to the first "--limit" points of each scan. Useful for'
        ' evaluating single scan from aggregated pointcloud.'
        ' Defaults to %(default)s',
    )
    parser.add_argument(
        '--codalab',
        dest='codalab',
        type=str,
        default=None,
        help='Exports "scores.txt" to given output directory for codalab'
        'Defaults to %(default)s',
    )
    parser.add_argument(
        '--radius', '-r',
        type=int,
        required=False,
        default=-1,
        help='Calculated range radius, -1 means using all points.'
    )
    return parser


if __name__ == '__main__':
    
    parser = get_args()
    FLAGS, unparsed = parser.parse_known_args()
    
    assert(FLAGS.split in splits)     # assert split
    assert(FLAGS.backend in backends) # assert backend

    # fill in real predictions dir
    if FLAGS.predictions is None:
        FLAGS.predictions = FLAGS.dataset

    # print summary of what we will do
    print("*" * 80)
    print("  INTERFACE:")
    print("  Data: ", FLAGS.dataset)
    print("  Predictions: ", FLAGS.predictions)
    print("  Backend: ", FLAGS.backend)
    print("  Split: ", FLAGS.split)
    print("  Config: ", FLAGS.datacfg)
    print("  Limit: ", FLAGS.limit)
    print("  Codalab: ", FLAGS.codalab)
    print("  Radius: ", FLAGS.radius)
    print("*" * 80)

    print(f"Opening data config file {FLAGS.datacfg}")
    DATA = yaml.safe_load(open(FLAGS.datacfg, 'r'))

    # get number of interest classes, and the label mappings
    class_strings = DATA["labels"]
    class_remap = DATA["moving_learning_map"]
    class_inv_remap = DATA["moving_learning_map_inv"]
    class_ignore = DATA["learning_ignore"]
    nr_classes = len(class_inv_remap)

    # make lookup table for mapping
    maxkey = max(class_remap.keys())
    
    # +100 hack making lut bigger just in case there are unknown labels
    remap_lut = np.zeros((maxkey + 100), dtype=np.int32)
    remap_lut[list(class_remap.keys())] = list(class_remap.values())
    # print(remap_lut)

    # create evaluator
    ignore = []
    for cl, ign in class_ignore.items():
        if ign:
            x_cl = int(cl)
            ignore.append(x_cl)
            print("Ignoring xentropy class ", x_cl, " in IoU evaluation")

    # create evaluator
    if FLAGS.backend == "torch":
        from auxiliary.torch_ioueval import iouEval
        evaluator = iouEval(nr_classes, ignore)
        frame_evaluator = iouEval(nr_classes, ignore)
    elif FLAGS.backend == "numpy":
        from auxiliary.np_ioueval import iouEval
        evaluator = iouEval(nr_classes, ignore)
        frame_evaluator = iouEval(nr_classes, ignore)
    else:
        print("Backend for evaluator should be one of ", str(backends))
        quit()

    evaluator.reset()
    frame_evaluator.reset()
    # get test set
    test_sequences = DATA["split"][FLAGS.split]

    # get label paths
    label_names = []
    lidar_names = []
    for sequence in test_sequences:
        sequence = '{0:02d}'.format(int(sequence))
        label_paths = os.path.join(FLAGS.dataset, "sequences", str(sequence), "labels")
        # populate the label names
        seq_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
            os.path.expanduser(label_paths)) for f in fn if ".label" in f]
        seq_label_names.sort()
        label_names.extend(seq_label_names)

        if FLAGS.radius != -1:
            lidar_paths = os.path.join(FLAGS.dataset, "sequences", str(sequence), "velodyne")
            seq_lidar_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
                os.path.expanduser(lidar_paths)) for f in fn if ".bin" in f]
            seq_lidar_names.sort()
            assert len(seq_label_names) == len(seq_lidar_names)
            lidar_names.extend(seq_lidar_names)
    # print(label_names)

    # get predictions paths
    pred_names = []
    for sequence in test_sequences:
        sequence = '{0:02d}'.format(int(sequence))
        pred_paths = os.path.join(FLAGS.predictions, "sequences", sequence, "predictions")
        # populate the label names
        seq_pred_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
            os.path.expanduser(pred_paths)) for f in fn if ".label" in f]
        seq_pred_names.sort()
        pred_names.extend(seq_pred_names)
    # print(pred_names)

    # check that I have the same number of files

    print("labels: ", len(label_names))
    print("predictions: ", len(pred_names))
    assert(len(label_names) == len(pred_names))
    if FLAGS.radius != -1:
        print("lidars: ", len(lidar_names))
        print(f"\033[32m Only use the points in radius <= {FLAGS.radius}m. \033[0m")

    # open each file, get the tensor, and make the iou comparison
    # for lidar_file, label_file, pred_file in zip(lidar_names[:], label_names[:], pred_names[:]):
    for f_id in tqdm(range(len(label_names[:])), desc="Evaluating sequences:", ncols=80):
        label_file = label_names[f_id]
        pred_file = pred_names[f_id]
        frame_evaluator.reset()

        if FLAGS.radius != -1:
            pc_xyz = np.fromfile(lidar_names[f_id], dtype=np.float32).reshape((-1, 4))[:, :3]
            depth = np.linalg.norm(pc_xyz, 2, axis=1)
            
            radius_mask = np.ones((pc_xyz.shape[0]), dtype=bool)
            if FLAGS.radius > 0:
                radius_mask = np.logical_and(depth <= FLAGS.radius, depth >= 2)

        # open label
        label = np.fromfile(label_file, dtype=np.int32)
        label = label.reshape((-1))  # reshape to vector
        label = label & 0xFFFF       # get lower half for semantics
        if FLAGS.limit is not None:
            label = label[:FLAGS.limit]  # limit to desired length
        label = remap_lut[label]         # remap to xentropy format

        # open prediction
        pred = np.fromfile(pred_file, dtype=np.int32)
        pred = pred.reshape((-1))    # reshape to vector
        pred = pred & 0xFFFF         # get lower half for semantics
        if FLAGS.limit is not None:
            pred = pred[:FLAGS.limit]  # limit to desired length
        pred = remap_lut[pred]         # remap to xentropy format
        # add single scan to evaluation

        if FLAGS.radius != -1:
            pred = pred[radius_mask]
            label = label[radius_mask]

        evaluator.addBatch(pred, label) # shape: (n, ) (n, ), type: int32
        # m_jaccard, class_jaccard = evaluator.getIoU()
        # frame_evaluator.addBatch(pred, label)
        # m_jaccard, class_jaccard = frame_evaluator.getIoU()

        seq = label_file.split('/')[-3]
        ind = label_file.split('/')[-1]
        # print(f"count: {count} || {seq}/{ind} || m_jaccard: {m_jaccard}, class_jaccard: {class_jaccard}")

    m_accuracy = evaluator.getacc()
    m_jaccard, class_jaccard = evaluator.getIoU()

    # print for spreadsheet
    print("*" * 80)
    print("below can be copied straight for paper table")
    for i, jacc in enumerate(class_jaccard):
        if i not in ignore:
            if int(class_inv_remap[i]) == 9:
                sys.stdout.write('iou_static: {jacc:.6f}\n'.format(jacc=jacc.item()))
            if int(class_inv_remap[i]) > 250:
                sys.stdout.write('iou_moving: {jacc:.6f}'.format(jacc=jacc.item()))
    sys.stdout.write('\n')
    sys.stdout.flush()

    # if codalab is necessary, then do it
    # for moving object detection, we only care about the results of moving objects
    if FLAGS.codalab is not None:
        results = {}
        for i, jacc in enumerate(class_jaccard):
            if i not in ignore:
                if int(class_inv_remap[i]) > 250:
                    results["iou_moving"] = float(jacc)
        # save to file
        output_filename = os.path.join(FLAGS.codalab, 'scores.txt')
        with open(output_filename, 'w') as yaml_file:
            yaml.dump(results, yaml_file, default_flow_style=False)

utils/evaluate_mos.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# Developed by Xieyuanli Chen
# This file is covered by the LICENSE file in the root of this project.
# Brief: This script generates residual images

import os
os.environ["OMP_NUM_THREADS"] = "4"
import sys
import yaml
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

from kitti_utils import load_poses, load_calib, load_files, load_vertex

try:
    from c_gen_virtual_scan import gen_virtual_scan as range_projection
except:
    print("Using clib by $export PYTHONPATH=$PYTHONPATH:<path-to-library>")
    print("Currently using python-lib to generate range images.")
    from kitti_utils import range_projection


if __name__ == '__main__':
    # load config file
    config_filename = 'config/data_preparing.yaml'
    if len(sys.argv) > 1:
        config_filename = sys.argv[1]
    
    if yaml.__version__ >= '5.1':
        config = yaml.load(open(config_filename), Loader=yaml.FullLoader)
    else:
        config = yaml.load(open(config_filename))
    
    # specify parameters
    num_frames = config['num_frames']
    debug = config['debug']
    normalize = config['normalize']
    num_last_n = config['num_last_n']
    visualize = config['visualize']
    visualization_folder = config['visualization_folder']
    
    # specify the output folders
    residual_image_folder = config['residual_image_folder']
    if not os.path.exists(residual_image_folder):
        os.makedirs(residual_image_folder)
        
    if visualize:
        if not os.path.exists(visualization_folder):
            os.makedirs(visualization_folder)
    
    # load poses
    pose_file = config['pose_file']
    poses = np.array(load_poses(pose_file))
    inv_frame0 = np.linalg.inv(poses[0])
    
    # load calibrations
    calib_file = config['calib_file']
    T_cam_velo = load_calib(calib_file)
    T_cam_velo = np.asarray(T_cam_velo).reshape((4, 4))
    T_velo_cam = np.linalg.inv(T_cam_velo)
    
    # convert kitti poses from camera coord to LiDAR coord
    new_poses = []
    for pose in poses:
        new_poses.append(T_velo_cam.dot(inv_frame0).dot(pose).dot(T_cam_velo))
    poses = np.array(new_poses)
    
    # load LiDAR scans
    scan_folder = config['scan_folder']
    scan_paths = load_files(scan_folder)
    
    # test for the first N scans
    if num_frames >= len(poses) or num_frames <= 0:
        print('generate training data for all frames with number of: ', len(poses))
    else:
        poses = poses[:num_frames]
        scan_paths = scan_paths[:num_frames]
    
    range_image_params = config['range_image']
    
    # generate residual images for the whole sequence
    for frame_idx in tqdm(range(len(scan_paths))):
        file_name = os.path.join(residual_image_folder, str(frame_idx).zfill(6))
        diff_image = np.full((range_image_params['height'], range_image_params['width']), 0,
                                                         dtype=np.float32)  # [H,W] range (0 is no data)
        
        # for the first N frame we generate a dummy file
        if frame_idx < num_last_n:
            np.save(file_name, diff_image)
        
        else:
            # load current scan and generate current range image
            current_pose = poses[frame_idx]
            current_scan = load_vertex(scan_paths[frame_idx])
            current_range = range_projection(current_scan.astype(np.float32),
                                             range_image_params['height'], range_image_params['width'],
                                             range_image_params['fov_up'], range_image_params['fov_down'],
                                             range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]
            
            # load last scan, transform into the current coord and generate a transformed last range image
            last_pose = poses[frame_idx - num_last_n]
            last_scan = load_vertex(scan_paths[frame_idx - num_last_n])
            last_scan_transformed = np.linalg.inv(current_pose).dot(last_pose).dot(last_scan.T).T
            last_range_transformed = range_projection(last_scan_transformed.astype(np.float32),
                                                      range_image_params['height'], range_image_params['width'],
                                                      range_image_params['fov_up'], range_image_params['fov_down'],
                                                      range_image_params['max_range'], range_image_params['min_range'])[:, :, 3]
            
            # generate residual image
            valid_mask = (current_range > range_image_params['min_range']) & \
                            (current_range < range_image_params['max_range']) & \
                            (last_range_transformed > range_image_params['min_range']) & \
                            (last_range_transformed < range_image_params['max_range'])
            difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask])
            
            if normalize:
                difference = np.abs(current_range[valid_mask] - last_range_transformed[valid_mask]) / current_range[valid_mask]

            diff_image[valid_mask] = difference
            
            if debug:
                fig, axs = plt.subplots(3)
                axs[0].imshow(last_range_transformed)
                axs[1].imshow(current_range)
                axs[2].imshow(diff_image, vmin=0, vmax=1)
                plt.show()
                
            if visualize:
                fig = plt.figure(frameon=False, figsize=(16, 10))
                fig.set_size_inches(20.48, 0.64)
                ax = plt.Axes(fig, [0., 0., 1., 1.])
                ax.set_axis_off()
                fig.add_axes(ax)
                ax.imshow(diff_image, vmin=0, vmax=1)
                image_name = os.path.join(visualization_folder, str(frame_idx).zfill(6))
                plt.savefig(image_name)
                plt.close()

            # save residual image
            np.save(file_name, diff_image)

utils/gen_residual_images.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# Developed by Jiadai Sun

import os
import sys
import glob
import yaml
import numpy as np

from tqdm import tqdm
from glob import glob
from icecream import ic


def get_frame_data(pc_path, label_path):
    pc_data = np.fromfile(pc_path, dtype=np.float32).reshape((-1, 4))
    label = np.fromfile(label_path, dtype=np.uint32).reshape((-1))
    sem_label = label & 0xFFFF
    ins_label = label >> 16

    return pc_data, sem_label, ins_label


class SeqKITTI():
    def __init__(self, dataset_path, split, data_cfg_path):
        self.dataset_path = dataset_path
        self.split = split  # valid train test
        self.data_yaml = yaml.load(open(data_cfg_path))

        print(split, self.data_yaml["split"][split])
        self.seqs = []
        for seq in self.data_yaml["split"][split]:
            seq = '{0:02d}'.format(int(seq))
            # print(split, seq)
            self.seqs.append(seq)

    def get_file_list(self, seq_id):
        velodyne_seq_path = os.path.join(self.dataset_path, "sequences", seq_id, "velodyne")
        velodyne_seq_files = sorted(glob.glob(os.path.join(velodyne_seq_path, "*.bin")))

        # load gt semantic segmentation files
        gtsemantic_seq_path = os.path.join(self.dataset_path, "sequences", seq_id, "labels")
        gtsemantic_seq_files = sorted(glob.glob(os.path.join(gtsemantic_seq_path, "*.label")))

        assert len(velodyne_seq_files) == len(gtsemantic_seq_files)

        return velodyne_seq_files, gtsemantic_seq_files

    def count_dynamic_frames(self, write_to_txt=False):

        if write_to_txt:
            fo = open("KITTI_train_split_dynamic_pointnumber.txt", "w")

        self.moving_threshold_num_points = 100

        for seq in self.seqs:

            velodyne_seq_files, gtsemantic_seq_files = self.get_file_list(
                seq_id=seq)
            num_moving_frames = 0

            for frame_idx in range(len(velodyne_seq_files)):

                f1_xyzi, f1_semlabel, f1_inslabel = \
                    get_frame_data(
                        pc_path=velodyne_seq_files[frame_idx], label_path=gtsemantic_seq_files[frame_idx])

                f1_moving_label_mask = (f1_semlabel > 250)

                if f1_moving_label_mask.sum() > self.moving_threshold_num_points:
                    num_moving_frames += 1
                
                if write_to_txt:
                    linestr = f"{seq} " + "%06d"%frame_idx + f" {f1_moving_label_mask.sum()}\n"
                    fo.write(linestr)

            print(f"Seq {seq} | Moving frames / all == {num_moving_frames}/{len(velodyne_seq_files)} = {num_moving_frames / len(velodyne_seq_files)}")

        pass

    def count_seqs_points(self,):

        for seq in self.seqs:

            length_min = 1000000
            length_max = -1

            velodyne_seq_files, gtsemantic_seq_files = self.get_file_list(seq_id=seq)
            # assert len(velodyne_seq_files) == len(gtsemantic_seq_files)

            for frame_idx in tqdm(range(len(velodyne_seq_files))):
                f1_xyzi = np.fromfile(velodyne_seq_files[frame_idx], dtype=np.float32).reshape((-1, 4))

                if f1_xyzi.shape[0] < length_min:
                    length_min = f1_xyzi.shape[0]
                if f1_xyzi.shape[0] > length_max:
                    length_max = f1_xyzi.shape[0]

            print(f"Seq {seq} | min: {length_min} / max: {length_max}")

    def count_moving_points_in_seqs(self,):

        for seq in self.seqs:

            length_min = 1000000
            length_max = -1
            # load point cloud files
            velodyne_seq_path = os.path.join(dataset_path, "sequences", seq, "velodyne")
            velodyne_seq_files = sorted(glob.glob(os.path.join(velodyne_seq_path, "*.bin")))

            velodyne_seq_files, gtsemantic_seq_files = self.get_file_list(seq_id=seq)
            # assert len(velodyne_seq_files) == len(gtsemantic_seq_files)

            num_moving_frames = 0
            for frame_idx in tqdm(range(len(velodyne_seq_files))):

                f1_xyzi, f1_semlabel, f1_inslabel = \
                    get_frame_data(pc_path=velodyne_seq_files[frame_idx], label_path=gtsemantic_seq_files[frame_idx])

                # mapping rae semantic labels to LiDAR-MOS labels
                f1_moving_label_mask = (f1_semlabel > 250)
                f1_semlabel[f1_moving_label_mask] = 251
                f1_semlabel[~f1_moving_label_mask] = 9

                a, b = np.unique(f1_semlabel, return_counts=True)
                print(a, b)

            print(f"Seq {seq} | min: {length_min} / max: {length_max}")


if __name__ == '__main__':

    dataset_path = '/home1/datasets/semantic_kitti/dataset'
    split = 'train'  # 'valid'
    data_cfg_path = 'config/labels/semantic-kitti-mos.yaml'
    seqKITTI = SeqKITTI(dataset_path, split, data_cfg_path)

    seqKITTI.count_seqs_points()
    # seqKITTI.count_dynamic_frames()
    # seqKITTI.count_dynamic_frames(write_to_txt=True)
    # seqKITTI.count_moving_points_in_seqs()


utils/kitti_mos_statistical_analysis.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# Developed by: Xieyuanli Chen
# Brief: clean the LiDAR scans using our LiDAR-based moving object segmentation method

import os
import sys
import yaml
import numpy as np
from tqdm import tqdm

from kitti_utils import load_vertex, load_labels, load_files

if __name__ == '__main__':
	# load config file
	config_filename = 'config/post-processing.yaml'
	data_yaml_filename = 'config/labels/semantic-kitti-mos.raw.yaml'
	if len(sys.argv) > 1:
		config_filename = sys.argv[1]

	if yaml.__version__ >= '5.1':
		config = yaml.load(open(config_filename), Loader=yaml.FullLoader)
		data_yaml = yaml.load(open(data_yaml_filename), Loader=yaml.FullLoader)
	else:
		config = yaml.load(open(config_filename))
		data_yaml = yaml.load(open(data_yaml_filename))

	# raw clean folder root
	scan_root = config['scan_root']

	# specify moving object segmentation results folder
	mos_pred_root = config['mos_pred_root']

	# create a new folder for combined results
	clean_scan_root = config['clean_scan_root']

	# specify the split
	split = config['split']

	# create output folder
	seqs = []
	if not os.path.exists(os.path.join(clean_scan_root, "sequences")):
		os.makedirs(os.path.join(clean_scan_root, "sequences"))

	if split == 'train':
		for seq in data_yaml["split"]["train"]:
			seq = '{0:02d}'.format(int(seq))
			print("train", seq)
			if not os.path.exists(os.path.join(clean_scan_root, "sequences", seq, "clean_scans")):
				os.makedirs(os.path.join(clean_scan_root, "sequences", seq, "clean_scans"))
			seqs.append(seq)
	if split == 'valid':
		for seq in data_yaml["split"]["valid"]:
			seq = '{0:02d}'.format(int(seq))
			print("train", seq)
			if not os.path.exists(os.path.join(clean_scan_root, "sequences", seq, "clean_scans")):
				os.makedirs(os.path.join(clean_scan_root,"sequences", seq, "clean_scans"))
			seqs.append(seq)
	if split == 'test':
		for seq in data_yaml["split"]["test"]:
			seq = '{0:02d}'.format(int(seq))
			print("train", seq)
			if not os.path.exists(os.path.join(clean_scan_root, "sequences", seq, "clean_scans")):
				os.makedirs(os.path.join(clean_scan_root, "sequences", seq, "clean_scans"))
			seqs.append(seq)

	for seq in seqs:
		# load moving object segmentation files
		mos_pred_seq_path = os.path.join(mos_pred_root, "sequences", seq, "predictions")
		mos_pred_files = load_files(mos_pred_seq_path)

		# load semantic segmentation files
		raw_scan_path = os.path.join(scan_root, "sequences", seq, "velodyne")
		raw_scan_files = load_files(raw_scan_path)

		print('processing seq:', seq)

		for frame_idx in tqdm(range(len(mos_pred_files))):
			# mos_pred should be 9/251 for static/dynamic
			mos_pred, _ = load_labels(mos_pred_files[frame_idx])
			current_scan = load_vertex(raw_scan_files[frame_idx])

			clean_scan = current_scan[mos_pred < 250]
			np.array(clean_scan, dtype=np.float32).tofile(os.path.join(
				clean_scan_root, "sequences", seq, "clean_scans", str(frame_idx).zfill(6) + '.bin'))

utils/scan_cleaner.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
import math
import yaml
import shutil
import argparse

from decimal import Decimal
from datetime import datetime
from pip._vendor.distlib.compat import raw_input


def remove_exponent(d):
    return d.quantize(Decimal(1)) if d == d.to_integral() else d.normalize()


def millify(n, precision=0, drop_nulls=True, prefixes=[]):
    millnames = ['', 'k', 'M', 'B', 'T', 'P', 'E', 'Z', 'Y']
    if prefixes:
        millnames = ['']
        millnames.extend(prefixes)
    n = float(n)
    millidx = max(0, min(len(millnames) - 1,
                         int(math.floor(0 if n == 0 else math.log10(abs(n)) / 3))))
    result = '{:.{precision}f}'.format(n / 10**(3 * millidx), precision=precision)
    if drop_nulls:
        result = remove_exponent(Decimal(result))
    return '{0}{dx}'.format(result, dx=millnames[millidx])


def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean expected')


def load_yaml(path):
    try:
        print(f"\033[32m Opening arch config file {path}\033[0m")
        yaml_data = yaml.safe_load(open(path, 'r'))
        return yaml_data
    except Exception as e:
        print(e)
        print(f"Error opening {path} yaml file.")
        quit()


def check_and_makedirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


def check_pretrained_dir(path):
    if path is not None:
        if os.path.isdir(path):
            print("\033[32m model folder exists! Using model from %s \033[0m" % (path))
        else:
            print("\033[32m model folder doesnt exist! Start with random weights...\033[0m")
    else:
        print("\033[32m No pretrained directory found.\033[0m")


def check_model_dir(path):
    if path is not None:
        if os.path.isdir(path):
            print("\033[32m model folder exists! Using model from %s \033[0m" % (path))
        else:
            print("\033[32m model folder doesnt exist! Can't infer...\033[0m")
            quit()
    else:
        print("\033[32m No model directory found.\033[0m")


def make_logdir(FLAGS, resume_train=False):
    try:
        if resume_train:
            if FLAGS.pretrained == "":
                FLAGS.pretrained = None
                if os.path.isdir(FLAGS.log):
                    if os.listdir(FLAGS.log):
                        # answer = raw_input("Log Directory is not empty. Do you want to proceed? [y/n]  ")
                        # if answer == 'n':
                        #     quit()
                        # else:
                        #     shutil.rmtree(FLAGS.log)
                        shutil.rmtree(FLAGS.log)
                os.makedirs(FLAGS.log)
            else:
                FLAGS.log = FLAGS.pretrained
                print("Not creating new log file. Using pretrained directory")
        else:
            if os.path.isdir(FLAGS.log):
                if os.listdir(FLAGS.log):
                    # answer = raw_input("Log Directory is not empty. Do you want to proceed? [y/n]  ")
                    # if answer == 'n':
                    #     quit()
                    # else:
                    #     shutil.rmtree(FLAGS.log)
                    shutil.rmtree(FLAGS.log)
            os.makedirs(FLAGS.log)
    except Exception as e:
        print(e)
        print("Error creating log directory. Check permissions!")
        quit()


def backup_to_logdir(FLAGS, pretrain_model=False):
    # copy all files to log folder (to remember what we did, and make inference
    # easier). Also, standardize name to be able to open it later
    try:
        print("Copying files to %s for further reference." % FLAGS.log)
        shutil.copyfile(FLAGS.arch_cfg, FLAGS.log + "/arch_cfg.yaml")
        shutil.copyfile(FLAGS.data_cfg, FLAGS.log + "/data_cfg.yaml")

        # Backup training code for later review
        code_backup_path = f"{FLAGS.log}/code"
        check_and_makedirs(code_backup_path)
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', '*.py'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'common'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'config'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'modules'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'script'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'train_yaml'))} {code_backup_path}")
        os.system(f"cp -r {os.path.abspath(os.path.join(os.path.dirname(__file__), '../', 'utils'))} {code_backup_path}")
        if pretrain_model:
            shutil.copyfile(FLAGS.pretrained + "/MFMOS_valid_best", FLAGS.log + "/MFMOS_valid_best")

    except Exception as e:
        print(e)
        print("Error copying files, check permissions. Exiting...")
        quit()


def make_predictions_dir(FLAGS, DATA, rm_old=False, save_movable=False):
    # create predictions file folder
    try:
        if rm_old:
            if os.path.isdir(FLAGS.log):
                shutil.rmtree(FLAGS.log)
            os.makedirs(FLAGS.log)
            os.makedirs(os.path.join(FLAGS.log, "sequences"))
        check_and_makedirs(os.path.join(FLAGS.log, "sequences"))
        
        for seq in DATA["split"][FLAGS.split]:
            seq = '{0:02d}'.format(int(seq))
            print(f"{FLAGS.split} : {seq}")
            check_and_makedirs(os.path.join(FLAGS.log, "sequences", seq, "predictions"))
            if save_movable:
                check_and_makedirs(os.path.join(FLAGS.log, "sequences", seq, "predictions_fuse"))
                check_and_makedirs(os.path.join(FLAGS.log, "sequences", seq, "predictions_movable"))

    except Exception as e:
        print(e)
        print("Error creating predictions directory. Check permissions!")
        raise

    pass


def get_args(flags=None):
    splits = ["train", "valid", "test"]

    if flags == "train":
        parser = argparse.ArgumentParser("./train.py")
    elif flags == "infer":
        parser = argparse.ArgumentParser("./infer.py")

    parser.add_argument(
        '--local_rank', default=0, type=int
    )

    parser.add_argument(
        '--dataset', '-d', type=str,
        required=False,
        default="DATAROOT",
        help='Dataset to train with. The parent directory of sequences. No Default.')
    parser.add_argument(
        '--log', '-l', type=str,
        default="./log/" + datetime.now().strftime("%Y-%-m-%d-%H:%M") + '/',
        help='Directory to put the log data. Default: ./log_default/date+time')
    parser.add_argument(
        '--name', '-n', type=str,
        default="",
        help='If you want to give an aditional discriptive name')
    # parser.add_argument(
    #     '--uncertainty', '-u',
    #     required=False,
    #     type=str2bool, nargs='?',
    #     const=True, default=True,
    #     help='Set this if you want to use the Uncertainty Version')
    
    if flags == "train":
        parser.add_argument(
            '--pretrained', '-p', type=str,
            required=False,
            default=None,
            help='Directory to get the pretrained model. If not passed, do from scratch!')
        parser.add_argument(
            '--arch_cfg', '-ac', type=str,
            required=False,
            default="",
            help='Architecture yaml cfg file. See /config/arch for sample. No default.')
        parser.add_argument(
            '--data_cfg', '-dc', type=str,
            required=False,
            default='',
            help='Classification yaml cfg file. See /config/labels for sample. No default.')

    if flags == "infer":
        # parser.add_argument(
        #     '--monte-carlo', '-c',
        #     type=int, default=30,
        #     help='Number of samplings per scan.')
        parser.add_argument(
            '--model', '-m',
            type=str,
            required=True,
            default=None,
            help='Directory to get the trained model.')
        parser.add_argument(
            '--split', '-s',
            type=str,
            required=False,
            default=None,
            help='Split to evaluate on. One of ' +
            str(splits) + '. Defaults to %(default)s',)
        parser.add_argument(
            '--pointrefine', '-prf',
            action='store_true',
            required=False,
            help='Whether to use the PointHead module to refine predictions')
        parser.add_argument(
            '--movable',
            action='store_true',
            required=False,
            help='Whether to save the label of movable objects')

    return parser

utils/utils.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# Developed by Xieyuanli Chen
# This file is covered by the LICENSE file in the root of this project.

import argparse
import os
import yaml
from auxiliary.laserscan import LaserScan, SemLaserScan
from auxiliary.laserscanvis import LaserScanVis
import copy

def get_args():
    parser = argparse.ArgumentParser("./visualize.py")
    parser.add_argument(
        '--dataset', '-d',
        type=str,
        required=True,
        help='Dataset to visualize. No Default',
    )
    parser.add_argument(
        '--config', '-c',
        type=str,
        required=False,
        default="config/labels/semantic-kitti-mos.yaml",
        help='Dataset config file. Defaults to %(default)s',
    )
    parser.add_argument(
        '--sequence', '-s',
        type=str,
        default="00",
        required=False,
        help='Sequence to visualize. Defaults to %(default)s',
    )
    parser.add_argument(
        '--predictions', '-p',
        type=str,
        default=None,
        required=False,
        help='Alternate location for labels, to use predictions folder. '
        'Must point to directory containing the predictions in the proper format '
        ' (see readme)'
        'Defaults to %(default)s',
    )
    parser.add_argument(
        '--ignore_semantics', '-i',
        dest='ignore_semantics',
        default=False,
        action='store_true',
        help='Ignore semantics. Visualizes uncolored pointclouds.'
        'Defaults to %(default)s',
    )
    parser.add_argument(
        '--do_instances', '-di',
        dest='do_instances',
        default=False,
        action='store_true',
        help='Visualize instances too. Defaults to %(default)s',
    )
    parser.add_argument(
        '--offset',
        type=int,
        default=0,
        required=False,
        help='Sequence to start. Defaults to %(default)s',
    )
    parser.add_argument(
        '--ignore_safety',
        dest='ignore_safety',
        default=False,
        action='store_true',
        help='Normally you want the number of labels and ptcls to be the same,'
        ', but if you are not done inferring this is not the case, so this disables'
        ' that safety.'
        'Defaults to %(default)s',
    )
    parser.add_argument(
        '--version', '-v',
        type=str,
        default="moving",
        choices=["moving", "movable", "fuse"],
        required=False,
        help="which version is selected to visualize",
    )
    parser.add_argument(
        '--gt_vis',
        default=False,
        action='store_true',
        help='Whether to visualize the ground truth',
    )
    return parser


if __name__ == '__main__':
    
    parser = get_args()
    FLAGS, unparsed = parser.parse_known_args()

    # print summary of what we will do
    print("*" * 80)
    print("  INTERFACE:")
    print("  Dataset:", FLAGS.dataset)
    print("  Config:", FLAGS.config)
    print("  Sequence:", FLAGS.sequence)
    print("  Predictions:", FLAGS.predictions)
    print("  ignore_semantics:", FLAGS.ignore_semantics)
    print("  do_instances:", FLAGS.do_instances)
    print("  ignore_safety:", FLAGS.ignore_safety)
    print("  offset:", FLAGS.offset)
    print("  visulization_version:", FLAGS.version)
    print("*" * 80)

    H, W = 64, 2048

    # open config file
    try:
        print("Opening config file %s" % FLAGS.config)
        CFG = yaml.safe_load(open(FLAGS.config, 'r'))
    except Exception as e:
        print(e)
        print("Error opening yaml file.")
        quit()

    # fix sequence name
    FLAGS.sequence = '{0:02d}'.format(int(FLAGS.sequence))

    # does sequence folder exist?
    scan_paths = os.path.join(FLAGS.dataset, "sequences", FLAGS.sequence, "velodyne")
    if os.path.isdir(scan_paths):
        print("Sequence folder exists! Using sequence from %s" % scan_paths)
    else:
        print("Sequence folder doesn't exist! Exiting...")
        quit()

    # populate the pointclouds
    scan_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(os.path.expanduser(scan_paths)) for f in fn]
    scan_names.sort()

    pred_label_names = None

    # does sequence folder exist?
    if not FLAGS.ignore_semantics:
        gt_label_paths = os.path.join(FLAGS.dataset, "sequences", FLAGS.sequence, "labels")

        if os.path.isdir(gt_label_paths):
            print("Labels folder exists! Using labels from %s" % gt_label_paths)
        else:
            print(gt_label_paths)
            print("Labels folder doesn't exist! Exiting...")
            quit()

        # populate the pointclouds
        gt_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
            os.path.expanduser(gt_label_paths)) for f in fn]
        gt_label_names.sort()

        # check that there are same amount of labels and scans
        if not FLAGS.ignore_safety:
            print(f"len(gt_label_names):{len(gt_label_names)}, len(scan_names)={len(scan_names)}")
            assert(len(gt_label_names) == len(scan_names))

        if FLAGS.predictions is not None:
            if FLAGS.version == "moving":
                pred_label_paths = os.path.join(FLAGS.predictions, "sequences", FLAGS.sequence, "predictions")
            elif FLAGS.version == "movable":
                pred_label_paths = os.path.join(FLAGS.predictions, "sequences", FLAGS.sequence, "predictions_movable")
            else:
                pred_label_paths = os.path.join(FLAGS.predictions, "sequences", FLAGS.sequence, "predictions_fuse")

            if os.path.isdir(pred_label_paths):
                print("Predictions labels folder exists! Using labels from %s" % pred_label_paths)
            else:
                raise FileNotFoundError("Predictions labels doesn't exist! Exiting...")
            pred_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
                os.path.expanduser(pred_label_paths)) for f in fn]
            pred_label_names.sort()

            # check that there are same amount of labels and scans
            if not FLAGS.ignore_safety:
                print(f"len(pred_label_names):{len(pred_label_names)}, len(scan_names)={len(scan_names)}")
                assert (len(pred_label_names) == len(scan_names))

        color_dict = CFG["color_map"]
        gt_color_dict = copy.deepcopy(color_dict)
        moving_learning_map = CFG["moving_learning_map"]
        movable_learning_map = CFG["movable_learning_map"]
        moving_learning_map_inv = CFG["moving_learning_map_inv"]
        movable_learning_map_inv = CFG["movable_learning_map_inv"]
        # # import pdb; pdb.set_trace()

        for key in color_dict.keys():
            if (key == 250) or (key in movable_learning_map.keys() and movable_learning_map[key] == 2):
                color_dict[key] = [255, 0, 0]
                if key != 250 and moving_learning_map[key] == 2:
                    color_dict[key] = [0, 0, 255]
            else:
                color_dict[key] = [255, 255, 255]
        nclasses = len(color_dict)
        print(color_dict)
        scan = SemLaserScan(nclasses, color_dict, H=H, W=W, project=True)
    else:
        gt_label_names = None
        scan = LaserScan(H=H, W=W, project=True)  # project all opened scans to spheric proj

    # create a visualizer
    vis = LaserScanVis(H=H, W=W,
                       scan=scan,
                       scan_names=scan_names,
                       gt_label_names=gt_label_names,
                       pred_label_names=pred_label_names,
                       offset=FLAGS.offset,
                       semantics=not FLAGS.ignore_semantics, instances=FLAGS.do_instances and not FLAGS.ignore_semantics)

    # print instructions
    print("To navigate:")
    print("\tb: back (previous scan)")
    print("\tn: next (next scan)")
    print("\tq: quit (exit program)")
    print("Describe:\n"
          "image:     \t-------------\n"
          "           \tRange image\n"
          "           \tGround Truth\n"
          "           \t{Predictions}\n"
          "           \t-------------\n"
          "PointCloud:\t[Ground Truth  ||  {Predictions}]")

    # run the visualizer
    vis.run()

utils/visualize_mos.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import os
import sys
import cv2
import numpy as np
from glob import glob


if __name__ == '__main__':
    
    path = "data/sequences/00/"

    save_path = os.path.join(path, "concate_residual")

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    files = sorted(glob(os.path.join(path, "velodyne/*.bin")))

    for fid in range(8, len(files)):

        for i in range(1, 9):
            tmp = f"{path}/visualization_{i}/{'%06d'%fid}.png"
            print(tmp)
            tmp_img = cv2.imread(tmp)
            if i == 1:
                img = tmp_img
            else:
                img = np.concatenate((img, tmp_img), axis=0)

        if False:
            cv2.imwrite(f"{save_path}/{'%06d'%fid}_8res.png", img)

        cv2.imshow(f"residual image", img)
        cv2.waitKey()

utils/viz_concate_residuals.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# @author: Jiadai Sun

import sys
import cv2
import yaml
import numpy as np

from kitti_utils import load_vertex, load_labels
from auxiliary.laserscan import LaserScan, SemLaserScan
from matplotlib import pyplot as plt
from utils import load_yaml, check_and_makedirs

def get_mpl_colormap(cmap_name):
    cmap = plt.get_cmap(cmap_name)
    # Initialize the matplotlib color map
    sm = plt.cm.ScalarMappable(cmap=cmap)
    # Obtain linear color range
    color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1]
    return color_range.reshape(256, 1, 3)


def remap(label, mapdict):
    # put label from original values to xentropy
    # or vice-versa, depending on dictionary values
    # make learning map a lookup table
    maxkey = 0
    for key, data in mapdict.items():
        if isinstance(data, list):
            nel = len(data)
        else:
            nel = 1
        if key > maxkey:
            maxkey = key
    # +100 hack making lut bigger just in case there are unknown labels
    if nel > 1:
        lut = np.zeros((maxkey + 100, nel), dtype=np.int32)
    else:
        lut = np.zeros((maxkey + 100), dtype=np.int32)
    for key, data in mapdict.items():
        try:
            lut[key] = data
        except IndexError:
            print("Wrong key ", key)
    # do the mapping
    return lut[label]


if __name__ == '__main__':

    seq = "08"
    frame_id = [98, 218, 222, 1630, 1641, 50]

    data_path = "/path_to_dataset/"
    save_path = "plot2d_result"
    check_and_makedirs(save_path)

    path = {
        "gtlabel": data_path,
        "method1": "/the prediction result of method_1",
        "method2": "/the prediction result of method_2",
        "ours": "/the prediction result of ours"
    }

    config_path = "config/labels/semantic-kitti-mos.yaml"
    CFG = load_yaml(config_path)

    color_dict = CFG["color_map"]
    nclasses = len(color_dict)

    scan = SemLaserScan(nclasses, color_dict, project=True, H=64, W=2048, fov_up=3.0, fov_down=-25.0)

    for f_id in frame_id:
        str_fid = "%06d"%(f_id)
        print(str_fid)

        scan_path = f'{data_path}/sequences/{seq}/velodyne/{str_fid}.bin'
        scan.open_scan(scan_path)

        for key, value in path.items():
            if key == 'gtlabel':
                label_path = f'{value}/sequences/{seq}/labels/{str_fid}.label'
            else:
                label_path = f'{value}/sequences/{seq}/predictions/{str_fid}.label'

            print(key)

            scan.open_label(label_path)
            scan.sem_label = remap(scan.sem_label, CFG["learning_map"])
            scan.sem_label = remap(scan.sem_label, CFG["learning_map_inv"])
            # print(scan.sem_label.max())
            scan.colorize()
            scan.do_label_projection()

            power = 16
            data = np.copy(scan.proj_range)

            data[data > 0] = data[data > 0]**(1 / power)
            data[data < 0] = data[data > 0].min()

            data = (data - data[data > 0].min()) / \
                (data.max() - data[data > 0].min()) * 255

            out_img = cv2.applyColorMap(data.astype(np.uint8), get_mpl_colormap('viridis'))
            imgpath = f'{save_path}/Range2D_{key}_seq{seq}_fid{f_id}.png'
            cv2.imwrite(imgpath, out_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])

            seg_vis2D = scan.proj_sem_color * 255 # (64, 1024, 3)
            imgpath = f'{save_path}/Pred2d_{key}_seq{seq}_fid{f_id}.png'
            cv2.imwrite(imgpath, seg_vis2D, [cv2.IMWRITE_PNG_COMPRESSION, 0])

utils/viz_mos_result_2d.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.
# @author: Jiadai Sun

import numpy as np
import open3d as o3d
from kitti_utils import load_vertex, load_labels

if __name__ == '__main__':

    seq = "08"
    frame_id = [98, 218, 222, 1630, 1641, 4016]

    data_path = "/path_to_dataset/"
    path = {
        "gtlabel": data_path,
        "method1": "/the prediction result of method_1",
        "method2": "/the prediction result of method_2",
        "ours": "/the prediction result of ours"
    }

    for f_id in frame_id:
        str_fid = "%06d" % (f_id)
        print(str_fid)

        scan_path = f'{data_path}/sequences/{seq}/velodyne/{str_fid}.bin'
        scan = load_vertex(scan_path)

        for key, value in path.items():
            if key == 'gtlabel':
                label_path = f'{value}/sequences/{seq}/labels/{str_fid}.label'
            else:
                label_path = f'{value}/sequences/{seq}/predictions/{str_fid}.label'

            print(key)
            label, _ = load_labels(label_path)

            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(scan[:, :3])
            pcd.paint_uniform_color([0.25, 0.25, 0.25])
            colors = np.array(pcd.colors)
            colors[label > 200] = [1.0, 0.0, 0.0]

            pcd.colors = o3d.utility.Vector3dVector(colors)

            vis = o3d.visualization.Visualizer()
            vis.create_window(
                window_name=f'{key}_seq{seq}_frame{f_id}', width=1000, height=1000)
            vis.add_geometry(pcd)
            # parameters = o3d.io.read_pinhole_camera_parameters("/home/user/Repo/LiDAR-MOS/ScreenCamera_2022-02-20-21-03-42.json")
            # ctr = vis.get_view_control()
            # ctr.convert_from_pinhole_camera_parameters(parameters)
            vis.run()
            vis.destroy_window()

utils/viz_mos_result_o3d.py
SCNU-RISLAB/MF-MOS
#!/usr/bin/env python3
# Developed by Xieyuanli Chen and Thomas Läbe
# This file is covered by the LICENSE file in the root of the project OverlapNet: https://github.com/PRBonn/OverlapNet
# Brief: a script to generate depth data

import os
import cv2
import numpy as np
import scipy.linalg as linalg

from kitti_utils import load_files, range_projection
from matplotlib import pyplot as plt


def get_mpl_colormap(cmap_name):
    cmap = plt.get_cmap(cmap_name)
    # Initialize the matplotlib color map
    sm = plt.cm.ScalarMappable(cmap=cmap)
    # Obtain linear color range
    color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1]
    return color_range.reshape(256, 1, 3)


def rotate_mat( axis, radian):
    rot_matrix = linalg.expm(np.cross(np.eye(3), axis / linalg.norm(axis) * radian))
    return rot_matrix
    # print(type(rot_matrix))


def gen_depth_data(scan_folder, dst_folder, normalize=False):
    """ Generate projected range data in the shape of (64, 900, 1).
        The input raw data are in the shape of (Num_points, 3).
    """
    # specify the goal folder
    dst_folder = os.path.join(dst_folder, 'depth')
    try:
        os.stat(dst_folder)
        print('generating depth data in: ', dst_folder)
    except:
        print('creating new depth folder: ', dst_folder)
        os.makedirs(dst_folder)

    # load LiDAR scan files
    scan_paths = load_files(scan_folder)

    depths = []
    axis_x, axis_y, axis_z = [1,0,0], [0,1,0], [0, 0, 1]

    # iterate over all scan files
    for idx in range(len(scan_paths)):
        # load a point cloud
        current_vertex = np.fromfile(scan_paths[idx], dtype=np.float32)
        current_vertex = current_vertex.reshape((-1, 4))

        # proj_vertex = range_projection(current_vertex, fov_up=14.8940, fov_down=-16.1760, proj_H=32, proj_W=900, max_range=50)
        proj_vertex = range_projection(current_vertex, fov_up=3, fov_down=-25, proj_H=64, proj_W=2048, max_range=50)

        proj_range = proj_vertex[:, :, -1]
        # normalize the image
        if normalize:
            proj_range = proj_range / np.max(proj_range) * 255

        # generate the destination path
        dst_path = os.path.join(dst_folder, str(idx).zfill(6))

        # np.save(dst_path, proj_range)
        filename = dst_path + ".png"
        color_img = cv2.applyColorMap(proj_range.astype(np.uint8),  get_mpl_colormap('viridis'))#cv2.COLORMAP_RAINBOW)

        cv2.imwrite(filename, color_img)
        print('finished generating depth data at: ', filename)

    return depths


if __name__ == '__main__':

    scan_folder = "data/sequences/05/velodyne"
    save_folder = 'viz_depth_result'

    depth_data = gen_depth_data(scan_folder, save_folder)

utils/viz_range_depth_img.py
SCNU-RISLAB/MF-MOS
import open3d as o3d
import numpy as np
import os
import time
import yaml
from tqdm import tqdm
class PointCloudPlayer():
    def __init__(self, velodyne_path, label_path, windows_w=1920, windows_h=1080,
                 background_color=[0, 0, 0], point_size=1.0, to_reset=True, color_cfg_path=None, viewfile_path="./viewfile.json"):
        self.velodyne_path = velodyne_path
        self.label_path = label_path

        self.windows_w = windows_w
        self.windows_h = windows_h
        self.background_color = background_color
        self.point_size = point_size
        self.to_reset = to_reset
        self.COLOR_CFG = yaml.safe_load(open(color_cfg_path, 'r'))
        self.viewfile_path = viewfile_path

        # self.get_vis()
        self.get_datas()
        self.set_colormap()

    def get_vis(self):
        self.vis = o3d.visualization.VisualizerWithKeyCallback()
        self.vis.create_window(width=1920, height=1080)  # 创建窗口
        # self.vis.set_full_screen(True)
        # #设置连续帧 雷达第一视角

        # self.pcd = o3d.geometry.PointCloud()
        # self.vis.add_geometry(self.pcd)

        render_option = self.vis.get_render_option()  # 渲染配置
        render_option.background_color = np.array(self.background_color)  # 设置点云渲染参数，背景颜色
        render_option.point_size = self.point_size  # 设置渲染点的大小

        self.playing = True

        def exit_callback():
            print("exit")
            self.vis.close()
            quit()
            return True

        def pause_callback():
            print("key pause")
            self.playing = not self.playing
            return True

        self.vis.register_key_callback(ord("q"), exit_callback)
        self.vis.register_key_callback(ord("c"), pause_callback)

    def get_datas(self):
        self.pc_datas = []
        self.labels = []
        assert len(os.listdir(self.velodyne_path)) == len(os.listdir(self.label_path))

        for name in os.listdir(self.velodyne_path):
            pc_file = os.path.join(self.velodyne_path, name)
            label_file = os.path.join(self.label_path, name.replace(".bin", ".label"))
            self.pc_datas.append(np.fromfile(pc_file, dtype=np.float32).reshape(-1, 4))
            self.labels.append(np.fromfile(label_file, dtype=np.uint32).reshape((-1)) & 0xFFFF)

    def set_colormap(self):
        self.color_dict = self.COLOR_CFG["color_map"]
        moving_learning_map = self.COLOR_CFG["moving_learning_map"]
        movable_learning_map = self.COLOR_CFG["movable_learning_map"]

        for key in self.color_dict.keys():
            if (key == 250) or (key in movable_learning_map.keys() and movable_learning_map[key] == 2):
                self.color_dict[key] = [0, 0, 255]
                if key != 250 and moving_learning_map[key] == 2:
                    self.color_dict[key] = [255, 0, 0]
            else:
                self.color_dict[key] = [255, 255, 255]

    def draw_color(self, point_xyz, label):
        colors = np.zeros((point_xyz.shape[0], 3))
        for i in list(set(label.tolist())):
            if i not in self.color_dict.keys():
                colors[label == i] = np.array([0, 255, 0])
                continue
            colors[label == i] = np.array(self.color_dict[i])

        return colors

    def save_view_file(self, pcd_numpy):
        vis = o3d.visualization.Visualizer()
        vis.create_window()
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(pcd_numpy)
        vis.add_geometry(pcd)
        axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1, origin=[0, 0, 0])
        vis.add_geometry(axis)
        vis.run()  # user changes the view and press "q" to terminate
        param = vis.get_view_control().convert_to_pinhole_camera_parameters()
        o3d.io.write_pinhole_camera_parameters(self.viewfile_path, param)
        vis.destroy_window()

    def play(self):
        if not os.path.exists(self.viewfile_path):
            self.save_view_file(- self.pc_datas[0][:, :3])

        pcds = []
        self.get_vis()
        with tqdm(total=len(self.pc_datas)) as pbar:
            pbar.set_description("Reading point cloud...")
            for pc_data, label in zip(self.pc_datas, self.labels):
                point_xyz = - pc_data[:, :3]  # x, y, z

                pcd = o3d.open3d.geometry.PointCloud()  # 创建点云对象
                pcd.points = o3d.utility.Vector3dVector(point_xyz)
                pcd.colors = o3d.utility.Vector3dVector(self.draw_color(point_xyz, label))

                pcds.append(pcd)
                # self.vis.add_geometry(pcd)

                pbar.update(1)

        ctr = self.vis.get_view_control()
        param = o3d.io.read_pinhole_camera_parameters(self.viewfile_path)
        ctr.convert_from_pinhole_camera_parameters(param)

        for pcd in pcds:
            self.vis.clear_geometries()
            # axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1, origin=[0, 0, 0])
            # self.vis.add_geometry(axis)
            self.vis.add_geometry(pcd)
            # self.vis.poll_events()
            # self.vis.update_renderer()
            ctr.convert_from_pinhole_camera_parameters(param)
            time.sleep(0.2)
            self.vis.run()
        self.vis.destroy_window()


if __name__ == "__main__":
    use_pred = True
    dataset_root_path = "/data1/MotionSeg3d_Dataset/data_odometry_velodyne/dataset/sequences"
    label_root_path = "/data1/MF-MOS/log/Valid/predictions/sequences"
    # label_root_path = "/data1/MotionSeg3d_Dataset/data_odometry_velodyne/dataset/sequences"

    seq = "07"
    if use_pred:
        labels_path = os.path.join(label_root_path, seq, "predictions_fuse")
    else:
        labels_path = os.path.join(label_root_path, seq, "labels")
    pcp = PointCloudPlayer(velodyne_path=os.path.join(dataset_root_path, seq, "velodyne"),
                           label_path=labels_path,
                           color_cfg_path="/data1/MF-MOS/config/labels/semantic-kitti-mos.raw.yaml",
                           point_size=1.2)

    pcp.play()
utils/viz_seqVideo.py
SCNU-RISLAB/MF-MOS