Class FreeMaskPreprocessing

datasets/preprocessing/freemask_preprocessing.py:14–241

from utils.point_cloud_utils import load_ply_with_normals

class FreeMaskPreprocessing(BasePreprocessing):

    FREEMASK_CLASS_IDS = (0, 1)
    FREEMASK_CLASS_NAMES = ('background', 'foreground')
    FREEMASK_COLOR_MAP = {0: (0, 0, 0), 1: (0, 0, 128)}
    FREEMASK_ORACLE_CLASS_IDS = (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39)

    def __init__(
            self,
            data_dir: str = "/canis/Datasets/ScanNet/public/v2",
            save_dir: str = "./data/processed/freemask",
            modes: tuple = ("train", "validation"),
            n_jobs: int = -1,
            git_repo: str = "data/Datasets/ScanNet/ScanNet",
            oracle: bool = False,
            freemask_dir: str = "/mnt/data/Datasets/ScanNetFreeMask"):

        super().__init__(data_dir, save_dir, modes, n_jobs)

        self.oracle = oracle
        git_repo = Path(git_repo)
        self.create_label_database(git_repo)
        self.freemask_base_path = Path(freemask_dir)

        for mode in self.modes:
            trainval_split_dir = git_repo / "Tasks" / "Benchmark"
            scannet_special_mode = "val" if mode == "validation" else mode
            with open(
                    trainval_split_dir / (f"scannetv2_{scannet_special_mode}.txt")
            ) as f:
                # -1 because the last one is always empty
                split_file = f.read().split("\n")[:-1]

            scans_folder = "scans_test" if mode == "test" else "scans"
            filepaths = []
            for scene in split_file:
                filepaths.append(
                    self.data_dir / scans_folder / scene / (scene + "_vh_clean_2.ply")
                )
            self.files[mode] = natsorted(filepaths)

    def create_label_database(self, git_repo):
        label_database = {}
        for row_id, class_id in enumerate(self.FREEMASK_CLASS_IDS):
            label_database[class_id] = {
                'color': self.FREEMASK_COLOR_MAP[class_id],
                'name': self.FREEMASK_CLASS_NAMES[row_id],
                'validation': True
            }
        self._save_yaml(self.save_dir / "label_database.yaml", label_database)
        return label_database

    def process_file(self, filepath, mode):
        """process_file.

       The first part is analogous to the scannet200 preprocessing.
       In second part we load the freemask data and masks proposals, which we assign to the points.