torch-pointcloud 0.0.1__py3-none-any.whl

torch_pointcloud/__init__.py

@@ -0,0 +1,23 @@
+"""PyTorch library for 3D point cloud deep learning: models, datasets, transforms, and pretrained weights."""
+
+from importlib.metadata import version as _version
+
+from . import config, datasets, inferers, layers, losses, models, transforms, utils
+from .models import create_model, list_models, register_model
+
+__version__ = _version("torch_pointcloud")
+
+__all__ = [
+    "__version__",
+    "config",
+    "create_model",
+    "datasets",
+    "inferers",
+    "layers",
+    "list_models",
+    "losses",
+    "models",
+    "register_model",
+    "transforms",
+    "utils",
+]

torch_pointcloud/config.py

@@ -0,0 +1,63 @@
+"""Environment-variable configuration of cache, model, and data directories and randomness defaults.
+
+This module contains the following global variables for configuration:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `HOME_DIR` | The home directory of the user. | `Path.home().as_posix()` |
+| `CACHE_DIR` | The cache directory for the package. | `Path(HOME_DIR, ".cache", "torch-pointcloud").as_posix()` |
+| `MODELS_DIR` | The directory for the models. | `Path(CACHE_DIR, "models").as_posix()` |
+| `DATA_DIR` | The directory for the data. | `"data"` |
+| `RANDOM_SEED` | The random seed for the package. | `None` |
+| `FPS_RANDOM_START` | Whether to start the random seed from the current time. | `None` |
+| `KNN_DENSE_BUDGET` | The dense budget for the KNN. | `16_000_000` |
+"""
+
+import os
+from pathlib import Path
+from typing import Optional
+
+
+def asbool(value: Optional[str]) -> Optional[bool]:
+    """Parse an environment variable as a boolean.
+
+    Args:
+        value: The raw variable value, or None when it is unset.
+
+    Returns:
+        True when the value reads as `true`, `1`, `yes` or `y` (case-insensitive), False for anything
+        else, and None when the variable is unset.
+    """
+    if value is None:
+        return None
+
+    return value.strip().lower() in ["true", "1", "yes", "y"]
+
+
+def asint(value: Optional[str]) -> Optional[int]:
+    """Parse an environment variable as an integer.
+
+    Args:
+        value: The raw variable value, or None when it is unset.
+
+    Returns:
+        The parsed integer, or None when the variable is unset or does not parse.
+    """
+    if value is None:
+        return None
+
+    try:
+        return int(value)
+    except ValueError:
+        return None
+
+
+HOME_DIR = Path.home().as_posix()
+CACHE_DIR = os.getenv("TORCH_POINTCLOUD_CACHE_DIR", Path(HOME_DIR, ".cache", "torch-pointcloud").as_posix())
+MODELS_DIR = os.getenv("TORCH_POINTCLOUD_MODELS_DIR", Path(CACHE_DIR, "models").as_posix())
+DATA_DIR = os.getenv("TORCH_POINTCLOUD_DATA_DIR", "data")
+
+# Some variables to affect how random operations are performed.
+RANDOM_SEED = asint(os.getenv("TORCH_POINTCLOUD_RANDOM_SEED", None))
+FPS_RANDOM_START = asbool(os.getenv("TORCH_POINTCLOUD_FPS_RANDOM_START", None))
+KNN_DENSE_BUDGET = asint(os.getenv("TORCH_POINTCLOUD_KNN_DENSE_BUDGET", None)) or 16_000_000

torch_pointcloud/datasets/__init__.py

@@ -0,0 +1,18 @@
+"""Benchmark point cloud datasets with disk caching and download helpers."""
+
+from .concat import ConcatDataset, SingleDatasetBatchSampler
+from .kitti import KITTI
+from .mix import MixDataset
+from .modelnet import ModelNet10, ModelNet40, ModelNet40Hdf5, ModelNetNormalResampled
+from .nuscenes import NuScenes, NuScenesMini
+from .parislille3d import ParisLille3D
+from .pointcloud import PointCloudDataset
+from .repeat import RepeatDataset
+from .s3dis import S3DIS, S3DISHdf5
+from .scannet import ScanNet, ScanNet20, ScanNet200
+from .scanobjectnn import ScanObjectNN
+from .semantic3d import Semantic3D
+from .semantickitti import SemanticKITTI
+from .shapenetpart import ShapeNetPart
+from .sunrgbd import SunRGBD
+from .toronto3d import Toronto3D

torch_pointcloud/datasets/concat.py

@@ -0,0 +1,161 @@
+"""Concatenate several datasets into one flat index space for multi-dataset joint training."""
+
+from bisect import bisect_right
+from typing import Any, Iterator, List, Optional, Sequence
+
+import torch
+from torch.utils.data import Dataset, Sampler
+
+
+class ConcatDataset(Dataset):
+    """Concatenates several datasets into one flat index space.
+
+    Each child dataset keeps its own `transform`, so datasets from different domains (e.g. ScanNet and
+    S3DIS) train jointly while stamping their own condition key and mapping to their native label space.
+    Pair it with `SingleDatasetBatchSampler` to keep every batch single-domain so per-dataset
+    normalization statistics (BatchNorm, PDNorm) stay clean.
+
+    Args:
+        datasets: The datasets to concatenate, in order. The first is treated as the main dataset by
+            `SingleDatasetBatchSampler` (its exhaustion ends the epoch).
+
+    Example:
+        ```python
+        from torch_pointcloud.datasets import ConcatDataset, S3DIS, ScanNet20
+
+        dataset = ConcatDataset([ScanNet20(root, split="train"), S3DIS(root, areas=["Area_1"])])
+        len(dataset)  # len(scannet) + len(s3dis)
+        dataset.sizes  # [len(scannet), len(s3dis)]
+        ```
+    """
+
+    def __init__(self, datasets: Sequence[Dataset]) -> None:
+        if len(datasets) == 0:
+            raise ValueError("ConcatDataset requires at least one dataset.")
+        self.datasets = list(datasets)
+        self.sizes = [len(d) for d in self.datasets]  # type: ignore[arg-type]
+        self.cumulative_sizes: List[int] = []
+        total = 0
+        for size in self.sizes:
+            total += size
+            self.cumulative_sizes.append(total)
+
+    def __len__(self) -> int:
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, index: int) -> Any:
+        if not -len(self) <= index < len(self):
+            raise IndexError(f"Index {index} is out of range for a dataset of length {len(self)}.")
+        if index < 0:
+            index += len(self)
+        dataset_index = bisect_right(self.cumulative_sizes, index)
+        start = self.cumulative_sizes[dataset_index - 1] if dataset_index > 0 else 0
+        return self.datasets[dataset_index][index - start]
+
+
+class SingleDatasetBatchSampler(Sampler[List[int]]):
+    """Yields batches whose global indices all come from one dataset of a `ConcatDataset`.
+
+    Given the per-dataset sizes and one positive integer ratio per dataset, it partitions each dataset's
+    indices into batches of `batch_size` and interleaves the datasets round-robin weighted by the ratios.
+    Within a round the first dataset yields `ratios[0]` batches, the second `ratios[1]`, and so on. The
+    first (main) dataset drives the epoch length: when it
+    is exhausted the epoch ends, while the other datasets restart (reshuffled) as needed. Because every
+    yielded batch is drawn from a single dataset, per-batch normalization (BatchNorm, PDNorm) sees a
+    single domain.
+
+    Args:
+        sizes: Number of samples in each child dataset, in `ConcatDataset` order.
+        ratios: One positive integer sampling weight per dataset, aligned with `sizes`.
+        batch_size: Number of indices per batch.
+        shuffle: Shuffle each dataset's indices, reshuffling on restart.
+        drop_last: Drop each dataset's trailing partial batch.
+        generator: Optional `torch.Generator` for shuffling.
+
+    Shape:
+        Each yielded value is a `List[int]` of length `batch_size` (or fewer for a trailing batch when
+        `drop_last` is `False`).
+
+    Example:
+        ```python
+        from torch.utils.data import DataLoader
+
+        from torch_pointcloud.datasets import ConcatDataset, SingleDatasetBatchSampler
+
+        dataset = ConcatDataset([scannet, s3dis])
+        sampler = SingleDatasetBatchSampler(dataset.sizes, ratios=[2, 1], batch_size=4)
+        loader = DataLoader(dataset, batch_sampler=sampler)
+        ```
+    """
+
+    def __init__(
+        self,
+        sizes: Sequence[int],
+        ratios: Sequence[int],
+        batch_size: int,
+        shuffle: bool = True,
+        drop_last: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> None:
+        if len(sizes) != len(ratios):
+            raise ValueError(f"sizes and ratios must have the same length, got {len(sizes)} and {len(ratios)}.")
+        if any(r <= 0 for r in ratios):
+            raise ValueError(f"ratios must be positive integers, got {list(ratios)}.")
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {batch_size}.")
+
+        self.sizes = list(sizes)
+        self.ratios = list(ratios)
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.generator = generator
+        self.offsets: List[int] = [sum(self.sizes[:i]) for i in range(len(self.sizes))]
+
+    def _num_batches(self, size: int) -> int:
+        if self.drop_last:
+            return size // self.batch_size
+        return (size + self.batch_size - 1) // self.batch_size
+
+    def _batches(self, dataset_index: int) -> List[List[int]]:
+        size = self.sizes[dataset_index]
+        offset = self.offsets[dataset_index]
+        if self.shuffle:
+            order = torch.randperm(size, generator=self.generator).tolist()
+        else:
+            order = list(range(size))
+
+        indices = [offset + i for i in order]
+        batches = [indices[start : start + self.batch_size] for start in range(0, size, self.batch_size)]
+
+        if self.drop_last and batches and len(batches[-1]) < self.batch_size:
+            batches.pop()
+        return batches
+
+    def __len__(self) -> int:
+        main = self._num_batches(self.sizes[0])
+        if main == 0:
+            return 0
+
+        full_rounds, remainder = divmod(main, self.ratios[0])
+        per_round = self.ratios[0] + sum(
+            self.ratios[i] for i in range(1, len(self.sizes)) if self._num_batches(self.sizes[i]) > 0
+        )
+        return full_rounds * per_round + remainder
+
+    def __iter__(self) -> Iterator[List[int]]:
+        iterators = [iter(self._batches(i)) for i in range(len(self.sizes))]
+        while True:
+            for dataset_index, ratio in enumerate(self.ratios):
+                for _ in range(ratio):
+                    batch = next(iterators[dataset_index], None)
+                    if batch is None:
+                        if dataset_index == 0:
+                            return
+
+                        iterators[dataset_index] = iter(self._batches(dataset_index))
+                        batch = next(iterators[dataset_index], None)
+                        if batch is None:
+                            break
+
+                    yield batch