import os import re import sys import os.path as osp import torch import random import logging import hashlib import warnings from tqdm import tqdm from datetime import datetime from itertools import product from tqdm.auto import tqdm as tq from typing import Any, List, Tuple, Union from torch_geometric.data import InMemoryDataset from torch_geometric.data.dataset import files_exist from torch_geometric.data.makedirs import makedirs from torch_geometric.data.dataset import _repr from torch_geometric.nn.pool.consecutive import consecutive_cluster from src.data import NAG from src.transforms import Transform, NAGSelectByKey, NAGRemoveKeys, \ SampleXYTiling, SampleRecursiveMainXYAxisTiling from src.visualization import show DIR = os.path.dirname(os.path.realpath(__file__)) log = logging.getLogger(__name__) __all__ = ['BaseDataset'] ######################################################################## # BaseDataset # ######################################################################## class BaseDataset(InMemoryDataset): """Base class for datasets. Child classes must overwrite the following methods (see respective docstrings for more details): ``` MyDataset(BaseDataset): def class_names(self): pass def num_classes(self): pass def stuff_classes(self): pass def class_colors(self): # Optional: only if you want to customize your color palette # for visualization pass def all_base_cloud_ids(self): pass def download_dataset(self): pass def read_single_raw_cloud(self): pass def raw_file_structure(self): # Optional: only if your raw or processed file structure # differs from the default pass def id_to_relative_raw_path(self): # Optional: only if your raw or processed file structure # differs from the default pass def processed_to_raw_path(self): # Optional: only if your raw or processed file structure # differs from the default pass ``` Parameters ---------- root : `str` Root directory where the dataset should be saved. stage : {'train', 'val', 'test', 'trainval'} transform : `callable` transform function operating on data. pre_transform : `callable` pre_transform function operating on data. pre_filter : `callable` pre_filter function operating on data. on_device_transform: `callable` on_device_transform function operating on data, in the 'on_after_batch_transfer' hook. This is where GPU-based augmentations should be, as well as any Transform you do not want to run in CPU-based DataLoaders val_mixed_in_train: bool whether the 'val' stage data is saved in the same clouds as the 'train' stage. This may happen when the stage splits are performed inside the clouds. In this case, an `on_device_transform` will be automatically created to separate stage-specific data upon reading test_mixed_in_val: bool whether the 'test' stage data is saved in the same clouds as the 'val' stage. This may happen when the stage splits are performed inside the clouds. In this case, an `on_device_transform` will be automatically created to separate stage-specific data upon reading custom_hash: str A user-chosen hash to be used for the dataset data directory. This will bypass the default behavior where the pre_transforms are used to generate a hash. It can be used, for instance, when one wants to instantiate a dataset with already-processed data, without knowing the exact config that was used to generate it in_memory: bool If True, the processed dataset will be entirely loaded in RAM upon instantiation. This will accelerate training and inference but requires large memory. WARNING: __getitem__ directly returns the data in memory, so any modification to the returned object will affect the `in_memory_data` too. Be careful to clone the object before modifying it. Besides, the `transform` are pre-applied to the in_memory data point_save_keys: list[str] List of point (ie level-0) attribute keys to save to disk at the end of preprocessing. Leaving to `None` will save all attributes by default point_no_save_keys: list[str] List of point (ie level-0) attribute keys to NOT save to disk at the end of preprocessing point_load_keys: list[str] List of point (ie level-0) attribute keys to load when reading data from disk segment_save_keys: list[str] List of segment (ie level-1+) attribute keys to save to disk at the end of preprocessing. Leaving to `None` will save all attributes by default segment_no_save_keys: list[str] List of segment (ie level-1+) attribute keys to NOT save to disk at the end of preprocessing segment_load_keys: list[str] List of segment (ie level-1+) attribute keys to load when reading data from disk """ def __init__( self, root: str, stage: str = 'train', transform: Transform = None, pre_transform: Transform = None, pre_filter: Transform = None, on_device_transform: Transform = None, save_y_to_csr: bool = True, save_pos_dtype: torch.dtype = torch.float, save_fp_dtype: torch.dtype = torch.half, xy_tiling: int = None, pc_tiling: int = None, val_mixed_in_train: bool = False, test_mixed_in_val: bool = False, custom_hash: str = None, in_memory: bool = False, point_save_keys: List[str] = None, point_no_save_keys: List[str] = None, point_load_keys: List[str] = None, segment_save_keys: List[str] = None, segment_no_save_keys: List[str] = None, segment_load_keys: List[str] = None, **kwargs): assert stage in ['train', 'val', 'trainval', 'test'] # Set these attributes before calling parent `__init__` because # some attributes will be needed in parent `download` and # `process` methods self._stage = stage self._save_y_to_csr = save_y_to_csr self._save_pos_dtype = save_pos_dtype self._save_fp_dtype = save_fp_dtype self._on_device_transform = on_device_transform self._val_mixed_in_train = val_mixed_in_train self._test_mixed_in_val = test_mixed_in_val self._custom_hash = custom_hash self._in_memory = in_memory self._point_save_keys = point_save_keys self._point_no_save_keys = point_no_save_keys self._point_load_keys = point_load_keys self._segment_save_keys = segment_save_keys self._segment_no_save_keys = segment_no_save_keys self._segment_load_keys = segment_load_keys if in_memory: log.warning( "'in_memory' was set to True. This means the entire dataset " "will be held in RAM. While this allows training and inference " "speedups, this means that the `transform' will only be " "applied once, upon loading the dataset to RAM. Hence, if you " "need augmentations or any other stochastic operations to be " "applied on your batches, make sure you moved them all to " "'on_device_transform'.") # Prepare tiling arguments. Can either be XY tiling of PC # tiling but not both. XY tiling will apply a regular grid along # the XY axes to the data, regardless of its orientation, shape # or density. The value of xy_tiling indicates the number of # tiles in each direction. So, if a single int is passed, each # cloud will be divided into xy_tiling**2 tiles. PC tiling will # recursively split the data wrt the principal component along # the XY plane. Each step splits the data in 2, wrt to its # geometry. The value of pc_tiling indicates the number of split # steps used. Hence, 2**pc_tiling tiles will be created. assert xy_tiling is None or pc_tiling is None, \ "Cannot apply both XY and PC tiling, please choose only one." if xy_tiling is None: self.xy_tiling = None elif isinstance(xy_tiling, int): self.xy_tiling = (xy_tiling, xy_tiling) if xy_tiling > 1 else None elif xy_tiling[0] > 1 or xy_tiling[1] > 1: self.xy_tiling = xy_tiling else: self.xy_tiling = None self.pc_tiling = pc_tiling if pc_tiling and pc_tiling >= 1 else None # Sanity check on the cloud ids. Ensures cloud ids are unique # across all stages, unless `val_mixed_in_train` or # `test_mixed_in_val` is True self.check_cloud_ids() # Initialization with downloading and all preprocessing root = osp.join(root, self.data_subdir_name) super().__init__(root, transform, pre_transform, pre_filter) # Display the dataset pre_transform_hash and full path path = osp.join(self.processed_dir, "", self.pre_transform_hash) log.info(f'Dataset hash: "{self.pre_transform_hash}"') log.info(f'Preprocessed data can be found at: "{path}"') # If `val_mixed_in_train` or `test_mixed_in_val`, we will need # to separate some stage-related data at reading time. # Since this operation can be computationally-costly, we prefer # postponing it to the `on_device_transform`. To this end, we # prepend the adequate transform to the dataset's # `on_device_transform`. Otherwise, if we have no mixed-stages, # we simply remove all `is_val` attributes in the # `on_device_transform` if self.stage == 'train' and self.val_mixed_in_train: t = NAGSelectByKey(key='is_val', negation=True) elif self.stage == 'val' and self.val_mixed_in_train or self.test_mixed_in_val: t = NAGSelectByKey(key='is_val', negation=False) elif self.stage == 'test' and self.test_mixed_in_val: t = NAGSelectByKey(key='is_val', negation=True) else: t = NAGRemoveKeys(level='all', keys=['is_val'], strict=False) # Make sure a NAGRemoveKeys for `is_val` does not already exist # in the `on_device_transform` before prepending the transform if not any( isinstance(odt, NAGSelectByKey) and odt.key == 'is_val' for odt in self.on_device_transform.transforms): self._on_device_transform.transforms = \ [t] + self._on_device_transform.transforms # Load the processed data, if the dataset must be in memory if self.in_memory: in_memory_data = [ NAG.load( self.processed_paths[i], keys_low=self.point_load_keys, keys=self.segment_load_keys) for i in range(len(self))] if self.transform is not None: in_memory_data = [self.transform(x) for x in in_memory_data] self._in_memory_data = in_memory_data else: self._in_memory_data = None @property def class_names(self) -> List[str]: """List of string names for dataset classes. This list must be one-item larger than `self.num_classes`, with the last label corresponding to 'void', 'unlabelled', 'ignored' classes, indicated as `y=self.num_classes` in the dataset labels. """ raise NotImplementedError @property def num_classes(self) -> int: """Number of classes in the dataset. Must be one-item smaller than `self.class_names`, to account for the last class name being used for 'void', 'unlabelled', 'ignored' classes, indicated as `y=self.num_classes` in the dataset labels. """ raise NotImplementedError @property def stuff_classes(self) -> List[int]: """List of 'stuff' labels for INSTANCE and PANOPTIC SEGMENTATION (setting this is NOT REQUIRED FOR SEMANTIC SEGMENTATION alone). By definition, 'stuff' labels are labels in `[0, self.num_classes-1]` which are not 'thing' labels. In instance segmentation, 'stuff' classes are not taken into account in performance metrics computation. In panoptic segmentation, 'stuff' classes are taken into account in performance metrics computation. Besides, each cloud/scene can only have at most one instance of each 'stuff' class. IMPORTANT: By convention, we assume `y ∈ [0, self.num_classes-1]` ARE ALL VALID LABELS (i.e. not 'ignored', 'void', 'unknown', etc), while `y < 0` AND `y >= self.num_classes` ARE VOID LABELS. """ raise NotImplementedError @property def thing_classes(self) -> List[int]: """List of 'thing' labels for instance and panoptic segmentation. By definition, 'thing' labels are labels in `[0, self.num_classes-1]` which are not 'stuff' labels. IMPORTANT: By convention, we assume `y ∈ [0, self.num_classes-1]` ARE ALL VALID LABELS (i.e. not 'ignored', 'void', 'unknown', etc), while `y < 0` AND `y >= self.num_classes` ARE VOID LABELS. """ return [i for i in range(self.num_classes) if i not in self.stuff_classes] @property def void_classes(self) -> List[int]: """List containing the 'void' labels. By default, we group all void/ignored/unknown class labels into a single `[self.num_classes]` label for simplicity. IMPORTANT: By convention, we assume `y ∈ [0, self.num_classes-1]` ARE ALL VALID LABELS (i.e. not 'ignored', 'void', 'unknown', etc), while `y < 0` AND `y >= self.num_classes` ARE VOID LABELS. """ return [self.num_classes] @property def class_colors(self) -> List[List[int]]: """Colors for visualization, if not None, must have the same length as `self.num_classes`. If None, the visualizer will use the label values in the data to generate random colors. """ return def print_classes(self) -> None: """Show the class names, labels and type (thing, stuff, void). """ for i, c in enumerate(self.class_names): try: class_type = \ 'stuff' if i in self.stuff_classes \ else 'thing' if i in self.thing_classes \ else 'void' except: class_type = '' print(f"{i:<3} {c:<20} {class_type}") @property def data_subdir_name(self) -> str: return self.__class__.__name__.lower() @property def stage(self) -> str: """Dataset stage. Expected to be 'train', 'val', 'trainval', or 'test' """ return self._stage @property def save_y_to_csr(self) -> bool: return self._save_y_to_csr @property def save_pos_dtype(self) -> bool: return self._save_pos_dtype @property def save_fp_dtype(self) -> bool: return self._save_fp_dtype @property def on_device_transform(self) -> Transform: return self._on_device_transform @property def val_mixed_in_train(self) -> bool: return self._val_mixed_in_train @property def test_mixed_in_val(self) -> bool: return self._test_mixed_in_val @property def custom_hash(self) -> str: return self._custom_hash @property def in_memory(self) -> bool: return self._in_memory @property def point_save_keys(self) -> List[str]: return self._point_save_keys @property def point_no_save_keys(self) -> List[str]: return self._point_no_save_keys @property def point_load_keys(self) -> List[str]: return self._point_load_keys @property def segment_save_keys(self) -> List[str]: return self._segment_save_keys @property def segment_no_save_keys(self) -> List[str]: return self._segment_no_save_keys @property def segment_load_keys(self) -> List[str]: return self._segment_load_keys @property def all_base_cloud_ids(self) -> List[str]: """Dictionary holding lists of clouds ids, for each stage. The following structure is expected: `{'train': [...], 'val': [...], 'test': [...]}` """ raise NotImplementedError @property def all_cloud_ids(self) -> List[str]: """Dictionary holding lists of clouds ids, for each stage. Unlike all_base_cloud_ids, these ids take into account the clouds tiling, if any. """ # If clouds are tiled, expand and append all cloud names with a # suffix indicating which tile it corresponds to if self.xy_tiling is not None: tx, ty = self.xy_tiling return { stage: [ f'{ci}__TILE_{x + 1}-{y + 1}_OF_{tx}-{ty}' for ci in ids for x, y in product(range(tx), range(ty))] for stage, ids in self.all_base_cloud_ids.items()} if self.pc_tiling is not None: return { stage: [ f'{ci}__TILE_{x + 1}_OF_{2**self.pc_tiling}' for ci in ids for x in range(2**self.pc_tiling)] for stage, ids in self.all_base_cloud_ids.items()} # If no tiling needed, return the all_base_cloud_ids return self.all_base_cloud_ids def id_to_base_id(self, id: str) -> str: """Given an ID, remove the tiling indications, if any. """ if self.xy_tiling is None and self.pc_tiling is None: return id return self.get_tile_from_path(id)[1] @property def cloud_ids(self) -> List[str]: """IDs of the dataset clouds, based on its `stage`. """ if self.stage == 'trainval': ids = self.all_cloud_ids['train'] + self.all_cloud_ids['val'] else: ids = self.all_cloud_ids[self.stage] return sorted(list(set(ids))) def check_cloud_ids(self) -> None: """Make sure the `all_cloud_ids` are valid. More specifically, the cloud ids must be unique across all stages, unless `val_mixed_in_train=True` or `test_mixed_in_val=True`, in which case some clouds may appear in several stages """ train = set(self.all_cloud_ids['train']) val = set(self.all_cloud_ids['val']) test = set(self.all_cloud_ids['test']) assert len(train.intersection(val)) == 0 or self.val_mixed_in_train, \ "Cloud ids must be unique across all the 'train' and 'val' " \ "stages, unless `val_mixed_in_train=True`" assert len(val.intersection(test)) == 0 or self.test_mixed_in_val, \ "Cloud ids must be unique across all the 'val' and 'test' " \ "stages, unless `test_mixed_in_val=True`" @property def raw_file_structure(self) -> str: """String to describe to the user the file structure of your dataset, at download time. """ return @property def raw_file_names(self) -> str: """The file paths to find in order to skip the download.""" return self.raw_file_names_3d @property def raw_file_names_3d(self) -> str: """Some file paths to find in order to skip the download. Those are not directly specified inside `self.raw_file_names` in case `self.raw_file_names` would need to be extended (e.g. with 3D bounding boxes files). """ return [self.id_to_relative_raw_path(x) for x in self.cloud_ids] def id_to_relative_raw_path(self, id: str) -> str: """Given a cloud id as stored in `self.cloud_ids`, return the path (relative to `self.raw_dir`) of the corresponding raw cloud. """ return self.id_to_base_id(id) + '.ply' @property def pre_transform_hash(self) -> str: """Produce a unique but stable hash based on the dataset's `pre_transform` attributes (as exposed by `_repr`). """ if self.custom_hash is not None: return self.custom_hash if self.pre_transform is None: return 'no_pre_transform' return hashlib.md5(_repr(self.pre_transform).encode()).hexdigest() @property def processed_file_names(self) -> List[str]: """The name of the files to find in the `self.processed_dir` folder in order to skip the processing """ # For 'trainval', we use files from 'train' and 'val' to save # memory if self.stage == 'trainval' and self.val_mixed_in_train: return [ osp.join('train', self.pre_transform_hash, f'{w}.h5') for s in ('train', 'val') for w in self.all_cloud_ids[s]] if self.stage == 'trainval': return [ osp.join(s, self.pre_transform_hash, f'{w}.h5') for s in ('train', 'val') for w in self.all_cloud_ids[s]] return [ osp.join(self.stage, self.pre_transform_hash, f'{w}.h5') for w in self.cloud_ids] def processed_to_raw_path(self, processed_path: str) -> str: """Given a processed cloud path from `self.processed_paths`, return the absolute path to the corresponding raw cloud. Overwrite this method if your raw data does not follow the default structure. """ # Extract useful information from stage, hash_dir, cloud_id = \ osp.splitext(processed_path)[0].split(os.sep)[-3:] # Remove the tiling in the cloud_id, if any base_cloud_id = self.id_to_base_id(cloud_id) # Read the raw cloud data raw_ext = osp.splitext(self.raw_file_names_3d[0])[1] raw_path = osp.join(self.raw_dir, base_cloud_id + raw_ext) return raw_path @property def in_memory_data(self) -> Any: """If the `self.in_memory`, this will return all processed data, loaded in memory. Returns None otherwise. """ return self._in_memory_data @property def submission_dir(self) -> str: """Submissions are saved in the `submissions` folder, in the same hierarchy as `raw` and `processed` directories. Each submission has a subdirectory of its own, named based on the date and time of creation. """ submissions_dir = osp.join(self.root, "submissions") date = '-'.join([ f'{getattr(datetime.now(), x)}' for x in ['year', 'month', 'day']]) time = '-'.join([ f'{getattr(datetime.now(), x)}' for x in ['hour', 'minute', 'second']]) submission_name = f'{date}_{time}' path = osp.join(submissions_dir, submission_name) return path def download(self) -> None: self.download_warning() self.download_dataset() def download_dataset(self) -> None: """Download the dataset data. Modify this method to implement your own `BaseDataset` child class. """ raise NotImplementedError def download_warning(self, interactive: bool = False) -> None: # Warning message for the user about to download log.info( f"WARNING: You must download the raw data for the " f"{self.__class__.__name__} dataset.") if self.raw_file_structure is not None: log.info("Files must be organized in the following structure:") log.info(self.raw_file_structure) log.info("") if interactive: log.info("Press any key to continue, or CTRL-C to exit.") input("") log.info("") def download_message(self, msg: str) -> None: log.info(f'Downloading "{msg}" to {self.raw_dir}...') def _process(self) -> None: """Overwrites torch-geometric's Dataset._process. This simply removes the 'pre_transform.pt' file used for checking whether the pre-transforms have changed. This is possible thanks to our `pre_transform_hash` mechanism. """ f = osp.join(self.processed_dir, 'pre_filter.pt') if osp.exists(f) and torch.load(f) != _repr(self.pre_filter): warnings.warn( "The `pre_filter` argument differs from the one used in " "the pre-processed version of this dataset. If you want to " "make use of another pre-filtering technique, make sure to " "delete '{self.processed_dir}' first") if files_exist(self.processed_paths): # pragma: no cover return if self.log and 'pytest' not in sys.modules: print('Processing...', file=sys.stderr) makedirs(self.processed_dir) self.process() path = osp.join(self.processed_dir, 'pre_filter.pt') torch.save(_repr(self.pre_filter), path) if self.log and 'pytest' not in sys.modules: print('Done!', file=sys.stderr) def process(self) -> None: # If some stages have mixed clouds (they rely on the same cloud # files and the split is operated at reading time by # `on_device_transform`), we create symlinks between the # necessary folders, to avoid duplicate preprocessing # computation hash_dir = self.pre_transform_hash train_dir = osp.join(self.processed_dir, 'train', hash_dir) val_dir = osp.join(self.processed_dir, 'val', hash_dir) test_dir = osp.join(self.processed_dir, 'test', hash_dir) if not osp.exists(train_dir): os.makedirs(train_dir, exist_ok=True) if not osp.exists(val_dir): if self.val_mixed_in_train: os.makedirs(osp.dirname(val_dir), exist_ok=True) os.symlink(train_dir, val_dir, target_is_directory=True) else: os.makedirs(val_dir, exist_ok=True) if not osp.exists(test_dir): if self.test_mixed_in_val: os.makedirs(osp.dirname(test_dir), exist_ok=True) os.symlink(val_dir, test_dir, target_is_directory=True) else: os.makedirs(test_dir, exist_ok=True) # Process clouds one by one for p in tq(self.processed_paths): self._process_single_cloud(p) def _process_single_cloud(self, cloud_path: str) -> None: """Internal method called by `self.process` to preprocess a single cloud of 3D points. """ # If required files exist, skip processing if osp.exists(cloud_path): return # Create necessary parent folders if need be os.makedirs(osp.dirname(cloud_path), exist_ok=True) # Read the raw cloud corresponding to the final processed # `cloud_path` and convert it to a Data object raw_path = self.processed_to_raw_path(cloud_path) data = self.sanitized_read_single_raw_cloud(raw_path) # If the cloud path indicates a tiling is needed, apply it here if self.xy_tiling is not None: tile = self.get_tile_from_path(cloud_path)[0] data = SampleXYTiling(x=tile[0], y=tile[1], tiling=tile[2])(data) elif self.pc_tiling is not None: tile = self.get_tile_from_path(cloud_path)[0] data = SampleRecursiveMainXYAxisTiling(x=tile[0], steps=tile[1])(data) # Apply pre_transform if self.pre_transform is not None: nag = self.pre_transform(data) else: nag = NAG([data]) # To save some disk space, we discard some level-0 attributes if self.point_save_keys is not None: keys = set(nag[0].keys) - set(self.point_save_keys) nag = NAGRemoveKeys(level=0, keys=keys)(nag) elif self.point_no_save_keys is not None: nag = NAGRemoveKeys(level=0, keys=self.point_no_save_keys)(nag) if self.segment_save_keys is not None: keys = set(nag[1].keys) - set(self.segment_save_keys) nag = NAGRemoveKeys(level='1+', keys=keys)(nag) elif self.segment_no_save_keys is not None: nag = NAGRemoveKeys(level='1+', keys=self.segment_no_save_keys)(nag) # Save pre_transformed data to the processed dir/ # TODO: is you do not throw away level-0 neighbors, make sure # that they contain no '-1' empty neighborhoods, because if # you load them for batching, the pyg reindexing mechanism will # break indices will not index update nag.save( cloud_path, y_to_csr=self.save_y_to_csr, pos_dtype=self.save_pos_dtype, fp_dtype=self.save_fp_dtype) del nag @staticmethod def get_tile_from_path(path: str) -> Tuple[Tuple, str, str]: # Search the XY tiling suffix pattern out_reg = re.search('__TILE_(\d+)-(\d+)_OF_(\d+)-(\d+)', path) if out_reg is not None: x, y, x_tiling, y_tiling = [int(g) for g in out_reg.groups()] suffix = f'__TILE_{x}-{y}_OF_{x_tiling}-{y_tiling}' prefix = path.replace(suffix, '') return (x - 1, y - 1, (x_tiling, y_tiling)), prefix, suffix # Search the PC tiling suffix pattern out_reg = re.search('__TILE_(\d+)_OF_(\d+)', path) if out_reg is not None: x, num = [int(g) for g in out_reg.groups()] suffix = f'__TILE_{x}_OF_{num}' prefix = path.replace(suffix, '') steps = torch.log2(torch.tensor(num)).int().item() return (x - 1, steps), prefix, suffix return def read_single_raw_cloud(self, raw_cloud_path: str) -> 'Data': """Read a single raw cloud and return a `Data` object, ready to be passed to `self.pre_transform`. This `Data` object should contain the following attributes: - `pos`: point coordinates - `y`: OPTIONAL point semantic label - `obj`: OPTIONAL `InstanceData` object with instance labels - `rgb`: OPTIONAL point color - `intensity`: OPTIONAL point LiDAR intensity IMPORTANT: By convention, we assume `y ∈ [0, self.num_classes-1]` ARE ALL VALID LABELS (i.e. not 'ignored', 'void', 'unknown', etc), while `y < 0` AND `y >= self.num_classes` ARE VOID LABELS. This applies to both `Data.y` and `Data.obj.y`. """ raise NotImplementedError def sanitized_read_single_raw_cloud(self, raw_cloud_path: str) -> 'Data': """Wrapper around the actual `self.read_single_raw_cloud`. This function ensures that the semantic and instance segmentation labels returned by the reader are sanitized. More specifically, we assume `[0, self.num_classes-1]` ARE ALL VALID LABELS (i.e. not 'ignored', 'void', 'unknown', etc), while `y < 0` AND `y >= self.num_classes` ARE VOID LABELS. To this end, this function maps all labels outside `[0, self.num_classes-1]` to `y = self.num_classes`. Hence, we actually have `self.num_classes + 1` labels in the data. This allows identifying the points to be ignored at metric computation time. Besides, this function ensures that there is at most 1 instance of each stuff (and void) class in each scene/cloud/tile, as described in: - https://arxiv.org/abs/1801.00868 - https://arxiv.org/abs/1905.01220 """ data = self.read_single_raw_cloud(raw_cloud_path) # Set all void labels to self.num_classes in the semantic # segmentation labels if getattr(data, 'y', None) is not None: data.y[data.y < 0] = self.num_classes data.y[data.y > self.num_classes] = self.num_classes # Set all void labels to self.num_classes in the # instance/panoptic segmentation annotations if getattr(data, 'obj', None) is not None: data.obj.y[data.obj.y < 0] = self.num_classes data.obj.y[data.obj.y > self.num_classes] = self.num_classes # For each cloud/scene and each stuff/void class, group # annotations into a single instance for i in self.stuff_classes + self.void_classes: idx = torch.where(data.obj.y == i)[0] if idx.numel() == 0: continue data.obj.obj[idx] = data.obj.obj[idx].min() return data def debug_instance_data(self, level: int = 1) -> None: """Sanity check to make sure at most 1 instance of each stuff class per scene/cloud. :param level: int NAG level which to inspect """ problematic_clouds = [] for i_cloud, nag in tqdm(enumerate(self)): _, perm = consecutive_cluster(nag[level].obj.obj) y = nag[level].obj.y[perm] y_count = torch.bincount(y, minlength=self.num_classes + 1) for c in self.stuff_classes + self.void_classes: if y_count[c] > 1: problematic_clouds.append(i_cloud) break assert len(problematic_clouds) == 0, \ f"The following clouds have more than 1 instance of for a stuff " \ f"or void class:\n{problematic_clouds}" def get_class_weight(self, smooth: str='sqrt') -> torch.Tensor: """Compute class weights based on the labels distribution in the dataset. Optionally a 'smooth' function may be passed to smoothen the weights' statistics. """ assert smooth in [None, 'sqrt', 'log'] # Read the first NAG just to know how many levels we have in the # preprocessed NAGs. nag = self[0] low = nag.num_levels - 1 # Make sure the dataset has labels if nag[low].y is None: return None del nag # To be as fast as possible, we read only the last level of each # NAG, and accumulate the class counts from the label histograms counts = torch.zeros(self.num_classes) for i in range(len(self)): if self.in_memory: y = self.in_memory_data[i][low].y else: y = NAG.load( self.processed_paths[i], low=low, keys_low=['y'])[0].y counts += y.sum(dim=0)[:self.num_classes] # Compute the class weights. Optionally, a 'smooth' function may # be applied to smoothen the weights statistics if smooth == 'sqrt': counts = counts.sqrt() if smooth == 'log': counts = counts.log() weights = 1 / (counts + 1) weights /= weights.sum() return weights def __len__(self) -> int: """Number of clouds in the dataset.""" return len(self.cloud_ids) def __getitem__(self, idx: int) -> Union['NAG', 'Data']: """Load a preprocessed NAG from disk and apply `self.transform` if any. Optionally, one may pass a tuple (idx, bool) where the boolean indicates whether the data should be loaded from disk, if `self.in_memory=True`. """ # Prepare from_hdd from_hdd = False if isinstance(idx, tuple): assert len(idx) == 2 and isinstance(idx[1], bool), \ "Only supports indexing with `int` or `(int, bool)` where the" \ " boolean indicates whether the data should be loaded from " \ "disk, when `self.in_memory=True`." idx, from_hdd = idx # Get the processed NAG directly from RAM if self.in_memory and not from_hdd: # TODO: careful, this means the transforms are only run # once. So no augmentations, samplings, etc in the # transforms... return self.in_memory_data[idx] # Read the NAG from HDD nag = NAG.load( self.processed_paths[idx], keys_low=self.point_load_keys, keys=self.segment_load_keys) # Apply transforms nag = nag if self.transform is None else self.transform(nag) return nag def make_submission( self, idx: int, pred: torch.Tensor, pos: torch.Tensor, submission_dir: str = None ) -> None: """Implement this if your dataset needs to produce data in a given format for submission. This is typically needed for datasets with held-out test sets. """ raise NotImplementedError def finalize_submission(self, submission_dir: str) -> None: """Implement this if your dataset needs to produce data in a given format for submission. This is typically needed for datasets with held-out test sets. """ raise NotImplementedError def show_examples( self, label: int, radius: float = 4, max_examples: int = 5, shuffle: bool = True, **kwargs ) -> None: """Interactive plots of some examples centered on points of the provided `label`. At most one example per cloud/tile/scene in the dataset will be shown. :param label: int or str Label of the class of interest, may be provided as an int or a string corresponding to the class name :param radius: float Radius of the spherical sampling to draw around the point of interest :param max_examples: int Maximum number of samples to draw :param shuffle: bool If True, the candidate samples will be shuffled every time :param kwargs: Kwargs to be passed to the visualization `show()` function :return: """ if isinstance(label, str): assert label in self.class_names, \ f"Label must be within {self.class_names}]" label = self.class_names.index(label) assert label >= 0 and label <= self.num_classes, \ f"Label must be within [0, {self.num_classes + 1}]" # Gather some clouds ids with the desired class cloud_list = [] iterator = list(range(len(self))) if shuffle: random.shuffle(iterator) for i_cloud in iterator: if len(cloud_list) >= max_examples: break if (self[i_cloud][1].y.argmax(dim=1) == label).any(): cloud_list.append(i_cloud) # If no cloud was found with the desired class, return here if len(cloud_list) == 0: print( f"Could not find any cloud with points of label={label} in the " f"dataset.") return # Display some found examples for i, i_cloud in enumerate(cloud_list): if i >= max_examples: break # Load the cloud nag = self[i_cloud] # Search for points with the desired label point_idx = torch.where(nag[0].y.argmax(dim=1) == label)[0].tolist() # Pick only on of the points as visualization center for the # cloud at hand if shuffle: random.shuffle(point_idx) i_point = point_idx[0] # Draw the scene center = nag[0].pos[i_point].cpu().tolist() title = f"Label={label} - Cloud={i_cloud} - Center={center}" print(f"\n{title}") show( nag, center=center, radius=radius, title=title, class_names=self.class_names, class_colors=self.class_colors, stuff_classes=self.stuff_classes, num_classes=self.num_classes, **kwargs)