base_dataset.py

import os
from abc import ABC, abstractmethod
import logging
import functools
from functools import partial
import numpy as np
import torch
import torch_geometric
from torch_geometric.transforms import Compose, FixedPoints
import copy

from torch_points3d.models import model_interface
from torch_points3d.core.data_transform import instantiate_transforms, MultiScaleTransform
from torch_points3d.core.data_transform import instantiate_filters
from torch_points3d.datasets.batch import SimpleBatch
from torch_points3d.datasets.multiscale_data import MultiScaleBatch
from torch_points3d.utils.enums import ConvolutionFormat
from torch_points3d.utils.config import ConvolutionFormatFactory
from torch_points3d.utils.colors import COLORS, colored_print

# A logger for this file
log = logging.getLogger(__name__)


def explode_transform(transforms):
    """ Returns a flattened list of transform
    Arguments:
        transforms {[list | T.Compose]} -- Contains list of transform to be added

    Returns:
        [list] -- [List of transforms]
    """
    out = []
    if transforms is not None:
        if isinstance(transforms, Compose):
            out = copy.deepcopy(transforms.transforms)
        elif isinstance(transforms, list):
            out = copy.deepcopy(transforms)
        else:
            raise Exception("Transforms should be provided either within a list or a Compose")
    return out


def save_used_properties(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        # Save used_properties for mocking dataset when calling pretrained registry
        result = func(self, *args, **kwargs)
        if isinstance(result, torch.Tensor):
            self.used_properties[func.__name__] = result.numpy().tolist()
        elif isinstance(result, np.ndarray):
            self.used_properties[func.__name__] = result.tolist()
        else:
            self.used_properties[func.__name__] = result
        return result

    return wrapper


class BaseDataset:
    def __init__(self, dataset_opt):
        self.dataset_opt = dataset_opt

        # Default dataset path
        dataset_name = dataset_opt.get("dataset_name", None)
        if dataset_name:
            self._data_path = os.path.join(dataset_opt.dataroot, dataset_name)
        else:
            class_name = self.__class__.__name__.lower().replace("dataset", "")
            self._data_path = os.path.join(dataset_opt.dataroot, class_name)
        self._batch_size = None
        self.strategies = {}
        self._contains_dataset_name = False

        self.train_sampler = None
        self.test_sampler = None
        self.val_sampler = None

        self._train_dataset = None
        self._test_dataset = None
        self._val_dataset = None

        self.train_pre_batch_collate_transform = None
        self.val_pre_batch_collate_transform = None
        self.test_pre_batch_collate_transform = None

        BaseDataset.set_transform(self, dataset_opt)
        self.set_filter(dataset_opt)

        self.used_properties = {}

    @staticmethod
    def remove_transform(transform_in, list_transform_class):
        """ Remove a transform if within list_transform_class

        Arguments:
            transform_in {[type]} -- [Compose | List of transform]
            list_transform_class {[type]} -- [List of transform class to be removed]

        Returns:
            [type] -- [description]
        """
        if isinstance(transform_in, Compose) or isinstance(transform_in, list):
            if len(list_transform_class) > 0:
                transform_out = []
                transforms = transform_in.transforms if isinstance(transform_in, Compose) else transform_in
                for t in transforms:
                    if not isinstance(t, tuple(list_transform_class)):
                        transform_out.append(t)
                transform_out = Compose(transform_out)
        else:
            transform_out = transform_in
        return transform_out

    @staticmethod
    def set_transform(obj, dataset_opt):
        """This function create and set the transform to the obj as attributes
        """
        obj.pre_transform = None
        obj.test_transform = None
        obj.train_transform = None
        obj.val_transform = None
        obj.inference_transform = None

        for key_name in dataset_opt.keys():
            if "transform" in key_name:
                new_name = key_name.replace("transforms", "transform")
                try:
                    transform = instantiate_transforms(getattr(dataset_opt, key_name))
                except Exception:
                    log.exception("Error trying to create {}, {}".format(new_name, getattr(dataset_opt, key_name)))
                    continue
                setattr(obj, new_name, transform)

        inference_transform = explode_transform(obj.pre_transform)
        inference_transform += explode_transform(obj.test_transform)
        obj.inference_transform = Compose(inference_transform) if len(inference_transform) > 0 else None

    def set_filter(self, dataset_opt):
        """This function create and set the pre_filter to the obj as attributes
        """
        self.pre_filter = None
        for key_name in dataset_opt.keys():
            if "filter" in key_name:
                new_name = key_name.replace("filters", "filter")
                try:
                    filt = instantiate_filters(getattr(dataset_opt, key_name))
                except Exception:
                    log.exception("Error trying to create {}, {}".format(new_name, getattr(dataset_opt, key_name)))
                    continue
                setattr(self, new_name, filt)

    @staticmethod
    def _collate_fn(batch, collate_fn=None, pre_collate_transform=None):
        if pre_collate_transform:
            batch = pre_collate_transform(batch)
        return collate_fn(batch)

    @staticmethod
    def _get_collate_function(conv_type, is_multiscale, pre_collate_transform=None):
        is_dense = ConvolutionFormatFactory.check_is_dense_format(conv_type)
        if is_multiscale:
            if conv_type.lower() == ConvolutionFormat.PARTIAL_DENSE.value.lower():
                fn = MultiScaleBatch.from_data_list
            else:
                raise NotImplementedError(
                    "MultiscaleTransform is activated and supported only for partial_dense format"
                )
        else:
            if is_dense:
                fn = SimpleBatch.from_data_list
            else:
                fn = torch_geometric.data.batch.Batch.from_data_list
        return partial(BaseDataset._collate_fn, collate_fn=fn, pre_collate_transform=pre_collate_transform)

    @staticmethod
    def get_num_samples(batch, conv_type):
        is_dense = ConvolutionFormatFactory.check_is_dense_format(conv_type)
        if is_dense:
            return batch.pos.shape[0]
        else:
            return batch.batch.max() + 1

    @staticmethod
    def get_sample(batch, key, index, conv_type):
        assert hasattr(batch, key)
        is_dense = ConvolutionFormatFactory.check_is_dense_format(conv_type)
        if is_dense:
            return batch[key][index]
        else:
            return batch[key][batch.batch == index]

    def create_dataloaders(
        self,
        model: model_interface.DatasetInterface,
        batch_size: int,
        shuffle: bool,
        num_workers: int,
        precompute_multi_scale: bool,
    ):
        """ Creates the data loaders. Must be called in order to complete the setup of the Dataset
        """
        conv_type = model.conv_type
        self._batch_size = batch_size

        if self.train_sampler:
            log.info(self.train_sampler)

        if self.train_dataset:
            self._train_loader = self._dataloader(
                self.train_dataset,
                self.train_pre_batch_collate_transform,
                conv_type,
                precompute_multi_scale,
                batch_size=batch_size,
                shuffle=shuffle and not self.train_sampler,
                num_workers=num_workers,
                sampler=self.train_sampler,
            )

        if self.test_dataset:
            self._test_loaders = [
                self._dataloader(
                    dataset,
                    self.test_pre_batch_collate_transform,
                    conv_type,
                    precompute_multi_scale,
                    batch_size=batch_size,
                    shuffle=False,
                    num_workers=num_workers,
                    sampler=self.test_sampler,
                )
                for dataset in self.test_dataset
            ]

        if self.val_dataset:
            self._val_loader = self._dataloader(
                self.val_dataset,
                self.val_pre_batch_collate_transform,
                conv_type,
                precompute_multi_scale,
                batch_size=batch_size,
                shuffle=False,
                num_workers=num_workers,
                sampler=self.val_sampler,
            )

        if precompute_multi_scale:
            self.set_strategies(model)

    def _dataloader(self, dataset, pre_batch_collate_transform, conv_type, precompute_multi_scale, **kwargs):
        batch_collate_function = self.__class__._get_collate_function(
            conv_type, precompute_multi_scale, pre_batch_collate_transform
        )
        num_workers = kwargs.get("num_workers", 0)
        persistent_workers = (num_workers > 0)
        dataloader = partial(
            torch.utils.data.DataLoader, collate_fn=batch_collate_function, worker_init_fn=np.random.seed,
            persistent_workers=persistent_workers
        )
        return dataloader(dataset, **kwargs)

    @property
    def has_train_loader(self):
        return hasattr(self, "_train_loader")

    @property
    def has_val_loader(self):
        return hasattr(self, "_val_loader")

    @property
    def has_test_loaders(self):
        return hasattr(self, "_test_loaders")

    @property
    def train_dataset(self):
        return self._train_dataset

    @train_dataset.setter
    def train_dataset(self, value):
        self._train_dataset = value
        if not hasattr(self._train_dataset, "name"):
            setattr(self._train_dataset, "name", "train")

    @property
    def val_dataset(self):
        return self._val_dataset

    @val_dataset.setter
    def val_dataset(self, value):
        self._val_dataset = value
        if not hasattr(self._val_dataset, "name"):
            setattr(self._val_dataset, "name", "val")

    @property
    def test_dataset(self):
        return self._test_dataset

    @test_dataset.setter
    def test_dataset(self, value):
        if isinstance(value, list):
            self._test_dataset = value
        else:
            self._test_dataset = [value]

        for i, dataset in enumerate(self._test_dataset):
            if not hasattr(dataset, "name"):
                if self.num_test_datasets > 1:
                    setattr(dataset, "name", "test_%i" % i)
                else:
                    setattr(dataset, "name", "test")
            else:
                self._contains_dataset_name = True

        # Check for uniqueness
        all_names = [d.name for d in self.test_dataset]
        if len(set(all_names)) != len(all_names):
            raise ValueError("Datasets need to have unique names. Current names are {}".format(all_names))

    @property
    def train_dataloader(self):
        return self._train_loader

    @property
    def val_dataloader(self):
        return self._val_loader

    @property
    def test_dataloaders(self):
        if self.has_test_loaders:
            return self._test_loaders
        else:
            return []

    @property
    def _loaders(self):
        loaders = []
        if self.has_train_loader:
            loaders += [self.train_dataloader]
        if self.has_val_loader:
            loaders += [self.val_dataloader]
        if self.has_test_loaders:
            loaders += self.test_dataloaders
        return loaders

    @property
    def num_test_datasets(self):
        return len(self._test_dataset) if self._test_dataset else 0

    @property
    def _test_datatset_names(self):
        if self.test_dataset:
            return [d.name for d in self.test_dataset]
        else:
            return []

    @property
    def available_stage_names(self):
        out = self._test_datatset_names
        if self.has_val_loader:
            out += [self._val_dataset.name]
        return out

    @property
    def available_dataset_names(self):
        return ["train"] + self.available_stage_names

    def get_raw_data(self, stage, idx, **kwargs):
        assert stage in self.available_dataset_names
        dataset = self.get_dataset(stage)
        if hasattr(dataset, "get_raw_data"):
            return dataset.get_raw_data(idx, **kwargs)
        else:
            raise Exception("Dataset {} doesn t have a get_raw_data function implemented".format(dataset))

    def has_labels(self, stage: str) -> bool:
        """ Tests if a given dataset has labels or not

        Parameters
        ----------
        stage : str
            name of the dataset to test
        """
        assert stage in self.available_dataset_names
        dataset = self.get_dataset(stage)
        if hasattr(dataset, "has_labels"):
            return dataset.has_labels

        sample = dataset[0]
        if hasattr(sample, "y"):
            return sample.y is not None
        return False

    @property  # type: ignore
    @save_used_properties
    def is_hierarchical(self):
        """ Used by the metric trackers to log hierarchical metrics
        """
        return False

    @property  # type: ignore
    @save_used_properties
    def class_to_segments(self):
        """ Use this property to return the hierarchical map between classes and segment ids, example:
        {
            'Airplaine': [0,1,2],
            'Boat': [3,4,5]
        }
        """
        return None

    @property  # type: ignore
    @save_used_properties
    def num_classes(self):
        return self.train_dataset.num_classes

    @property
    def weight_classes(self):
        return getattr(self.train_dataset, "weight_classes", None)

    @property  # type: ignore
    @save_used_properties
    def feature_dimension(self):
        if self.train_dataset:
            return self.train_dataset.num_features
        elif self.test_dataset is not None:
            if isinstance(self.test_dataset, list):
                return self.test_dataset[0].num_features
            else:
                return self.test_dataset.num_features
        elif self.val_dataset is not None:
            return self.val_dataset.num_features
        else:
            raise NotImplementedError()

    @property
    def batch_size(self):
        return self._batch_size

    @property
    def num_batches(self):
        out = {
            self.train_dataset.name: len(self._train_loader),
            "val": len(self._val_loader) if self.has_val_loader else 0,
        }
        if self.test_dataset:
            for loader in self._test_loaders:
                stage_name = loader.dataset.name
                out[stage_name] = len(loader)
        return out

    def get_dataset(self, name):
        """ Get a dataset by name. Raises an exception if no dataset was found

        Parameters
        ----------
        name : str
        """
        all_datasets = [self.train_dataset, self.val_dataset]
        if self.test_dataset:
            all_datasets += self.test_dataset
        for dataset in all_datasets:
            if dataset is not None and dataset.name == name:
                return dataset
        raise ValueError("No dataset with name %s was found." % name)

    def _set_composed_multiscale_transform(self, attr, transform):
        current_transform = getattr(attr.dataset, "transform", None)
        if current_transform is None:
            setattr(attr.dataset, "transform", transform)
        else:
            if (
                isinstance(current_transform, Compose) and transform not in current_transform.transforms
            ):  # The transform contains several transformations
                current_transform.transforms += [transform]
            elif current_transform != transform:
                setattr(
                    attr.dataset, "transform", Compose([current_transform, transform]),
                )

    def _set_multiscale_transform(self, transform):
        for _, attr in self.__dict__.items():
            if isinstance(attr, torch.utils.data.DataLoader):
                self._set_composed_multiscale_transform(attr, transform)
        for loader in self.test_dataloaders:
            self._set_composed_multiscale_transform(loader, transform)

    def set_strategies(self, model):
        strategies = model.get_spatial_ops()
        transform = MultiScaleTransform(strategies)
        self._set_multiscale_transform(transform)

    @abstractmethod
    def get_tracker(self, wandb_log: bool, tensorboard_log: bool):
        pass

    def resolve_saving_stage(self, selection_stage):
        """This function is responsible to determine if the best model selection
        is going to be on the validation or test datasets
        """
        log.info(
            "Available stage selection datasets: {} {} {}".format(
                COLORS.IPurple, self.available_stage_names, COLORS.END_NO_TOKEN
            )
        )

        if self.num_test_datasets > 1 and not self._contains_dataset_name:
            msg = "If you want to have better trackable names for your test datasets, add a "
            msg += COLORS.IPurple + "name" + COLORS.END_NO_TOKEN
            msg += " attribute to them"
            log.info(msg)

        if selection_stage == "":
            if self.has_val_loader:
                selection_stage = self.val_dataset.name
            else:
                selection_stage = self.test_dataset[0].name
        log.info(
            "The models will be selected using the metrics on following dataset: {} {} {}".format(
                COLORS.IPurple, selection_stage, COLORS.END_NO_TOKEN
            )
        )
        return selection_stage

    def add_weights(self, dataset_name="train", class_weight_method="sqrt"):
        """ Add class weights to a given dataset that are then accessible using the `class_weights` attribute
        """
        L = self.num_classes
        weights = torch.ones(L)
        dataset = self.get_dataset(dataset_name)
        idx_classes, counts = torch.unique(dataset.data.y, return_counts=True)

        dataset.idx_classes = torch.arange(L).long()
        weights[idx_classes] = counts.float()
        weights = weights.float()
        weights = weights.mean() / weights
        if class_weight_method == "sqrt":
            weights = torch.sqrt(weights)
        elif str(class_weight_method).startswith("log"):
            weights = torch.log(1.1 + weights / weights.sum())
        else:
            raise ValueError("Method %s not supported" % class_weight_method)

        weights /= torch.sum(weights)
        log.info("CLASS WEIGHT : {}".format([np.round(weight.item(), 4) for weight in weights]))
        setattr(dataset, "weight_classes", weights)

        return dataset

    def __repr__(self):
        message = "Dataset: %s \n" % self.__class__.__name__
        for attr in self.__dict__:
            if "transform" in attr:
                message += "{}{} {}= {}\n".format(COLORS.IPurple, attr, COLORS.END_NO_TOKEN, getattr(self, attr))
        for attr in self.__dict__:
            if attr.endswith("_dataset"):
                dataset = getattr(self, attr)
                if isinstance(dataset, list):
                    if len(dataset) > 1:
                        size = ", ".join([str(len(d)) for d in dataset])
                    else:
                        size = len(dataset[0])
                elif dataset:
                    size = len(dataset)
                else:
                    size = 0
                if attr.startswith("_"):
                    attr = attr[1:]
                message += "Size of {}{} {}= {}\n".format(COLORS.IPurple, attr, COLORS.END_NO_TOKEN, size)
        for key, attr in self.__dict__.items():
            if key.endswith("_sampler") and attr:
                message += "{}{} {}= {}\n".format(COLORS.IPurple, key, COLORS.END_NO_TOKEN, attr)
        message += "{}Batch size ={} {}".format(COLORS.IPurple, COLORS.END_NO_TOKEN, self.batch_size)
        return message