code/pointcept/datasets/semantic_kitti_kevin.py

"""
Semantic KITTI dataset

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import numpy as np

from .builder import DATASETS
from .defaults import DefaultDataset


@DATASETS.register_module()
class SemanticKITTIDataset(DefaultDataset):
    def __init__(
        self,
        split="train",
        data_root="/media/changbryan/BC_T7/itriDataset/",
        transform=None,
        test_mode=False,
        test_cfg=None,
        loop=1,
        ignore_index=-1,
    ):
        self.ignore_index = ignore_index
        self.learning_map = self.get_learning_map(ignore_index)
        self.learning_map_inv = self.get_learning_map_inv(ignore_index)
        super().__init__(
            split=split,
            data_root=data_root,
            transform=transform,
            test_mode=test_mode,
            test_cfg=test_cfg,
            loop=loop,
        )

    def get_data_list(self):
        # # Generate sequence numbers 0-406
        # all_sequences = list(range(407))
        
        # # Randomly shuffle the sequences
        # np.random.seed(42)  # For reproducibility
        # np.random.shuffle(all_sequences)
        
        # # Split into train (60%), val (20%), test (20%)
        # train_size = int(0.7 * len(all_sequences))
        # val_size = int(0.15 * len(all_sequences))
        
        split2seq = dict(
        #     train=all_sequences[:train_size],
        #     val=all_sequences[train_size:train_size + val_size],
        #     test=all_sequences[train_size + val_size:],
            train=[00],  # Commented out debug values
            val=[00],
            test=[00],
        )
        
        if isinstance(self.split, str):
            seq_list = split2seq[self.split]
        elif isinstance(self.split, list):
            seq_list = []
            for split in self.split:
                seq_list += split2seq[split]
        else:
            raise NotImplementedError

        data_list = []
        for seq in seq_list:
            seq = str(seq).zfill(2)
            seq_folder = os.path.join(self.data_root, "dataset", "sequences", seq)
            seq_files = sorted(os.listdir(os.path.join(seq_folder, "velodyne")))
            
            # Determine the number of frames for each split
            num_files = len(seq_files)
            num_train = int(num_files * 0.7)
            num_val = int(num_files * 0.15)
            num_test = num_files - num_train - num_val  # Remaining frames for test

            # Split frames within each sequence
            frame_sequences = np.array(range(num_files))
            np.random.seed(42)  # For reproducibility
            np.random.shuffle(frame_sequences)
            train_files = sorted(np.take(seq_files, frame_sequences[:num_train]))
            val_files = sorted(np.take(seq_files, frame_sequences[num_train:num_train + num_val]))
            test_files = sorted(np.take(seq_files, frame_sequences[num_train + num_val:]))

            np.save('frameSequence.npy', frame_sequences)

            # Add files to data_list based on the specified split
            if self.split == "train":
                data_list += [os.path.join(seq_folder, "velodyne", file) for file in train_files]
            elif self.split == "val":
                data_list += [os.path.join(seq_folder, "velodyne", file) for file in val_files]
            elif self.split == "test":
                data_list += [os.path.join(seq_folder, "velodyne", file) for file in test_files]

        return data_list

    def get_data(self, idx):
        data_path = self.data_list[idx % len(self.data_list)]
        with open(data_path, "rb") as b:
            scan = np.fromfile(b, dtype=np.float32).reshape(-1, 4)
        coord = scan[:, :3]
        strength = scan[:, -1].reshape([-1, 1])
        #modified for MT
        #strength = scan[:, -1].reshape([-1, 1])/255.

        label_file = data_path.replace("velodyne", "labels").replace(".bin", ".label")
        if os.path.exists(label_file):
            with open(label_file, "rb") as a:
                segment = np.fromfile(a, dtype=np.int32).reshape(-1)
                segment = np.vectorize(self.learning_map.__getitem__)(
                    segment & 0xFFFF
                ).astype(np.int32)
        else:
            segment = np.zeros(scan.shape[0]).astype(np.int32)
        data_dict = dict(coord=coord, strength=strength, segment=segment)
        return data_dict

    def get_data_name(self, idx):
        file_path = self.data_list[idx % len(self.data_list)]
        dir_path, file_name = os.path.split(file_path)
        sequence_name = os.path.basename(os.path.dirname(dir_path))
        frame_name = os.path.splitext(file_name)[0]
        data_name = f"{sequence_name}_{frame_name}"
        return data_name

    @staticmethod
    def get_learning_map(ignore_index):
        learning_map = {
            # 0: ignore_index,  # "unlabeled"
            # 1: ignore_index,  # "outlier" mapped to "unlabeled" --------------------------mapped
            # 10: 0,  # "car"
            # 11: 1,  # "bicycle"
            # 13: 4,  # "bus" mapped to "other-vehicle" --------------------------mapped
            # 15: 2,  # "motorcycle"
            # 16: 4,  # "on-rails" mapped to "other-vehicle" ---------------------mapped
            # 18: 3,  # "truck"
            # 20: 4,  # "other-vehicle"
            # 30: 5,  # "person"
            # 31: 6,  # "bicyclist"
            # 32: 7,  # "motorcyclist"
            # 40: 8,  # "road"
            # 44: 9,  # "parking"
            # 48: 10,  # "sidewalk"
            # 49: 11,  # "other-ground"
            # 50: 12,  # "building"
            # 51: 13,  # "fence"
            # 52: ignore_index,  # "other-structure" mapped to "unlabeled" ------------------mapped
            # 60: 8,  # "lane-marking" to "road" ---------------------------------mapped
            # 70: 14,  # "vegetation"
            # 71: 15,  # "trunk"
            # 72: 16,  # "terrain"
            # 80: 17,  # "pole"
            # 81: 18,  # "traffic-sign"
            # 99: ignore_index,  # "other-object" to "unlabeled" ----------------------------mapped
            # 252: 0,  # "moving-car" to "car" ------------------------------------mapped
            # 253: 6,  # "moving-bicyclist" to "bicyclist" ------------------------mapped
            # 254: 5,  # "moving-person" to "person" ------------------------------mapped
            # 255: 7,  # "moving-motorcyclist" to "motorcyclist" ------------------mapped
            # 256: 4,  # "moving-on-rails" mapped to "other-vehicle" --------------mapped
            # 257: 4,  # "moving-bus" mapped to "other-vehicle" -------------------mapped
            # 258: 3,  # "moving-truck" to "truck" --------------------------------mapped
            # 259: 4,  # "moving-other"-vehicle to "other-vehicle" ----------------mapped
            


            0: 0,  # "unlabeled"
            # used the labeling from the shp file 
            # map the specific road lines to dahed, solid, and double
            1: 1,  
            # 2: 1,
            # 3: 3,
            # 4: 1,
            # 5: 1,
            # 6: 1,

            # 11: 2,
            # 12: 1,
            # 13: 3,
            # 14: 1,
            # 15: 1,
            # 16: 1,

            # 21: 2,
            # 22: 1,
            # 23: 3,
            # 24: 1,
            # 25: 1,
            # 26: 1,
            # #labeling used in ITRI semantic maps
            # 41: 1,
            # 42: 2,
            # 43: 3,
            # 44: 1,
            # 45: 2,
            # 46: 3,
            # 47: 2,
        }
        return learning_map

    @staticmethod
    def get_learning_map_inv(ignore_index):
        learning_map_inv = {
            # ignore_index: ignore_index,  # "unlabeled"
            # 0: 10,  # "car"
            # 1: 11,  # "bicycle"
            # 2: 15,  # "motorcycle"
            # 3: 18,  # "truck"
            # 4: 20,  # "other-vehicle"
            # 5: 30,  # "person"
            # 6: 31,  # "bicyclist"
            # 7: 32,  # "motorcyclist"
            # 8: 40,  # "road"
            # 9: 44,  # "parking"
            # 10: 48,  # "sidewalk"
            # 11: 49,  # "other-ground"
            # 12: 50,  # "building"
            # 13: 51,  # "fence"
            # 14: 70,  # "vegetation"
            # 15: 71,  # "trunk"
            # 16: 72,  # "terrain"
            # 17: 80,  # "pole"
            # 18: 81,  # "traffic-sign"


            ignore_index: ignore_index,
            0: 0,  # "unlabeled"
            1: 1,  
            # 2: 2,
            # 3: 3,
        }
        return learning_map_inv