Initial release: PT-v3m1 segmentation model (BIMStruct3D pipeline)

.gitignore

@@ -0,0 +1,3 @@
+.venv/
+__pycache__/
+*.pyc

README.md

@@ -0,0 +1,154 @@
+---
+library_name: pytorch
+tags:
+  - point-cloud
+  - semantic-segmentation
+  - point-transformer-v3
+  - lidar
+  - 3d
+datasets:
+  - RPTU-FGMB/DeKH
+license: cc-by-nc-sa-4.0
+---
+
+# BIMStruct3D-segmentation
+
+Point-cloud semantic segmentation model (PT-v3m1) from the
+[BIMStruct3D](#citation) scan-to-BIM pipeline.
+
+[![License: CC BY-NC-SA 4.0](https://img.shields.io/badge/License-CC%20BY--NC--SA%204.0-lightgrey.svg)](https://creativecommons.org/licenses/by-nc-sa/4.0/) [![arXiv](https://img.shields.io/badge/arXiv-2604.24311-b31b1b.svg)](https://arxiv.org/abs/2604.24311) [![EC3 2026](https://img.shields.io/badge/EC3-2026-blue.svg)](https://ec-3.org/publication/ec32026_197/) [![Dataset](https://img.shields.io/badge/%F0%9F%A4%97%20Dataset-DeKH-ffd21e.svg)](https://huggingface.co/datasets/RPTU-FGMB/DeKH)
+
+This folder is a self-contained package: model weights, config, a minimal copy
+of the [Pointcept](https://github.com/Pointcept/Pointcept) codebase (which
+this model was trained with), and one script (`segment_scan.py`) that takes a
+raw `.las`/`.laz`/`.ply` scan and produces a segmented file of the same
+format back out.
+
+## What's in here
+
+```
+configs/model_config.py   model architecture + class list
+weights/model_best.pth    trained checkpoint
+pointcept/                trimmed copy of the Pointcept codebase (see below)
+segment_scan.py           chunk -> infer -> merge inference script
+pyproject.toml / uv.lock  environment definition
+```
+
+## Model
+
+- Backbone: **PT-v3m1** (Point Transformer V3), ~46M parameters.
+- Pretrained jointly on Structured3D + ScanNet + S3DIS (Point Prompt
+  Training / PPT-v1m1) with a `RandomColorDrop` augmentation (points
+  randomly lose their color during training, so the model doesn't over-rely
+  on it).
+- Fine-tuned on **CV4AEC** training data (indoor/construction point clouds).
+- 10 output classes: `clutter, floor, ceiling, wall, column, door, window,
+  stairs, railing, lights` (see `configs/model_config.py` -> `data.names`).
+
+The `pointcept/` directory here is deliberately **not** a full checkout of
+the Pointcept codebase or of the internal repo this model was trained in —
+it's trimmed to only what's needed to build this specific model and run
+inference (backbone + generic dataset/transform machinery), so it doesn't
+pull in unrelated dataset loaders or the CV4AEC-specific dataset class
+(`CV4AECDataset`) at all -- `segment_scan.py` builds tiles as plain
+`pointcept.datasets.DefaultDataset` items instead, so you don't need the
+CV4AEC dataset or its raw data to run this. If you want the full upstream
+codebase instead (e.g. to train further), see the notes at the bottom.
+
+## Setup
+
+Requires a CUDA GPU (tested on an RTX 4090) and
+[`uv`](https://docs.astral.sh/uv/):
+
+```bash
+uv sync
+```
+
+This creates `.venv/` with Python 3.12, PyTorch 2.7 (CUDA 12.6), spconv,
+torch-scatter/torch-cluster, and the handful of other packages
+`pointcept/` needs -- pinned in `uv.lock` for a reproducible install.
+`flash-attn` is deliberately **not** installed (it requires a slow
+from-source build); the model falls back to standard attention, which this
+config already reflects (`enable_flash=False`).
+
+## Running it
+
+```bash
+uv run segment_scan.py --input your_scan.las --output your_scan_segmented.las
+```
+
+`--input`/`--output` accept `.las`, `.laz`, or `.ply` (mix and match --
+input and output don't need to be the same format).
+
+What it does:
+1. Reads the point cloud and estimates per-point normals (the model expects
+   color + normals as input features).
+2. Splits it into overlapping ~20m tiles (large scans don't fit the model in
+   one shot) -- `--tile-size` / `--overlap` to change this.
+3. Runs each tile through the model with 10-way test-time augmentation
+   (5 scales x flip/no-flip), softmax-averaged.
+4. Merges tiles back into one point cloud: for points covered by more than
+   one tile (the overlap region), takes a majority vote across tiles.
+5. Writes a file with the original coordinates/colors plus a per-point
+   `classification` field (0-9, matching the class list above -- for `.ply`
+   this is a custom scalar field, readable in CloudCompare/MeshLab) and a
+   companion `<output>_labels.npy`.
+
+On a ~6.6M point scan this takes about 8 minutes end-to-end on an RTX 4090
+(most of it is the 10x test-time-augmentation passes per tile). Use
+`--tile-size` to trade off tile count vs. points-per-tile if you need to fit
+tighter GPU memory, or edit `TEST_CFG["aug_transform"]` at the top of
+`segment_scan.py` to drop some of the TTA passes for a faster, slightly
+less accurate run.
+
+## If you want the full upstream Pointcept instead
+
+This delivery's `pointcept/` is intentionally minimal. If you need the rest
+(other model architectures, training scripts, the full dataset zoo), clone
+<https://github.com/Pointcept/Pointcept> and drop `configs/model_config.py`
++ `weights/model_best.pth` into it -- the checkpoint loads cleanly against
+current upstream (`strict=True`), with one config-level rename to be aware
+of: newer Pointcept renamed the backbone's `cls_mode` argument to
+`enc_mode` (already reflected in `configs/model_config.py` here).
+
+## License
+
+- **Code** (`segment_scan.py`, `pointcept/`, configs): MIT.
+- **Weights** (`weights/model_best.pth`): CC-BY-NC-SA-4.0.
+
+## Citation
+
+If you use this model, please cite the paper it was built for:
+
+```bibtex
+@inproceedings{chamseddine2026bimstruct3d,
+  author    = {Chamseddine, Mahdi and Kaufmann, Fabian and Schellen, Marius and Glock, Christian and Stricker, Didier and Rambach, Jason},
+  booktitle    = {Proceedings of the European Conference on Computing in Construction (EC3)},
+  organization = {European Council on Computing in Construction},
+  title     = {BIMStruct3D: A Fully Automated Hybrid Learning Scan-to-BIM Pipeline with Integrated Topology Refinement},
+  year      = {2026},
+  doi       = {10.35490/EC3.2026}
+}
+```
+
+Please also credit the [CV4AEC Scan-to-BIM Competition 2024](https://cv4aec.github.io/cvpr2024)
+([challenge repo](https://github.com/GradientSpaces/cv4aec-challenge), Gradient
+Spaces @ Stanford) if you use data or results tied to that benchmark -- this
+is the specific challenge whose data the model was fine-tuned on.
+
+## Related dataset
+
+This model was used to produce the semantic segmentation underlying the
+BIMStruct3D pipeline's results on the **DeKH (German Hospital) dataset**:
+real hospital point clouds (four scenes across three buildings) with
+semantic annotations and ground-truth IFC BIM models, released alongside
+the paper above.
+
+- Dataset: <https://huggingface.co/datasets/RPTU-FGMB/DeKH>
+- License: CC BY-NC-SA 4.0 (non-commercial)
+
+## Acknowledgement
+
+This research was funded by the European Union as part of the projects:
+HumanTech (Grant Agreement 101058236) and ShieldBOT (Grant Agreement
+101235093).

configs/model_config.py

@@ -0,0 +1,114 @@
+# Resolved training/model config for the bundled checkpoint (weights/model_best.pth).
+# Trimmed down from the original training config: the original `data.train` /
+# `data.val` / `data.test` blocks referenced our internal dataset class and
+# training-data paths, which aren't needed to run inference and have been
+# removed. `segment_scan.py` builds the model straight from `model` below and
+# applies its own preprocessing pipeline (see PREPROCESS_TRANSFORM /
+# TEST_TIME_AUGMENTATIONS in that script) instead of going through a dataset
+# class, so nothing here is unused cruft other than the two notes below.
+
+weight = "weights/model_best.pth"  # relative to this delivery folder's root
+resume = False
+evaluate = True
+test_only = False
+seed = 37743921
+
+# save_path is only consulted by the standard Pointcept tools/train.py /
+# tools/test.py entry points (dataset-based train/eval), which this delivery
+# doesn't use. Left here only so this file stays loadable by those scripts if
+# you ever want to wire up your own dataset for standard evaluation.
+save_path = "./output"
+
+num_worker = 8
+batch_size = 1
+batch_size_val = None
+batch_size_test = None
+epoch = 3000
+eval_epoch = 100
+sync_bn = False
+enable_amp = True
+empty_cache = False
+empty_cache_per_epoch = False
+find_unused_parameters = False
+mix_prob = 0.8
+param_dicts = [dict(keyword="block", lr=0.0003)]
+hooks = [
+    dict(type="CheckpointLoader"),
+    dict(type="IterationTimer", warmup_iter=2),
+    dict(type="InformationWriter"),
+    dict(type="SemSegEvaluator"),
+    dict(type="CheckpointSaver", save_freq=None),
+    dict(type="PreciseEvaluator", test_last=False),
+]
+train = dict(type="DefaultTrainer")
+test = dict(type="SemSegTester", verbose=True)
+model = dict(
+    type="DefaultSegmentorV2",
+    num_classes=10,
+    backbone_out_channels=64,
+    backbone=dict(
+        type="PT-v3m1",
+        in_channels=6,
+        order=("z", "z-trans", "hilbert", "hilbert-trans"),
+        stride=(2, 2, 2, 2),
+        enc_depths=(2, 2, 2, 6, 2),
+        enc_channels=(32, 64, 128, 256, 512),
+        enc_num_head=(2, 4, 8, 16, 32),
+        enc_patch_size=(1024, 1024, 1024, 1024, 1024),
+        dec_depths=(2, 2, 2, 2),
+        dec_channels=(64, 64, 128, 256),
+        dec_num_head=(4, 4, 8, 16),
+        dec_patch_size=(1024, 1024, 1024, 1024),
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        drop_path=0.3,
+        shuffle_orders=True,
+        pre_norm=True,
+        enable_rpe=False,
+        enable_flash=False,  # flash-attn deliberately excluded from the minimal install; standard attention is used instead
+        upcast_attention=False,
+        upcast_softmax=False,
+        enc_mode=False,  # renamed from `cls_mode` in newer upstream Pointcept
+        pdnorm_bn=False,
+        pdnorm_ln=False,
+        pdnorm_decouple=True,
+        pdnorm_adaptive=False,
+        pdnorm_affine=True,
+        pdnorm_conditions=("ScanNet", "S3DIS", "Structured3D"),
+    ),
+    criteria=[
+        dict(type="CrossEntropyLoss", loss_weight=1.0, ignore_index=-1),
+        dict(type="LovaszLoss", mode="multiclass", loss_weight=1.0, ignore_index=-1),
+    ],
+)
+optimizer = dict(type="AdamW", lr=0.003, weight_decay=0.05)
+scheduler = dict(
+    type="OneCycleLR",
+    max_lr=[0.003, 0.0003],
+    pct_start=0.05,
+    anneal_strategy="cos",
+    div_factor=10.0,
+    final_div_factor=1000.0,
+)
+
+# Class scheme the model was trained on. Predicted label indices from
+# segment_scan.py map directly onto this list (no remap table needed).
+data = dict(
+    num_classes=10,
+    ignore_index=-1,
+    names=[
+        "clutter",
+        "floor",
+        "ceiling",
+        "wall",
+        "column",
+        "door",
+        "window",
+        "stairs",
+        "railing",
+        "lights",
+    ],
+)

pointcept/datasets/__init__.py

@@ -0,0 +1,9 @@
+# Trimmed for inference-only use: segment_scan.py only needs the generic
+# DefaultDataset (directory-of-.npy-assets loader) plus the transform/collate
+# machinery. The named dataset classes below each need their own source
+# data on disk to even be useful, and some pull in extra devkits (nuscenes,
+# waymo, ...) we don't need here.
+from .defaults import DefaultDataset, ConcatDataset
+from .builder import build_dataset
+from .utils import point_collate_fn, collate_fn
+from .transform import *

pointcept/datasets/aeo.py

@@ -0,0 +1,29 @@
+"""
+AEO 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 .defaults import DefaultDataset
+from .builder import DATASETS
+
+
+@DATASETS.register_module()
+class AEODataset(DefaultDataset):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        label_mapping = np.ones(41, dtype=int) * -1
+        label_mapping[[0, 1, 3, 4, 13, 16, 19, 21, 22, 28, 29, 34, 36, 37, 38, 39]] = (
+            np.arange(16)
+        )
+        self.label_mapping = label_mapping
+
+    def get_data(self, idx):
+        data_dict = super().get_data(idx)
+        data_dict["segment"] = self.label_mapping[data_dict["segment"]]
+        return data_dict

pointcept/datasets/builder.py

@@ -0,0 +1,15 @@
+"""
+Dataset Builder
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+from pointcept.utils.registry import Registry
+
+DATASETS = Registry("datasets")
+
+
+def build_dataset(cfg):
+    """Build datasets."""
+    return DATASETS.build(cfg)

pointcept/datasets/cap3d.py

@@ -0,0 +1,245 @@
+"""
+Objaverse Point Cloud Dataset Loader
+
+Load pre-processed .pt files.
+Each .pt file is assumed to be a dictionary of PyTorch tensors.
+"""
+
+import os
+import torch
+import numpy as np
+import open3d as o3d
+from PIL import Image
+import json
+from collections.abc import Sequence
+import random
+
+from .builder import DATASETS
+from .defaults import DefaultDataset, DefaultImagePointDataset
+
+
+@DATASETS.register_module()
+class Cap3DDataset(DefaultDataset):
+    """
+    Cap3D Dataset for loading .pt files.
+    """
+
+    def __init__(
+        self,
+        data_num=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if data_num:
+            self.data_list = self.data_list[:data_num]
+
+    @staticmethod
+    def get_normals(center, coords):
+        Cs = np.repeat(center.reshape((1, -1)), coords.shape[0], axis=0)
+        view_dirs = coords - Cs
+        view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True)
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(coords)
+        pcd.estimate_normals()
+        normals = np.asarray(pcd.normals)
+        dot_product = np.sum(normals * view_dirs, axis=-1)
+        flip_mask = dot_product < 0
+        normals[flip_mask] = -normals[flip_mask]
+
+        # Normalize normals a nd m
+        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
+        return normals
+
+    def get_data(self, idx):
+        """
+        The core data loading function.
+
+        Args:
+            idx (int): The index of the data sample.
+
+        Returns:
+            dict: A dictionary containing the loaded and processed point cloud data.
+        """
+        data_path = self.data_list[idx % len(self.data_list)]
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+
+        try:
+            data = torch.load(data_path, map_location="cpu")
+        except Exception as e:
+            raise IOError(f"Failed to load a .pt file at {data_path}. Error: {e}")
+
+        data = data.numpy().transpose()
+        color = data[:, 3:]
+        coord = data[:, :3]
+        normal = self.get_normals(np.mean(coord, axis=0), coord)
+
+        data_dict = {}
+
+        data_dict["coord"] = coord
+        data_dict["color"] = color
+        data_dict["normal"] = normal
+
+        data_dict["name"] = name
+        data_dict["split"] = split
+
+        return data_dict
+
+
+@DATASETS.register_module()
+class Cap3DImagePointDataset(DefaultImagePointDataset):
+    def __init__(
+        self,
+        data_num=-1,
+        **kwargs,
+    ):
+        self.data_num = data_num
+        super().__init__(**kwargs)
+
+    def get_data_list(self):
+        split_list = {}
+        if isinstance(self.split, str):
+            data_path = os.path.join(self.data_root, "splits", f"{self.split}.json")
+            with open(data_path, "r", encoding="utf-8") as file:
+                data_list = json.load(file)
+            split_list[self.split] = list(data_list.keys())
+        elif isinstance(self.split, Sequence):
+            data_list = {}
+            for split in self.split:
+                data_path = os.path.join(self.data_root, "splits", f"{split}.json")
+                with open(data_path, "r", encoding="utf-8") as file:
+                    data_split_dict = json.load(file)
+                    data_list.update(data_split_dict)
+                split_list[split] = list(data_split_dict.keys())
+        else:
+            raise NotImplementedError
+
+        if self.data_num < 0 or self.data_num >= len(data_list):
+            selected_data_list = data_list
+            selected_split_list = split_list
+            return selected_data_list, selected_split_list
+
+        all_keys = list(data_list.keys())
+        selected_keys = random.sample(all_keys, self.data_num)
+        selected_data_list = {key: data_list[key] for key in selected_keys}
+        valid_keys = set(selected_keys)
+
+        selected_split_list = {}
+        for split_name, key_list in split_list.items():
+            filtered_keys = [key for key in key_list if key in valid_keys]
+            if filtered_keys:
+                selected_split_list[split_name] = filtered_keys
+        return selected_data_list, selected_split_list
+
+    @staticmethod
+    def get_normals(center, coords):
+        Cs = np.repeat(center.reshape((1, -1)), coords.shape[0], axis=0)
+        view_dirs = coords - Cs
+        view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True)
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(coords)
+        pcd.estimate_normals()
+        normals = np.asarray(pcd.normals)
+        dot_product = np.sum(normals * view_dirs, axis=-1)
+        flip_mask = dot_product < 0
+        normals[flip_mask] = -normals[flip_mask]
+        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
+        return normals
+
+    def get_pc_data(self, data_path):
+        try:
+            data = torch.load(data_path, map_location="cpu")
+        except Exception as e:
+            raise IOError(f"Failed to load a .pt file at {data_path}. Error: {e}")
+
+        data = data.numpy().transpose()
+        color = data[:, 3:]
+        coord = data[:, :3]
+        normal = self.get_normals(np.mean(coord, axis=0), coord)
+
+        data_dict = {}
+
+        data_dict["coord"] = coord
+        data_dict["color"] = color
+        data_dict["normal"] = normal
+
+        return data_dict
+
+    def get_data(self, idx):
+        data_dict = {}
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        data_dict["name"] = name
+        data_dict["split"] = split
+        data_path = self.data_list[name]
+
+        pointclouds_path = data_path["pointclouds"]
+        pc_dict = self.get_pc_data(pointclouds_path)
+        data_dict.update(pc_dict)
+        imgs_path = data_path["images"]
+        imgs = [Image.open(asset).convert("RGB") for asset in imgs_path]
+        if len(imgs) > 0:
+            img_width, img_height = imgs[0].size
+            div_w = img_width // self.patch_w
+            div_h = img_height // self.patch_h
+            div_min = max(min(div_w, div_h), 1)
+            crop_img_width = div_min * self.patch_w
+            crop_img_height = div_min * self.patch_h
+            left = int((img_width - crop_img_width) / 2)
+            top = int((img_height - crop_img_height) / 2)
+            right = int((img_width + crop_img_width) / 2)
+            bottom = int((img_height + crop_img_height) / 2)
+            imgs = [img.crop((left, top, right, bottom)) for img in imgs]
+            imgs = [self.transform_img(img) for img in imgs]
+            imgs_list = torch.stack(imgs)
+            data_dict["images"] = imgs_list.float()
+        else:
+            data_dict["images"] = torch.empty(
+                (0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size)
+            )
+        data_dict["img_num"] = np.array([data_dict["images"].shape[0]], dtype=np.int32)
+
+        correspondences_path = data_path["correspondences"]
+        correspondence_infos = np.ones(
+            (data_dict["coord"].shape[0], len(correspondences_path), 2),
+            dtype=np.float32,
+        ) * (-1)
+        for asset_id, asset in enumerate(correspondences_path):
+            correspondence_info = np.load(asset).astype(np.float32)
+            if np.array_equal(correspondence_info, -np.ones((1, 3))):
+                continue
+            correspondence_info = self.resize_correspondence_info(
+                correspondence_info,
+                (self.patch_h * self.patch_size, self.patch_w * self.patch_size),
+                (img_height, img_width),
+                (left, top, right, bottom),
+                self.patch_size,
+            )
+            correspondence_infos[
+                correspondence_info[:, -1].astype(np.int32), asset_id, :
+            ] = correspondence_info[:, :-1]
+        data_dict["correspondence"] = correspondence_infos
+
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+        return data_dict

pointcept/datasets/dataloader.py

@@ -0,0 +1,205 @@
+from functools import partial
+import weakref
+import torch
+import math
+import torch.utils.data
+import torch.distributed as dist
+from torch.utils.data import Sampler
+
+import pointcept.utils.comm as comm
+from pointcept.datasets.utils import point_collate_fn
+from pointcept.datasets import ConcatDataset
+from pointcept.utils.env import set_seed
+
+
+class MultiDatasetDummySampler:
+    def __init__(self):
+        self.dataloader = None
+
+    def set_epoch(self, epoch):
+        if comm.get_world_size() > 1:
+            for dataloader in self.dataloader.dataloaders:
+                dataloader.sampler.set_epoch(epoch)
+        return
+
+
+class MultiDatasetDataloader:
+    """
+    Multiple Datasets Dataloader, batch data from a same dataset and mix up ratio determined by loop of each sub dataset.
+    The overall length is determined by the main dataset (first) and loop of concat dataset.
+    """
+
+    def __init__(
+        self,
+        concat_dataset: ConcatDataset,
+        batch_size_per_gpu: int,
+        num_worker_per_gpu: int,
+        mix_prob=0,
+        seed=None,
+    ):
+        self.datasets = concat_dataset.datasets
+        self.ratios = [dataset.loop for dataset in self.datasets]
+        # reset data loop, original loop serve as ratios
+        for dataset in self.datasets:
+            dataset.loop = 1
+        # determine union training epoch by main dataset
+        self.datasets[0].loop = concat_dataset.loop
+        # build sub-dataloaders
+        num_workers = num_worker_per_gpu // len(self.datasets)
+        self.dataloaders = []
+        for dataset_id, dataset in enumerate(self.datasets):
+            if comm.get_world_size() > 1:
+                sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            else:
+                sampler = None
+
+            init_fn = (
+                partial(
+                    self._worker_init_fn,
+                    dataset_id=dataset_id,
+                    num_workers=num_workers,
+                    num_datasets=len(self.datasets),
+                    rank=comm.get_rank(),
+                    seed=seed,
+                )
+                if seed is not None
+                else None
+            )
+            self.dataloaders.append(
+                torch.utils.data.DataLoader(
+                    dataset,
+                    batch_size=batch_size_per_gpu,
+                    shuffle=(sampler is None),
+                    num_workers=num_worker_per_gpu,
+                    sampler=sampler,
+                    collate_fn=partial(point_collate_fn, mix_prob=mix_prob),
+                    pin_memory=True,
+                    worker_init_fn=init_fn,
+                    drop_last=True,
+                    persistent_workers=True,
+                )
+            )
+        self.sampler = MultiDatasetDummySampler()
+        self.sampler.dataloader = weakref.proxy(self)
+
+    def __iter__(self):
+        iterator = [iter(dataloader) for dataloader in self.dataloaders]
+        while True:
+            for i in range(len(self.ratios)):
+                for _ in range(self.ratios[i]):
+                    try:
+                        batch = next(iterator[i])
+                    except StopIteration:
+                        if i == 0:
+                            return
+                        else:
+                            iterator[i] = iter(self.dataloaders[i])
+                            batch = next(iterator[i])
+                    yield batch
+
+    def __len__(self):
+        main_data_loader_length = len(self.dataloaders[0])
+        return (
+            main_data_loader_length // self.ratios[0] * sum(self.ratios)
+            + main_data_loader_length % self.ratios[0]
+        )
+
+    @staticmethod
+    def _worker_init_fn(worker_id, num_workers, dataset_id, num_datasets, rank, seed):
+        worker_seed = (
+            num_workers * num_datasets * rank
+            + num_workers * dataset_id
+            + worker_id
+            + seed
+        )
+        set_seed(worker_seed)
+
+
+class DistributedImbalancedSampler(Sampler):
+    def __init__(
+        self,
+        dataset,
+        sampled_dataset_index,
+        sampled_dataset_limit,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+        seed=0,
+    ):
+        if num_replicas is None:
+            if not dist.is_available():
+                num_replicas = 1
+                rank = 0
+            else:
+                try:
+                    num_replicas = dist.get_world_size()
+                    rank = dist.get_rank()
+                except:
+                    num_replicas = 1
+                    rank = 0
+
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.seed = seed
+        self.sampled_dataset_index = sampled_dataset_index
+        self.sampled_dataset_limit = sampled_dataset_limit
+        self.shuffle = shuffle
+
+        self.lengths = [len(d) for d in dataset.datasets]
+        self.offsets = [0] + list(
+            torch.cumsum(torch.tensor(self.lengths), 0)[:-1].numpy()
+        )
+        self.total_size_per_epoch = (
+            sum(self.lengths)
+            - self.lengths[self.sampled_dataset_index]
+            + self.sampled_dataset_limit
+        )
+        self.num_samples = math.ceil(self.total_size_per_epoch / self.num_replicas)
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch)
+
+        indices = []
+        for i, length in enumerate(self.lengths):
+            if i == self.sampled_dataset_index:
+                sub_indices = (
+                    torch.randperm(length, generator=g)[: self.sampled_dataset_limit]
+                    + self.offsets[i]
+                )
+                indices.append(sub_indices)
+            else:
+                sub_indices = torch.arange(length) + self.offsets[i]
+                indices.append(sub_indices)
+
+        final_indices = torch.cat(indices)
+
+        if self.shuffle:
+            final_indices = final_indices[
+                torch.randperm(len(final_indices), generator=g)
+            ]
+
+        final_indices = final_indices.tolist()
+        padding_size = self.total_size - len(final_indices)
+        if padding_size <= len(final_indices):
+            final_indices += final_indices[:padding_size]
+        else:
+            final_indices += (
+                final_indices * math.ceil(padding_size / len(final_indices))
+            )[:padding_size]
+
+        assert len(final_indices) == self.total_size
+
+        offset = self.num_samples * self.rank
+        subsample_indices = final_indices[offset : offset + self.num_samples]
+
+        return iter(subsample_indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

pointcept/datasets/defaults.py

@@ -0,0 +1,676 @@
+"""
+Default Datasets
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import glob
+import json
+import numpy as np
+import torch
+from copy import deepcopy
+from torch.utils.data import Dataset
+from collections.abc import Sequence
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+from torchvision.transforms import transforms as T
+
+from pointcept.utils.logger import get_root_logger
+from pointcept.utils.cache import shared_dict
+
+from .builder import DATASETS, build_dataset
+from .transform import Compose, TRANSFORMS
+
+INTERPOLATION_MODE = {
+    "bilinear": InterpolationMode.BILINEAR,
+    "bicubic": InterpolationMode.BICUBIC,
+    "nearest": InterpolationMode.NEAREST,
+}
+
+
+@DATASETS.register_module()
+class DefaultDataset(Dataset):
+    VALID_ASSETS = [
+        "coord",
+        "color",
+        "normal",
+        "strength",
+        "segment",
+        "instance",
+        "pose",
+    ]
+
+    def __init__(
+        self,
+        split="train",
+        data_root="data/dataset",
+        transform=None,
+        test_mode=False,
+        test_cfg=None,
+        cache=False,
+        ignore_index=-1,
+        loop=1,
+    ):
+        super(DefaultDataset, self).__init__()
+        self.data_root = data_root
+        self.split = split
+        self.transform = Compose(transform)
+        self.cache = cache
+        self.ignore_index = ignore_index
+        self.loop = (
+            loop if not test_mode else 1
+        )  # force make loop = 1 while in test mode
+        self.test_mode = test_mode
+        self.test_cfg = test_cfg if test_mode else None
+        if test_mode:
+            self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize)
+            self.test_crop = (
+                TRANSFORMS.build(self.test_cfg.crop) if self.test_cfg.crop else None
+            )
+            self.post_transform = Compose(self.test_cfg.post_transform)
+            self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform]
+
+        self.data_list = self.get_data_list()
+        logger = get_root_logger()
+        logger.info(
+            "Totally {} x {} samples in {} {} set.".format(
+                len(self.data_list), self.loop, os.path.basename(self.data_root), split
+            )
+        )
+
+    def get_data_list(self):
+        if isinstance(self.split, str):
+            split_list = [self.split]
+        elif isinstance(self.split, Sequence):
+            split_list = self.split
+        else:
+            raise NotImplementedError
+
+        data_list = []
+        for split in split_list:
+            if os.path.isfile(os.path.join(self.data_root, split)):
+                with open(os.path.join(self.data_root, split)) as f:
+                    data_list += [
+                        os.path.join(self.data_root, data) for data in json.load(f)
+                    ]
+            else:
+                data_list += glob.glob(os.path.join(self.data_root, split, "*"))
+        return data_list
+
+    def get_data(self, idx):
+        data_path = self.data_list[idx % len(self.data_list)]
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        if self.cache:
+            cache_name = f"pointcept-{name}"
+            return shared_dict(cache_name)
+
+        data_dict = {}
+        assets = os.listdir(data_path)
+        for asset in assets:
+            if not asset.endswith(".npy"):
+                continue
+            if asset[:-4] not in self.VALID_ASSETS:
+                continue
+            data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
+        data_dict["name"] = name
+        data_dict["split"] = split
+
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+        return data_dict
+
+    def get_data_name(self, idx):
+        return os.path.basename(self.data_list[idx % len(self.data_list)])
+
+    def get_split_name(self, idx):
+        return os.path.basename(
+            os.path.dirname(self.data_list[idx % len(self.data_list)])
+        )
+
+    def prepare_train_data(self, idx):
+        # load data
+        data_dict = self.get_data(idx)
+        data_dict = self.transform(data_dict)
+        return data_dict
+
+    def prepare_test_data(self, idx):
+        # load data
+        data_dict = self.get_data(idx)
+        data_dict = self.transform(data_dict)
+        result_dict = dict(segment=data_dict.pop("segment"), name=data_dict.pop("name"))
+        if "origin_segment" in data_dict:
+            assert "inverse" in data_dict
+            result_dict["origin_segment"] = data_dict.pop("origin_segment")
+            result_dict["inverse"] = data_dict.pop("inverse")
+
+        data_dict_list = []
+        for aug in self.aug_transform:
+            data_dict_list.append(aug(deepcopy(data_dict)))
+
+        fragment_list = []
+        for data in data_dict_list:
+            if self.test_voxelize is not None:
+                data_part_list = self.test_voxelize(data)
+            else:
+                data["index"] = np.arange(data["coord"].shape[0])
+                data_part_list = [data]
+            for data_part in data_part_list:
+                if self.test_crop is not None:
+                    data_part = self.test_crop(data_part)
+                else:
+                    data_part = [data_part]
+                fragment_list += data_part
+
+        for i in range(len(fragment_list)):
+            fragment_list[i] = self.post_transform(fragment_list[i])
+        result_dict["fragment_list"] = fragment_list
+        return result_dict
+
+    def __getitem__(self, idx):
+        if self.test_mode:
+            return self.prepare_test_data(idx)
+        else:
+            return self.prepare_train_data(idx)
+
+    def __len__(self):
+        return len(self.data_list) * self.loop
+
+
+@DATASETS.register_module()
+class DefaultImagePointDataset(Dataset):
+    PC_VALID_ASSETS = [
+        "coord",
+        "color",
+        "normal",
+        "superpoint",
+        "pose",
+        "segment",
+        "instance",
+        "strength",
+    ]
+
+    def __init__(
+        self,
+        split="train",
+        data_root="data/dataset",
+        transform=None,
+        test_mode=False,
+        test_cfg=None,
+        cache=False,
+        ignore_index=-1,
+        loop=1,
+        crop_h=630,
+        crop_w=1120,
+        patch_size=14,
+        interpolation="bilinear",
+        if_img=True,
+    ):
+        super(DefaultImagePointDataset, self).__init__()
+        self.if_img = if_img
+        self.data_root = data_root
+        self.split = split
+        self.transform = Compose(transform)
+        self.cache = cache
+        self.ignore_index = ignore_index
+        self.loop = (
+            loop if not test_mode else 1
+        )  # force make loop = 1 while in test mode
+        self.test_mode = test_mode
+        self.test_cfg = test_cfg if test_mode else None
+
+        if test_mode:
+            self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize)
+            self.test_crop = (
+                TRANSFORMS.build(self.test_cfg.crop) if self.test_cfg.crop else None
+            )
+            self.post_transform = Compose(self.test_cfg.post_transform)
+            self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform]
+
+        self.data_list, self.split_list = self.get_data_list()
+        if isinstance(self.data_list, dict):
+            self.data_name = list(self.data_list.keys())
+        elif isinstance(self.data_list[0], dict):
+            self.data_name = list([data["token"] for data in self.data_list])
+        else:
+            self.data_name = self.data_list
+        logger = get_root_logger()
+        logger.info(
+            "Totally {} x {} samples in {} {} set.".format(
+                len(self.data_name), self.loop, os.path.basename(self.data_root), split
+            )
+        )
+
+        self.crop_h = crop_h
+        self.crop_w = crop_w
+        self.patch_size = patch_size
+        self.patch_h = crop_h // patch_size
+        self.patch_w = crop_w // patch_size
+        self.transform_img = T.Compose(
+            [
+                T.Resize(
+                    (self.patch_h * self.patch_size, self.patch_w * self.patch_size),
+                    interpolation=INTERPOLATION_MODE[interpolation],
+                ),
+                T.ToTensor(),
+            ]
+        )
+
+    def resize_correspondence_info(
+        self, correspondence, size, size0, crop_size, _alignment
+    ):
+        h, w = size
+        h0, w0 = size0
+        left, top, right, bottom = crop_size
+        crop_h = bottom - top
+        crop_w = right - left
+        mask_crop = (
+            (correspondence[:, 1] >= top)
+            & (correspondence[:, 1] < bottom)
+            & (correspondence[:, 0] >= left)
+            & (correspondence[:, 0] < right)
+        )
+        correspondence = correspondence[mask_crop]
+        correspondence[:, 1] -= top
+        correspondence[:, 0] -= left
+        correspondence[:, 1] = (correspondence[:, 1] * h / crop_h / _alignment).astype(
+            np.float32
+        )
+        correspondence[:, 0] = (correspondence[:, 0] * w / crop_w / _alignment).astype(
+            np.float32
+        )
+        correspondence = correspondence[:, [1, 0, 2]]
+        correspondence = np.unique(correspondence, axis=0)
+        return correspondence
+
+    def get_data_list(self):
+        split_list = {}
+        if isinstance(self.split, str):
+            data_path = os.path.join(self.data_root, "splits", f"{self.split}.json")
+            with open(data_path, "r", encoding="utf-8") as file:
+                data_list = json.load(file)
+            split_list[self.split] = list(data_list.keys())
+        elif isinstance(self.split, Sequence):
+            data_list = {}
+            for split in self.split:
+                data_path = os.path.join(self.data_root, "splits", f"{split}.json")
+                with open(data_path, "r", encoding="utf-8") as file:
+                    data_split_dict = json.load(file)
+                    data_list.update(data_split_dict)
+                split_list[split] = list(data_split_dict.keys())
+        else:
+            raise NotImplementedError
+        return data_list, split_list
+
+    def get_data_name(self, idx):
+        return self.data_name[idx % len(self.data_name)]
+
+    def get_split_name(self, idx):
+        for split, names in self.split_list.items():
+            if self.data_name[idx % len(self.data_name)] in names:
+                return split
+        return None
+
+    def get_data(self, idx):
+        data_dict = {}
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        data_dict["name"] = name
+        data_dict["split"] = split
+        data_path = self.data_list[name]
+
+        pointclouds_path = data_path["pointclouds"]
+        assets = os.listdir(pointclouds_path)
+        for asset in assets:
+            if not asset.endswith(".npy"):
+                continue
+            if asset[:-4] not in self.PC_VALID_ASSETS:
+                continue
+            data_dict[asset[:-4]] = np.load(os.path.join(pointclouds_path, asset))
+        if self.if_img:
+            imgs_path = data_path["images"]
+            imgs = [Image.open(asset) for asset in imgs_path]
+            if len(imgs) > 0:
+                img_width, img_height = imgs[0].size
+                div_w = img_width // self.patch_w
+                div_h = img_height // self.patch_h
+                div_min = max(min(div_w, div_h), 1)
+                crop_img_width = div_min * self.patch_w
+                crop_img_height = div_min * self.patch_h
+                left = int((img_width - crop_img_width) / 2)
+                top = int((img_height - crop_img_height) / 2)
+                right = int((img_width + crop_img_width) / 2)
+                bottom = int((img_height + crop_img_height) / 2)
+                imgs = [img.crop((left, top, right, bottom)) for img in imgs]
+                imgs = [self.transform_img(img) for img in imgs]
+                imgs_list = torch.stack(imgs)
+                data_dict["images"] = imgs_list.float()
+            else:
+                data_dict["images"] = torch.empty(
+                    (
+                        0,
+                        3,
+                        self.patch_h * self.patch_size,
+                        self.patch_w * self.patch_size,
+                    )
+                )
+            data_dict["img_num"] = np.array(
+                [data_dict["images"].shape[0]], dtype=np.int32
+            )
+
+            correspondences_path = data_path["correspondences"]
+            correspondence_infos = np.ones(
+                (data_dict["coord"].shape[0], len(correspondences_path), 2),
+                dtype=np.float32,
+            ) * (-1)
+            for asset_id, asset in enumerate(correspondences_path):
+                correspondence_info = np.load(asset).astype(np.float32)
+                if np.array_equal(correspondence_info, -np.ones((1, 3))):
+                    continue
+                correspondence_info = self.resize_correspondence_info(
+                    correspondence_info,
+                    (self.patch_h * self.patch_size, self.patch_w * self.patch_size),
+                    (img_height, img_width),
+                    (left, top, right, bottom),
+                    self.patch_size,
+                )
+                correspondence_infos[
+                    correspondence_info[:, -1].astype(np.int32), asset_id, :
+                ] = correspondence_info[:, :-1]
+            data_dict["correspondence"] = correspondence_infos  # .reshape(-1, 2)
+
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+        return data_dict
+
+    def prepare_train_data(self, idx):
+        # load data
+        data_dict = self.get_data(idx)
+        data_dict = self.transform(data_dict)
+        return data_dict
+
+    def prepare_test_data(self, idx):
+        # load data
+        data_dict = self.get_data(idx)
+        data_dict = self.transform(data_dict)
+        result_dict = dict(segment=data_dict.pop("segment"), name=data_dict.pop("name"))
+        if "origin_segment" in data_dict:
+            assert "inverse" in data_dict
+            result_dict["origin_segment"] = data_dict.pop("origin_segment")
+            result_dict["inverse"] = data_dict.pop("inverse")
+
+        data_dict_list = []
+        for aug in self.aug_transform:
+            data_dict_list.append(aug(deepcopy(data_dict)))
+
+        fragment_list = []
+        for data in data_dict_list:
+            if self.test_voxelize is not None:
+                data_part_list = self.test_voxelize(data)
+            else:
+                data["index"] = np.arange(data["coord"].shape[0])
+                data_part_list = [data]
+            for data_part in data_part_list:
+                if self.test_crop is not None:
+                    data_part = self.test_crop(data_part)
+                else:
+                    data_part = [data_part]
+                fragment_list += data_part
+
+        for i in range(len(fragment_list)):
+            fragment_list[i] = self.post_transform(fragment_list[i])
+        result_dict["fragment_list"] = fragment_list
+        return result_dict
+
+    def __getitem__(self, idx):
+        if self.test_mode:
+            return self.prepare_test_data(idx)
+        else:
+            return self.prepare_train_data(idx)
+
+    def __len__(self):
+        return len(self.data_list) * self.loop
+
+
+@DATASETS.register_module()
+class DefaultMultiViewImagePointDataset(DefaultImagePointDataset):
+    def __init__(self, depth_scale=1000.0, if_img=True, **kwargs):
+        self.depth_scale = depth_scale
+        self.if_img = if_img
+        super().__init__(if_img=if_img, **kwargs)
+
+    def get_data_list(self):
+        split_list = {}
+        if isinstance(self.split, str):
+            splits_to_load = [self.split]
+        elif isinstance(self.split, Sequence):
+            splits_to_load = self.split
+        else:
+            raise NotImplementedError
+
+        data_list = {}
+        for split in splits_to_load:
+            data_path = os.path.join(self.data_root, "splits", f"{split}.json")
+            if not os.path.exists(data_path):
+                raise FileNotFoundError(f"Split file not found: {data_path}")
+
+            with open(data_path, "r", encoding="utf-8") as file:
+                data_split_dict = json.load(file)
+                data_list.update(data_split_dict)
+            split_list[split] = list(data_split_dict.keys())
+        return data_list, split_list
+
+    @staticmethod
+    def get_normals(center, coords):
+        normals = coords - center[None, :]
+        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
+        return normals
+
+    def get_data(self, idx):
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        data_info = self.data_list[name]
+
+        all_points_world, all_colors, camera_params = [], [], []
+        correspondences = []
+        coord_num = 0
+        for i in range(len(data_info["images"])):
+            rgb_path = data_info["images"][i]
+            depth_path = data_info["depths"][i]
+            pose_path = data_info["Ts"][i]
+            intr_path = data_info["Ks"][i]
+
+            try:
+                with Image.open(rgb_path) as img:
+                    color_img = np.array(img)
+                with Image.open(depth_path) as img:
+                    depth_img = np.array(img)
+                cam_pose_world = np.load(pose_path)
+                cam_intr = np.load(intr_path)
+                camera_params.append({"pose": cam_pose_world, "intr": cam_intr})
+            except FileNotFoundError:
+                continue
+
+            depth_m = depth_img.astype(np.float32) / self.depth_scale
+            h, w = depth_m.shape
+            fx, fy, cx, cy = (
+                cam_intr[0, 0],
+                cam_intr[1, 1],
+                cam_intr[0, 2],
+                cam_intr[1, 2],
+            )
+            v, u = np.indices((h, w))
+            valid_mask = (depth_m > 0) & (depth_m < 0.8)
+            z_cam = depth_m[valid_mask]
+            x_cam, y_cam = (u[valid_mask] - cx) * z_cam / fx, (
+                v[valid_mask] - cy
+            ) * z_cam / fy
+            points_cam = np.vstack((x_cam, y_cam, z_cam)).T
+            points_world = (
+                cam_pose_world
+                @ np.hstack([points_cam, np.ones((points_cam.shape[0], 1))]).T
+            ).T[:, :3]
+            all_points_world.append(points_world)
+            all_colors.append(color_img[valid_mask])
+
+            pixel = np.vstack([u[valid_mask], v[valid_mask]]).T
+            pixel = np.concatenate([pixel, np.arange(pixel.shape[0])[:, None]], axis=1)
+            pixel[:, 2] += coord_num
+            coord_num += points_world.shape[0]
+            correspondences.append(pixel)
+        coord = (
+            np.concatenate(all_points_world, axis=0)
+            if all_points_world
+            else np.zeros((0, 3))
+        )
+        coord[:, 2] = -coord[:, 2]
+        color = np.concatenate(all_colors, axis=0) if all_colors else np.zeros((0, 3))
+        normal = self.get_normals(np.array([0, 0, 0]), coord)
+
+        if self.if_img:
+            imgs_path = data_info["images"]
+            imgs = [Image.open(asset) for asset in imgs_path]
+            img_width, img_height = imgs[0].size
+            div_w = img_width // self.patch_w
+            div_h = img_height // self.patch_h
+            div_min = max(min(div_w, div_h), 1)
+            crop_img_width = div_min * self.patch_w
+            crop_img_height = div_min * self.patch_h
+            left = int((img_width - crop_img_width) / 2)
+            top = int((img_height - crop_img_height) / 2)
+            right = int((img_width + crop_img_width) / 2)
+            bottom = int((img_height + crop_img_height) / 2)
+            imgs = [img.crop((left, top, right, bottom)) for img in imgs]
+            imgs = [self.transform_img(img) for img in imgs]
+            if len(imgs) > 0:
+                imgs_list = torch.stack(imgs)
+                imgs_tensor = imgs_list.float()
+            else:
+                imgs_tensor = torch.empty(
+                    (
+                        0,
+                        3,
+                        self.patch_h * self.patch_size,
+                        self.patch_w * self.patch_size,
+                    )
+                )
+            img_num = np.array([imgs_tensor.shape[0]], dtype=np.int32)
+
+            num_points, num_imgs = coord.shape[0], len(imgs)
+            correspondence_infos = (
+                np.ones((num_points, num_imgs, 2), dtype=np.float32) * -1
+            )
+
+            for asset_id, correspondence_info in enumerate(correspondences):
+                correspondence_info = self.resize_correspondence_info(
+                    correspondence_info,
+                    (self.patch_h * self.patch_size, self.patch_w * self.patch_size),
+                    (img_height, img_width),
+                    (left, top, right, bottom),
+                    self.patch_size,
+                )
+                correspondence_infos[
+                    correspondence_info[:, -1].astype(np.int32), asset_id, :
+                ] = correspondence_info[:, :-1]
+
+            data_dict = {
+                "name": name,
+                "split": split,
+                "coord": coord.astype(np.float32),
+                "color": color.astype(np.float32),
+                "normal": normal.astype(np.float32),
+                "images": imgs_tensor,
+                "img_num": img_num,
+                "correspondence": correspondence_infos,
+            }
+        else:
+            data_dict = {
+                "name": name,
+                "split": split,
+                "coord": coord.astype(np.float32),
+                "color": color.astype(np.float32),
+                "normal": normal.astype(np.float32),
+            }
+        return data_dict
+
+
+@DATASETS.register_module()
+class ConcatDataset(Dataset):
+    def __init__(self, datasets, loop=1):
+        super(ConcatDataset, self).__init__()
+        self.datasets = [build_dataset(dataset) for dataset in datasets]
+        self.loop = loop
+        self.data_list = self.get_data_list()
+        logger = get_root_logger()
+        logger.info(
+            "Totally {} x {} samples in the concat set.".format(
+                len(self.data_list), self.loop
+            )
+        )
+
+    def get_data_list(self):
+        data_list = []
+        for i in range(len(self.datasets)):
+            data_list.extend(
+                zip(
+                    np.ones(len(self.datasets[i]), dtype=int) * i,
+                    np.arange(len(self.datasets[i])),
+                )
+            )
+        return data_list
+
+    def get_data(self, idx):
+        dataset_idx, data_idx = self.data_list[idx % len(self.data_list)]
+        return self.datasets[dataset_idx][data_idx]
+
+    def get_data_name(self, idx):
+        dataset_idx, data_idx = self.data_list[idx % len(self.data_list)]
+        return self.datasets[dataset_idx].get_data_name(data_idx)
+
+    def __getitem__(self, idx):
+        return self.get_data(idx)
+
+    def __len__(self):
+        return len(self.data_list) * self.loop

pointcept/datasets/hk.py

@@ -0,0 +1,94 @@
+"""
+HK Remote Data Dataset
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import numpy as np
+import torch
+
+from pointcept.utils.cache import shared_dict
+from .builder import DATASETS
+from .defaults import DefaultDataset
+import numpy as np
+from scipy.spatial import cKDTree
+
+
+def find_closest_pair_kdtree(points):
+    n = points.shape[0]
+    if n < 2:
+        return float("inf"), None, None
+
+    tree = cKDTree(points)
+    distances, indices = tree.query(points, k=2)
+    min_dist = np.min(distances[:, 1])
+    min_idx_in_distances = np.argmin(distances[:, 1])
+    point1_idx = min_idx_in_distances
+    point2_idx = indices[min_idx_in_distances, 1]
+
+    return min_dist, (points[point1_idx], points[point2_idx])
+
+
+@DATASETS.register_module()
+class HKDataset(DefaultDataset):
+    def __init__(self, crop_h=630, crop_w=1120, patch_size=14, if_img=True, **kwargs):
+        self.crop_h = crop_h
+        self.crop_w = crop_w
+        self.patch_size = patch_size
+        self.patch_h = crop_h // patch_size
+        self.patch_w = crop_w // patch_size
+        self.if_img = if_img
+        super().__init__(**kwargs)
+
+    def get_data(self, idx):
+        data_path = self.data_list[idx % len(self.data_list)]
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        if self.cache:
+            cache_name = f"pointcept-{name}"
+            return shared_dict(cache_name)
+
+        data_dict = {}
+        assets = os.listdir(data_path)
+        for asset in assets:
+            if not asset.endswith(".npy"):
+                continue
+            if asset[:-4] not in self.VALID_ASSETS:
+                continue
+            data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
+        data_dict["name"] = name
+        data_dict["split"] = split
+        if self.if_img:
+            data_dict["images"] = torch.empty(
+                (0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size)
+            )
+            data_dict["img_num"] = np.array([0], dtype=np.int32)
+            data_dict["correspondence"] = np.ones(
+                (data_dict["coord"].shape[0], 0, 2), dtype=np.float32
+            ) * (-1)
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        return data_dict

pointcept/datasets/hm3d.py

@@ -0,0 +1,45 @@
+"""
+Habitat-Matterport 3D Dataset
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import glob
+import os
+from collections.abc import Sequence
+from .builder import DATASETS
+from .defaults import DefaultDataset
+
+
+@DATASETS.register_module()
+class HM3DDataset(DefaultDataset):
+    def __init__(
+        self,
+        force_label=True,
+        **kwargs,
+    ):
+        # if force_label, only load data with label
+        self.force_label = force_label
+        super().__init__(**kwargs)
+
+    def get_single_data_list(self, split):
+        if self.force_label:
+            data_list = glob.glob(
+                os.path.join(self.data_root, split, "*", "segment.npy")
+            )
+            data_list = [os.path.dirname(data) for data in data_list]
+        else:
+            data_list = glob.glob(os.path.join(self.data_root, split, "*"))
+        return data_list
+
+    def get_data_list(self):
+        if isinstance(self.split, str):
+            data_list = self.get_single_data_list(self.split)
+        elif isinstance(self.split, Sequence):
+            data_list = []
+            for split in self.split:
+                data_list += self.get_single_data_list(split)
+        else:
+            raise NotImplementedError
+        return data_list

pointcept/datasets/modelnet.py

@@ -0,0 +1,156 @@
+"""
+ModelNet40 Dataset
+
+get sampled point clouds of ModelNet40 (XYZ and normal from mesh, 10k points per shape)
+at "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip"
+
+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
+import copy
+import pointops
+import torch
+from torch.utils.data import Dataset
+from copy import deepcopy
+
+
+from pointcept.utils.logger import get_root_logger
+from .builder import DATASETS
+from .transform import Compose
+
+
+@DATASETS.register_module()
+class ModelNetDataset(Dataset):
+    def __init__(
+        self,
+        split="train",
+        data_root="data/modelnet40",
+        class_names=None,
+        transform=None,
+        num_points=8192,
+        uniform_sampling=True,
+        save_record=True,
+        test_mode=False,
+        test_cfg=None,
+        loop=1,
+        if_color=False,
+    ):
+        super().__init__()
+        self.data_root = data_root
+        self.class_names = dict(zip(class_names, range(len(class_names))))
+        self.split = split
+        self.num_point = num_points
+        self.uniform_sampling = uniform_sampling
+        self.transform = Compose(transform)
+        self.loop = (
+            loop if not test_mode else 1
+        )  # force make loop = 1 while in test mode
+        self.if_color = if_color
+        self.test_mode = test_mode
+        self.test_cfg = test_cfg if test_mode else None
+        if test_mode:
+            self.post_transform = Compose(self.test_cfg.post_transform)
+            self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform]
+
+        self.data_list = self.get_data_list()
+        logger = get_root_logger()
+        logger.info(
+            "Totally {} x {} samples in {} set.".format(
+                len(self.data_list), self.loop, split
+            )
+        )
+
+        # check, prepare record
+        record_name = f"modelnet40_{self.split}"
+        if num_points is not None:
+            record_name += f"_{num_points}points"
+            if uniform_sampling:
+                record_name += "_uniform"
+        record_path = os.path.join(self.data_root, f"{record_name}.pth")
+        if os.path.isfile(record_path):
+            logger.info(f"Loading record: {record_name} ...")
+            self.data = torch.load(record_path, weights_only=False)
+        else:
+            logger.info(f"Preparing record: {record_name} ...")
+            self.data = {}
+            for idx in range(len(self.data_list)):
+                data_name = self.data_list[idx]
+                logger.info(f"Parsing data [{idx}/{len(self.data_list)}]: {data_name}")
+                self.data[data_name] = self.get_data(idx)
+            if save_record:
+                torch.save(self.data, record_path)
+
+    def get_data(self, idx):
+        data_idx = idx % len(self.data_list)
+        data_name = self.data_list[data_idx]
+        if data_name in self.data.keys():
+            return copy.deepcopy(self.data[data_name])
+        else:
+            data_shape = "_".join(data_name.split("_")[0:-1])
+            data_path = os.path.join(
+                self.data_root, data_shape, self.data_list[data_idx] + ".txt"
+            )
+            data = np.loadtxt(data_path, delimiter=",").astype(np.float32)
+            if self.num_point is not None:
+                if self.uniform_sampling:
+                    with torch.no_grad():
+                        mask = pointops.farthest_point_sampling(
+                            torch.tensor(data).float().cuda(),
+                            torch.tensor([len(data)]).long().cuda(),
+                            torch.tensor([self.num_point]).long().cuda(),
+                        )
+                    data = data[mask.cpu()]
+                else:
+                    data = data[: self.num_point]
+            coord, normal = data[:, 0:3], data[:, 3:6]
+            category = np.array([self.class_names[data_shape]])
+            if self.if_color:
+                color = np.zeros_like(coord)
+                return dict(coord=coord, color=color, normal=normal, category=category)
+            return dict(coord=coord, normal=normal, category=category)
+
+    def get_data_list(self):
+        assert isinstance(self.split, str)
+        split_path = os.path.join(
+            self.data_root, "modelnet40_{}.txt".format(self.split)
+        )
+        data_list = np.loadtxt(split_path, dtype="str")
+        return data_list
+
+    def get_data_name(self, idx):
+        data_idx = idx % len(self.data_list)
+        return self.data_list[data_idx]
+
+    def __getitem__(self, idx):
+        if self.test_mode:
+            return self.prepare_test_data(idx)
+        else:
+            return self.prepare_train_data(idx)
+
+    def __len__(self):
+        return len(self.data_list) * self.loop
+
+    def prepare_train_data(self, idx):
+        data_dict = self.get_data(idx)
+        data_dict = self.transform(data_dict)
+        return data_dict
+
+    def prepare_test_data(self, idx):
+        assert idx < len(self.data_list)
+        data_dict = self.get_data(idx)
+        category = data_dict.pop("category")
+        data_dict = self.transform(data_dict)
+        data_dict_list = []
+        for aug in self.aug_transform:
+            data_dict_list.append(aug(copy.deepcopy(data_dict)))
+        for i in range(len(data_dict_list)):
+            data_dict_list[i] = self.post_transform(data_dict_list[i])
+        data_dict = dict(
+            voting_list=data_dict_list,
+            category=category,
+            name=self.get_data_name(idx),
+        )
+        return data_dict

pointcept/datasets/nuscenes.py

@@ -0,0 +1,583 @@
+"""
+nuScenes Dataset
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Zheng Zhang
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import numpy as np
+from collections.abc import Sequence
+import pickle
+from PIL import Image
+import open3d as o3d
+import torch
+
+from .builder import DATASETS
+from .defaults import DefaultDataset, DefaultImagePointDataset
+
+os.environ["OMP_NUM_THREADS"] = "1"
+
+
+@DATASETS.register_module()
+class NuScenesDataset(DefaultDataset):
+    def __init__(self, sweeps=10, ignore_index=-1, **kwargs):
+        self.sweeps = sweeps
+        self.ignore_index = ignore_index
+        self.learning_map = self.get_learning_map(ignore_index)
+        super().__init__(ignore_index=ignore_index, **kwargs)
+
+    def get_info_path(self, split):
+        assert split in ["train", "val", "test"]
+        if split == "train":
+            return os.path.join(
+                self.data_root, "info", f"nuscenes_infos_{self.sweeps}sweeps_train.pkl"
+            )
+        elif split == "val":
+            return os.path.join(
+                self.data_root, "info", f"nuscenes_infos_{self.sweeps}sweeps_val.pkl"
+            )
+        elif split == "test":
+            return os.path.join(
+                self.data_root, "info", f"nuscenes_infos_{self.sweeps}sweeps_test.pkl"
+            )
+        else:
+            raise NotImplementedError
+
+    def get_data_list(self):
+        if isinstance(self.split, str):
+            info_paths = [self.get_info_path(self.split)]
+        elif isinstance(self.split, Sequence):
+            info_paths = [self.get_info_path(s) for s in self.split]
+        else:
+            raise NotImplementedError
+        data_list = []
+        for info_path in info_paths:
+            with open(info_path, "rb") as f:
+                info = pickle.load(f)
+                data_list.extend(info)
+        return data_list
+
+    def get_data(self, idx):
+        data = self.data_list[idx % len(self.data_list)]
+        lidar_path = os.path.join(self.data_root, "raw", data["lidar_path"])
+        points = np.fromfile(str(lidar_path), dtype=np.float32, count=-1).reshape(
+            [-1, 5]
+        )
+        coord = points[:, :3]
+        strength = points[:, 3].reshape([-1, 1]) / 255  # scale strength to [0, 1]
+
+        if "gt_segment_path" in data.keys():
+            gt_segment_path = os.path.join(
+                self.data_root, "raw", data["gt_segment_path"]
+            )
+            segment = np.fromfile(
+                str(gt_segment_path), dtype=np.uint8, count=-1
+            ).reshape([-1])
+            segment = np.vectorize(self.learning_map.__getitem__)(segment).astype(
+                np.int64
+            )
+        else:
+            segment = np.ones((points.shape[0],), dtype=np.int64) * self.ignore_index
+        data_dict = dict(
+            coord=coord,
+            strength=strength,
+            segment=segment,
+            name=self.get_data_name(idx),
+        )
+        return data_dict
+
+    def get_data_name(self, idx):
+        # return data name for lidar seg, optimize the code when need to support detection
+        return self.data_list[idx % len(self.data_list)]["lidar_token"]
+
+    @staticmethod
+    def get_learning_map(ignore_index):
+        learning_map = {
+            0: ignore_index,
+            1: ignore_index,
+            2: 6,
+            3: 6,
+            4: 6,
+            5: ignore_index,
+            6: 6,
+            7: ignore_index,
+            8: ignore_index,
+            9: 0,
+            10: ignore_index,
+            11: ignore_index,
+            12: 7,
+            13: ignore_index,
+            14: 1,
+            15: 2,
+            16: 2,
+            17: 3,
+            18: 4,
+            19: ignore_index,
+            20: ignore_index,
+            21: 5,
+            22: 8,
+            23: 9,
+            24: 10,
+            25: 11,
+            26: 12,
+            27: 13,
+            28: 14,
+            29: ignore_index,
+            30: 15,
+            31: ignore_index,
+        }
+        return learning_map
+
+
+@DATASETS.register_module()
+class NuScenesColorNormalDataset(NuScenesDataset):
+    @staticmethod
+    def estimate_normals(points, center=np.array([0, 0, 0])):
+        normals = points - center[None, :]
+        norms = np.linalg.norm(normals, axis=1, keepdims=True)
+        normals = normals / norms
+        return normals
+
+    def get_data(self, idx):
+        data = self.data_list[idx % len(self.data_list)]
+        lidar_path = os.path.join(self.data_root, "raw", data["lidar_path"])
+        points = np.fromfile(str(lidar_path), dtype=np.float32, count=-1).reshape(
+            [-1, 5]
+        )
+        coord = points[:, :3]
+        normal = self.estimate_normals(coord)
+        if "gt_segment_path" in data.keys():
+            gt_segment_path = os.path.join(
+                self.data_root, "raw", data["gt_segment_path"]
+            )
+            segment = np.fromfile(
+                str(gt_segment_path), dtype=np.uint8, count=-1
+            ).reshape([-1])
+            segment = np.vectorize(self.learning_map.__getitem__)(segment).astype(
+                np.int64
+            )
+        else:
+            segment = np.ones((points.shape[0],), dtype=np.int64) * self.ignore_index
+        data_dict = dict(
+            coord=coord,
+            color=np.zeros_like(coord),  # placeholder for color
+            normal=np.zeros_like(coord),  # placeholder for normal
+            segment=segment,
+            name=self.get_data_name(idx),
+        )
+        return data_dict
+
+
+@DATASETS.register_module()
+class NuScenesImagePointDataset(DefaultImagePointDataset):
+    CAMERA_TYPES = [
+        "CAM_FRONT",
+        "CAM_FRONT_RIGHT",
+        "CAM_FRONT_LEFT",
+        "CAM_BACK",
+        "CAM_BACK_LEFT",
+        "CAM_BACK_RIGHT",
+    ]
+
+    def __init__(
+        self,
+        if_img=False,
+        if_sweep=False,
+        sweeps_max=10,
+        sweeps=10,
+        sweep_gap=1,
+        ignore_index=-1,
+        img_num=4,
+        **kwargs,
+    ):
+        self.sweeps = sweeps
+        self.sweep_gap = sweep_gap
+        self.sweeps_max = sweeps_max
+        self.if_sweep = if_sweep
+        self.if_img = if_img
+        self.ignore_index = ignore_index
+        self.learning_map = self.get_learning_map(ignore_index)
+        frame_count = sweeps + 1 if if_sweep else 1
+        self.img_ratio = img_num / (6 * frame_count)
+        super().__init__(ignore_index=ignore_index, if_img=if_img, **kwargs)
+
+    @staticmethod
+    def project_lidar_to_image_with_color(
+        lidar_points,  # shape: (N, 3) or (N, 4)
+        image,  # shape: (H, W, 3), uint8 RGB
+        cam_intrinsic,  # shape: (3, 3)
+        lidar_to_cam,  # shape: (4, 4)
+        lidar_colors,
+    ):
+        """
+        Projects LiDAR points to the image, fetches pixel color and pixel coordinates.
+        Returns:
+            filtered_points: (M, 3) - 3D points in camera frame that project onto the image.
+            colors:          (M, 3) - RGB colors at projected 2D locations.
+            uv_coords:       (M, 2) - Integer pixel coordinates (u, v) on the image.
+            mask:            (N,)   - (optional) Boolean mask indicating which lidar points are used.
+        """
+        lidar_uv_coords = np.full(
+            (lidar_points.shape[0], 2), -1, dtype=int
+        )  # Default to (-1, -1)
+        lidar_points_coord = lidar_points[:, :3]
+        ones = np.ones((lidar_points_coord.shape[0], 1))
+        lidar_hom = np.concatenate([lidar_points_coord, ones], axis=1)  # (N, 4)
+
+        points_cam = (lidar_to_cam @ lidar_hom.T).T  # (N, 4)
+
+        valid = points_cam[:, 2] > 0
+        points_cam = points_cam[valid]
+
+        pts_2d = (cam_intrinsic @ points_cam[:, :3].T).T  # (N, 3)
+        pts_2d = pts_2d[:, :2] / pts_2d[:, 2:3]  # (N, 2) - pixel (u, v)
+
+        H, W = image.shape[:2]
+        u, v = pts_2d[:, 0], pts_2d[:, 1]
+        inside = (u >= 0) & (u < W) & (v >= 0) & (v < H)
+
+        u = u[inside].astype(int)
+        v = v[inside].astype(int)
+
+        mask = np.zeros(lidar_points.shape[0], dtype=bool)
+        mask[np.where(valid)[0][inside]] = True
+
+        lidar_colors[mask] = image[v, u, :]
+        lidar_uv_coords[mask] = np.stack([u, v], axis=1)  # (M, 2)
+        return lidar_colors, lidar_uv_coords, mask
+
+    @staticmethod
+    def transform_normal(normal, transform):
+        return (transform[:3, :3] @ normal.T).T
+
+    def get_info_path(self, split):
+        assert split in ["train", "val", "test"]
+        if split == "train":
+            return os.path.join(
+                self.data_root,
+                "info",
+                f"nuscenes_infos_{self.sweeps_max}sweeps_train.pkl",
+            )
+        elif split == "val":
+            return os.path.join(
+                self.data_root,
+                "info",
+                f"nuscenes_infos_{self.sweeps_max}sweeps_val.pkl",
+            )
+        elif split == "test":
+            return os.path.join(
+                self.data_root,
+                "info",
+                f"nuscenes_infos_{self.sweeps_max}sweeps_test.pkl",
+            )
+        else:
+            raise NotImplementedError
+
+    def get_data_list(self):
+        split_list = {}
+        if isinstance(self.split, str):
+            info_paths = [self.get_info_path(self.split)]
+            split = [self.split]
+        elif isinstance(self.split, Sequence):
+            split = self.split
+            info_paths = [self.get_info_path(s) for s in self.split]
+        else:
+            raise NotImplementedError
+        data_list = []
+        for info_path, split_i in zip(info_paths, split):
+            with open(info_path, "rb") as f:
+                info = pickle.load(f)
+                data_list.extend(info)
+                split_list[split_i] = list([i["token"] for i in info])
+        return data_list, split_list
+
+    def get_data(self, idx):
+        data = self.data_list[idx % len(self.data_list)]
+        lidar_path = os.path.join(self.data_root, "raw", data["lidar_path"])
+        points = np.fromfile(str(lidar_path), dtype=np.float32, count=-1).reshape(
+            [-1, 5]
+        )
+        keyframe_point_count = points.shape[0]
+
+        imgs = []
+        cam_coords = []
+        cam_normals = []
+        cam_colors = []
+        cam_strengths = []
+        cam_correspondences = []
+        correspondence_start = 0
+        frame_pcd_offset = []
+        lidar_colors = np.zeros((points.shape[0], 3), dtype=int)  # Default to black
+        for id, cam_name in enumerate(self.CAMERA_TYPES):
+            cam_info = data["cams"][cam_name]
+            cam_intrinsic = cam_info["camera_intrinsics"]
+            cam_image = Image.open(
+                os.path.join(self.data_root, "raw", data["cams"][cam_name]["data_path"])
+            )
+            cam_image_np = np.array(cam_image)
+            sensor2lidar = np.eye(4)
+            sensor2lidar[:3, :3] = cam_info["sensor2lidar_rotation"]
+            sensor2lidar[:3, 3] = cam_info["sensor2lidar_translation"]
+            lidar2sensor = np.linalg.inv(sensor2lidar)
+            lidar_colors, correspondence_info, _ = (
+                self.project_lidar_to_image_with_color(
+                    points, cam_image_np, cam_intrinsic, lidar2sensor, lidar_colors
+                )
+            )
+            correspondence_point_id = (
+                np.array(range(correspondence_info.shape[0])).reshape((-1, 1))
+                + correspondence_start
+            )
+            correspondence_info = np.hstack(
+                [correspondence_info, correspondence_point_id]
+            )
+            if np.random.rand() < self.img_ratio:
+                cam_correspondences.append(correspondence_info)
+                imgs.append(cam_image)
+        correspondence_start += points.shape[0]
+        cam_coord = points[:, :3]
+        cam_center = np.array([0, 0, 0])
+        cam_normal = self.get_normals(cam_center, cam_coord)
+        cam_normals.append(cam_normal)
+        cam_strength = points[:, 3].reshape([-1, 1]) / 255
+        cam_coords.append(cam_coord)
+        cam_colors.append(lidar_colors)
+        cam_strengths.append(cam_strength)
+
+        if self.if_sweep:
+            frame_pcd_offset.append(points.shape[0])
+            for id, sweep in enumerate(
+                data["sweeps"][: (self.sweep_gap * self.sweeps) : self.sweep_gap]
+            ):
+                lidar_path = os.path.join(self.data_root, "raw", sweep["lidar_path"])
+                points = np.fromfile(
+                    str(lidar_path), dtype=np.float32, count=-1
+                ).reshape([-1, 5])
+                lidar_colors = np.zeros(
+                    (points.shape[0], 3), dtype=int
+                )  # Default to black
+                cam_lidar_tm = (
+                    sweep["transform_matrix"]
+                    if sweep["transform_matrix"] is not None
+                    else np.eye(4)
+                )
+                for id, cam_name in enumerate(self.CAMERA_TYPES):
+                    cam_info = sweep["cams"][cam_name]
+                    # cam_image_np = np.array(imgs[id])
+                    cam_intrinsic = cam_info["camera_intrinsics"]
+                    cam_image = Image.open(
+                        os.path.join(self.data_root, "raw", cam_info["data_path"])
+                    )
+                    cam_image_np = np.array(cam_image)
+                    sensor2lidar = np.eye(4)
+                    sensor2lidar[:3, :3] = cam_info["sensor2lidar_rotation"]
+                    sensor2lidar[:3, 3] = cam_info["sensor2lidar_translation"]
+                    lidar2sensor = np.linalg.inv(sensor2lidar) @ cam_lidar_tm
+                    lidar_colors, correspondence_info, _ = (
+                        self.project_lidar_to_image_with_color(
+                            points,
+                            cam_image_np,
+                            cam_intrinsic,
+                            lidar2sensor,
+                            lidar_colors,
+                        )
+                    )
+                    correspondence_point_id = (
+                        np.array(range(correspondence_info.shape[0])).reshape((-1, 1))
+                        + correspondence_start
+                    )
+                    correspondence_info = np.hstack(
+                        [correspondence_info, correspondence_point_id]
+                    )
+                    if np.random.rand() < self.img_ratio:
+                        cam_correspondences.append(correspondence_info)
+                        imgs.append(cam_image)
+                correspondence_start += correspondence_info.shape[0]
+                frame_pcd_offset.append(correspondence_start)
+                cam_coord = points[:, :3]
+                cam_center = np.array([0, 0, 0])
+                cam_normal = self.get_normals(cam_center, cam_coord)
+                cam_normal = self.transform_normal(cam_normal, cam_lidar_tm)
+                cam_normals.append(cam_normal)
+                ones = np.ones((points.shape[0], 1))
+                cam_coord_hom = np.concatenate([cam_coord, ones], axis=1)  # (N, 4)
+                cam_coord = cam_coord_hom @ cam_lidar_tm.T
+                cam_coord = cam_coord[:, :3]
+                cam_strength = points[:, 3].reshape([-1, 1]) / 255
+                cam_coords.append(cam_coord)
+                cam_colors.append(lidar_colors)
+                cam_strengths.append(cam_strength)
+
+        coord = np.vstack(cam_coords)
+        color = np.vstack(cam_colors)
+        normal = np.vstack(cam_normals)
+        strength = np.vstack(cam_strengths)
+        frame_pcd_offset = np.array(frame_pcd_offset)
+
+        car_from_ref = np.linalg.inv(data["ref_from_car"])
+        coord_homo = np.hstack((coord, np.ones((coord.shape[0], 1))))
+        coord_homo = coord_homo @ car_from_ref.T
+        coord = coord_homo[:, :3]
+        normal = self.transform_normal(normal, car_from_ref)
+
+        img_assets = dict()
+        if self.if_img:
+            if len(imgs) > 0:
+                img_width, img_height = imgs[0].size
+                div_w = img_width // self.patch_w
+                div_h = img_height // self.patch_h
+                div_min = max(min(div_w, div_h), 1)
+                crop_img_width = div_min * self.patch_w
+                crop_img_height = div_min * self.patch_h
+                left = int((img_width - crop_img_width) / 2)
+                top = int((img_height - crop_img_height) / 2)
+                right = int((img_width + crop_img_width) / 2)
+                bottom = int((img_height + crop_img_height) / 2)
+                imgs = [img.crop((left, top, right, bottom)) for img in imgs]
+                imgs = [self.transform_img(img) for img in imgs]
+                imgs_list = torch.stack(imgs)
+                img_assets["images"] = imgs_list.float()
+            else:
+                img_assets["images"] = torch.empty(
+                    (
+                        0,
+                        3,
+                        self.patch_h * self.patch_size,
+                        self.patch_w * self.patch_size,
+                    )
+                )
+            img_assets["img_num"] = np.array(
+                [img_assets["images"].shape[0]], dtype=np.int32
+            )
+
+            correspondence_infos = np.ones(
+                (coord.shape[0], len(cam_correspondences), 2), dtype=np.int32
+            ) * (-1)
+            for id, correspondence_info in enumerate(cam_correspondences):
+                correspondence_info = self.resize_correspondence_info(
+                    correspondence_info,
+                    (self.patch_h * self.patch_size, self.patch_w * self.patch_size),
+                    (img_height, img_width),
+                    (left, top, right, bottom),
+                    self.patch_size,
+                )
+                correspondence_infos[correspondence_info[:, -1], id, :] = (
+                    correspondence_info[:, :-1]
+                )
+            img_assets["correspondence"] = correspondence_infos
+        if "gt_segment_path" in data.keys():
+            gt_segment_path = os.path.join(
+                self.data_root, "raw", data["gt_segment_path"]
+            )
+            segment = np.fromfile(
+                str(gt_segment_path), dtype=np.uint8, count=-1
+            ).reshape([-1])
+            segment = np.vectorize(self.learning_map.__getitem__)(segment).astype(
+                np.int64
+            )
+        else:
+            segment = (
+                np.ones((keyframe_point_count,), dtype=np.int64) * self.ignore_index
+            )
+
+        if segment.shape[0] != keyframe_point_count:
+            raise ValueError(
+                f"nuScenes sample {self.get_data_name(idx)} has "
+                f"{keyframe_point_count} keyframe points but "
+                f"{segment.shape[0]} segment labels."
+            )
+
+        if self.if_sweep:
+            sweep_point_count = coord.shape[0] - keyframe_point_count
+            sweep_segment = np.full(
+                (sweep_point_count,), self.ignore_index, dtype=segment.dtype
+            )
+            segment = np.concatenate((segment, sweep_segment), axis=0)
+
+        color = color.astype(np.float32)
+        if self.if_sweep:
+            data_dict = dict(
+                coord=coord,
+                color=color,
+                normal=normal,
+                strength=strength,
+                segment=segment,
+                frame_pcd_offset=frame_pcd_offset,
+                name=self.get_data_name(idx),
+            )
+        else:
+            data_dict = dict(
+                coord=coord,
+                color=color,
+                normal=normal,
+                strength=strength,
+                segment=segment,
+                name=self.get_data_name(idx),
+            )
+        data_dict.update(img_assets)
+        return data_dict
+
+    def get_data_name(self, idx):
+        return self.data_list[idx % len(self.data_list)]["lidar_token"]
+
+    @staticmethod
+    def get_normals(cam_center, coords):
+        Cs = np.repeat(cam_center.reshape((1, -1)), coords.shape[0], axis=0)
+        view_dirs = coords - Cs
+        view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True)
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(coords)
+        pcd.estimate_normals()
+        normals = np.asarray(pcd.normals)
+        dot_product = np.sum(normals * view_dirs, axis=-1)
+        flip_mask = dot_product > 0
+        normals[flip_mask] = -normals[flip_mask]
+        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
+        return normals
+
+    # or using the vector from LiDAR to points, which is faster for data loading
+    # @staticmethod
+    # def get_normals(center, points):
+    #     normals = points - center[None, :]
+    #     norms = np.linalg.norm(normals, axis=1, keepdims=True)
+    #     normals = normals / norms
+    #     return normals
+
+    @staticmethod
+    def get_learning_map(ignore_index):
+        learning_map = {
+            0: ignore_index,
+            1: ignore_index,
+            2: 6,
+            3: 6,
+            4: 6,
+            5: ignore_index,
+            6: 6,
+            7: ignore_index,
+            8: ignore_index,
+            9: 0,
+            10: ignore_index,
+            11: ignore_index,
+            12: 7,
+            13: ignore_index,
+            14: 1,
+            15: 2,
+            16: 2,
+            17: 3,
+            18: 4,
+            19: ignore_index,
+            20: ignore_index,
+            21: 5,
+            22: 8,
+            23: 9,
+            24: 10,
+            25: 11,
+            26: 12,
+            27: 13,
+            28: 14,
+            29: ignore_index,
+            30: 15,
+            31: ignore_index,
+        }
+        return learning_map

pointcept/datasets/partnet.py

@@ -0,0 +1,79 @@
+"""
+HK Remote Data Dataset
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import numpy as np
+import open3d as o3d
+import torch
+
+from .defaults import DefaultDataset
+from .builder import DATASETS
+from pointcept.utils.cache import shared_dict
+
+
+@DATASETS.register_module()
+class PartNetDataDataset(DefaultDataset):
+    def __init__(self, if_img=True, crop_h=630, crop_w=1120, patch_size=14, **kwargs):
+        self.crop_h = crop_h
+        self.crop_w = crop_w
+        self.patch_size = patch_size
+        self.patch_h = crop_h // patch_size
+        self.patch_w = crop_w // patch_size
+        self.if_img = if_img
+        super().__init__(**kwargs)
+
+    def get_data(self, idx):
+        data_path = self.data_list[idx % len(self.data_list)]
+        pc_data_path = os.path.join(
+            data_path, "point_sample", "sample-points-all-pts-nor-rgba-10000.ply"
+        )
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        if self.cache:
+            cache_name = f"pointcept-{name}"
+            return shared_dict(cache_name)
+
+        data_dict = {}
+        pc = o3d.io.read_point_cloud(pc_data_path)
+        data_dict["coord"] = np.asarray(pc.points)
+        data_dict["color"] = np.asarray(pc.colors)
+        data_dict["normal"] = np.asarray(pc.normals)
+        data_dict["name"] = name
+        data_dict["split"] = split
+
+        if self.if_img:
+            data_dict["images"] = torch.empty(
+                (0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size)
+            )
+            data_dict["img_num"] = np.array([0], dtype=np.int32)
+            data_dict["correspondence"] = np.ones(
+                (data_dict["coord"].shape[0], 0, 2), dtype=np.float32
+            ) * (-1)
+
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+        return data_dict

pointcept/datasets/partnete.py

@@ -0,0 +1,111 @@
+"""
+PartNetE Datasets
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import glob
+from collections.abc import Sequence
+import numpy as np
+import json
+
+from .defaults import DefaultDataset
+from .builder import DATASETS
+
+from pointcept.utils.cache import shared_dict
+
+
+@DATASETS.register_module()
+class PartNetEDataset(DefaultDataset):
+    def __init__(
+        self, class_names, num_parts, data_root, meta_path, split, *args, **kwargs
+    ):
+        self.categories = class_names
+        self.num_parts = num_parts
+        self.num_part_offset = np.concatenate(([0], np.cumsum(self.num_parts)))
+        with open(meta_path, "r", encoding="utf-8") as file:
+            meta_data = json.load(file)
+        self.category2part = {}
+        self.parts = []
+        for class_id, class_name in enumerate(self.categories):
+            self.category2part[class_name] = self.num_part_offset[class_id] + list(
+                range(self.num_parts[class_id])
+            )
+            category_part_name = [
+                class_name + "_" + part_name
+                for part_name in ["other"] + meta_data[class_name]
+            ]
+            self.parts.extend(category_part_name)
+        super().__init__(data_root=data_root, split=split, *args, **kwargs)
+
+    def get_data_list(self):
+        if isinstance(self.split, str):
+            data_list = glob.glob(os.path.join(self.data_root, self.split, "*/*"))
+        elif isinstance(self.split, Sequence):
+            data_list = []
+            for split in self.split:
+                data_list += glob.glob(os.path.join(self.data_root, split, "*/*"))
+        else:
+            raise NotImplementedError
+        return data_list
+
+    def get_data_name(self, idx):
+        file_path = self.data_list[idx % len(self.data_list)]
+        dir_path, id_name = os.path.split(file_path)
+        object_name = os.path.basename(dir_path)
+        data_name = f"{object_name}_{id_name}"
+        return data_name
+
+    def get_data(self, idx):
+        data_path = self.data_list[idx % len(self.data_list)]
+        name = self.get_data_name(idx)
+        split = self.get_split_name(idx)
+        if self.cache:
+            cache_name = f"pointcept-{name}"
+            return shared_dict(cache_name)
+
+        data_dict = {}
+        assets = os.listdir(data_path)
+        for asset in assets:
+            if not asset.endswith(".npy"):
+                continue
+            if asset[:-4] not in self.VALID_ASSETS:
+                continue
+            data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
+        data_dict["name"] = name
+        data_dict["split"] = split
+
+        object_name = name.split("_")[0]
+        cls_token = self.categories.index(object_name)
+        data_dict["cls_token"] = cls_token
+
+        if "coord" in data_dict.keys():
+            data_dict["coord"] = data_dict["coord"].astype(np.float32)
+
+        if "color" in data_dict.keys():
+            data_dict["color"] = data_dict["color"].astype(np.float32)
+
+        if "normal" in data_dict.keys():
+            data_dict["normal"] = data_dict["normal"].astype(np.float32)
+
+        if "segment" in data_dict.keys():
+            data_dict["segment"] = (
+                data_dict["segment"].reshape([-1]).astype(np.int32)
+                + self.num_part_offset[cls_token]
+                + 1
+            )
+        else:
+            data_dict["segment"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        if "instance" in data_dict.keys():
+            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
+        else:
+            data_dict["instance"] = (
+                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
+            )
+
+        return data_dict

pointcept/datasets/preprocessing/arkitscenes/preprocess_arkitscenes_mesh.py

@@ -0,0 +1,93 @@
+"""
+Preprocessing ArkitScenes
+"""
+
+import os
+import argparse
+import glob
+import plyfile
+import numpy as np
+import pandas as pd
+import multiprocessing as mp
+from concurrent.futures import ProcessPoolExecutor
+from itertools import repeat
+from pathlib import Path
+
+
+def read_plymesh(filepath):
+    """Read ply file and return it as numpy array. Returns None if emtpy."""
+    with open(filepath, "rb") as f:
+        plydata = plyfile.PlyData.read(f)
+    if plydata.elements:
+        vertices = pd.DataFrame(plydata["vertex"].data).values
+        faces = np.stack(plydata["face"].data["vertex_indices"], axis=0)
+        return vertices, faces
+
+
+def face_normal(vertex, face):
+    v01 = vertex[face[:, 1]] - vertex[face[:, 0]]
+    v02 = vertex[face[:, 2]] - vertex[face[:, 0]]
+    vec = np.cross(v01, v02)
+    length = np.sqrt(np.sum(vec**2, axis=1, keepdims=True)) + 1.0e-8
+    nf = vec / length
+    area = length * 0.5
+    return nf, area
+
+
+def vertex_normal(vertex, face):
+    nf, area = face_normal(vertex, face)
+    nf = nf * area
+
+    nv = np.zeros_like(vertex)
+    for i in range(face.shape[0]):
+        nv[face[i]] += nf[i]
+
+    length = np.sqrt(np.sum(nv**2, axis=1, keepdims=True)) + 1.0e-8
+    nv = nv / length
+    return nv
+
+
+def parse_scene(scene_path, output_dir):
+    print(f"Parsing scene {scene_path}")
+    split = os.path.basename(os.path.dirname(os.path.dirname(scene_path)))
+    scene_id = os.path.basename(os.path.dirname(scene_path))
+    vertices, faces = read_plymesh(scene_path)
+    coords = vertices[:, :3]
+    colors = vertices[:, 3:6]
+    normals = vertex_normal(coords, faces)
+    data_dict = dict(coord=coords, color=colors, normal=normals)
+    os.makedirs(output_dir / split / scene_id, exist_ok=True)
+    for key in data_dict.keys():
+        np.save(output_dir / split / scene_id / f"{key}.npy", data_dict[key])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ArkitScenes dataset containing 3dod folder",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    opt = parser.parse_args()
+    # Create output directories
+    train_output_dir = os.path.join(opt.output_root, "Training")
+    os.makedirs(train_output_dir, exist_ok=True)
+    val_output_dir = os.path.join(opt.output_root, "Validation")
+    os.makedirs(val_output_dir, exist_ok=True)
+    # Load scene paths
+    scene_paths = sorted(glob.glob(opt.dataset_root + "/3dod/*/*/*_mesh.ply"))
+    # Preprocess data.
+    pool = ProcessPoolExecutor(max_workers=opt.num_workers)
+    print("Processing scenes...")
+    _ = list(pool.map(parse_scene, scene_paths, repeat(Path(opt.output_root))))

pointcept/datasets/preprocessing/concerto/arkitscenes/preprocess_arkitscenes.py

@@ -0,0 +1,343 @@
+"""
+Preprocessing Script for ARKitScenes
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+from pathlib import Path
+import argparse
+import os
+import pandas as pd
+import camtools as ct
+import numpy as np
+from scipy.spatial import cKDTree
+import shutil
+import open3d as o3d
+import cv2
+import glob
+import numpy as np
+import os
+import plyfile
+import multiprocessing as mp
+from rotation import convert_angle_axis_to_matrix3
+
+
+def correspondenceGet(mesh, K, T, img_size, coords_gt):
+    height, width = img_size
+    if np.isnan(T).any() or np.isinf(T).any():
+        return None
+    depth = ct.raycast.mesh_to_im_depth(
+        mesh=mesh, K=K, T=np.linalg.inv(T), height=height, width=width
+    )
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    valid = ~np.isinf(depth).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict
+
+
+def correspondenceSave(mesh, scene_dir, coords_gt, output_dir, img_size, Ks, Ts):
+    os.makedirs(output_dir, exist_ok=True)
+    scene_dir = Path(scene_dir)
+    index_gt = [
+        img_name[:-4].split("_")[1]
+        for img_name in os.listdir(str(scene_dir / "color"))
+        if img_name.endswith(".png")
+    ]
+    index_gt = sorted(index_gt, key=lambda x: float(x))
+
+    coords_gt_ = coords_gt
+    pixels_ = []
+    coords_ = []
+
+    for i, (K, T) in enumerate(zip(Ks, Ts)):
+        coord_dict = correspondenceGet(mesh, K, T, img_size, coords_gt)
+        if coord_dict is None:
+            correspondences = -np.ones((1, 3))
+        else:
+            pixels_ = coord_dict["pixel"]
+            coords_ = coord_dict["coord"]
+            tree = cKDTree(coords_gt_)
+            dis, idx = tree.query(coords_, k=1)
+            idx_valid = idx[dis < 0.01]
+            pixels_valid = pixels_[dis < 0.01]
+            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+        np.save(Path(output_dir) / "{}.npy".format(index_gt[i]), correspondences)
+
+
+def read_plymesh(filepath):
+    """Read ply file and return it as numpy array. Returns None if emtpy."""
+    with open(filepath, "rb") as f:
+        plydata = plyfile.PlyData.read(f)
+    if plydata.elements:
+        vertices = pd.DataFrame(plydata["vertex"].data).values
+        faces = np.stack(plydata["face"].data["vertex_indices"], axis=0)
+        return vertices, faces
+
+
+def face_normal(vertex, face):
+    v01 = vertex[face[:, 1]] - vertex[face[:, 0]]
+    v02 = vertex[face[:, 2]] - vertex[face[:, 0]]
+    vec = np.cross(v01, v02)
+    length = np.sqrt(np.sum(vec**2, axis=1, keepdims=True)) + 1.0e-8
+    nf = vec / length
+    area = length * 0.5
+    return nf, area
+
+
+def vertex_normal(vertex, face):
+    nf, area = face_normal(vertex, face)
+    nf = nf * area
+
+    nv = np.zeros_like(vertex)
+    for i in range(face.shape[0]):
+        nv[face[i]] += nf[i]
+
+    length = np.sqrt(np.sum(nv**2, axis=1, keepdims=True)) + 1.0e-8
+    nv = nv / length
+    return nv
+
+
+def TrajStringToMatrix(traj_str):
+    """convert traj_str into translation and rotation matrices
+    Args:
+        traj_str: A space-delimited file where each line represents a camera position at a particular timestamp.
+        The file has seven columns:
+        * Column 1: timestamp
+        * Columns 2-4: rotation (axis-angle representation in radians)
+        * Columns 5-7: translation (usually in meters)
+
+    Returns:
+        ts: translation matrix
+        Rt: rotation matrix
+    """
+    # line=[float(x) for x in traj_str.split()]
+    # ts = line[0];
+    # R = cv2.Rodrigues(np.array(line[1:4]))[0];
+    # t = np.array(line[4:7]);
+    # Rt = np.concatenate((np.concatenate((R, t[:,np.newaxis]), axis=1), [[0.0,0.0,0.0,1.0]]), axis=0)
+    tokens = traj_str.split()
+    assert len(tokens) == 7
+    ts = tokens[0]
+    # Rotation in angle axis
+    angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
+    r_w_to_p = convert_angle_axis_to_matrix3(np.asarray(angle_axis))
+    # Translation
+    t_w_to_p = np.asarray([float(tokens[4]), float(tokens[5]), float(tokens[6])])
+    extrinsics = np.eye(4, 4)
+    extrinsics[:3, :3] = r_w_to_p
+    extrinsics[:3, -1] = t_w_to_p
+    Rt = np.linalg.inv(extrinsics)
+    return (ts, Rt)
+
+
+def st2_camera_intrinsics(filename):
+    w, h, fx, fy, hw, hh = np.loadtxt(filename)
+    return np.asarray([[fx, 0, hw], [0, fy, hh], [0, 0, 1]])
+
+
+def handle_process(
+    scene_path,
+    output_path,
+    pointclouds_root,
+    frame_gap,
+    parse_pointclouds,
+    parse_depths,
+):
+    print("loading from ", scene_path)
+    split = os.path.basename(os.path.dirname(scene_path))
+    scene_id = os.path.basename(scene_path)
+    img_folder = os.path.join(scene_path, f"{scene_id}_frames", "lowres_wide")
+    traj_file = os.path.join(scene_path, f"{scene_id}_frames", "lowres_wide.traj")
+    pc_output_path = os.path.join(output_path, split, f"{scene_id}")
+    pc_input_path = os.path.join(pointclouds_root, split, f"{scene_id}")
+    im_output_path = os.path.join(output_path, "images", split, f"{scene_id}")
+    os.makedirs(im_output_path, exist_ok=True)
+    if not os.path.exists(img_folder):
+        frame_ids = []
+    else:
+        images = sorted(
+            glob.glob(os.path.join(img_folder, "*.png")),
+            key=lambda x: float(x.split("/")[-1].split("_")[1][:-4]),
+        )
+        frame_ids = [os.path.basename(x) for x in images]
+        frame_ids = [x.split(".png")[0].split("_")[1] for x in frame_ids]
+        video_id = img_folder.split("/")[-2].split("_frames")[0]
+        frame_ids.sort()
+        frame_ids = frame_ids[::frame_gap]
+    if len(frame_ids) > 0:
+        height, width = cv2.imread(images[0]).shape[:2]
+        if os.path.exists(traj_file):
+            with open(traj_file) as f:
+                traj = f.readlines()
+            # convert traj to json dict
+            poses_from_traj = {}
+            for line in traj:
+                traj_timestamp = line.split(" ")[0]
+                poses_from_traj[f"{round(float(traj_timestamp), 3):.3f}"] = (
+                    TrajStringToMatrix(line)[1].tolist()
+                )
+
+            if os.path.exists(traj_file):
+                poses = poses_from_traj
+            else:
+                poses = {}
+
+            # get img dlc
+            poses_keys = list(poses.keys())
+            intrinsics = []
+            poses_list = []
+            Ks_path = os.path.join(im_output_path, "intrinsic")
+            Ts_path = os.path.join(im_output_path, "pose")
+            img_path = os.path.join(im_output_path, "color")
+
+            os.makedirs(Ks_path, exist_ok=True)
+            os.makedirs(Ts_path, exist_ok=True)
+            os.makedirs(img_path, exist_ok=True)
+            if parse_depths:
+                depth_path = os.path.join(im_output_path, "depth")
+                os.makedirs(depth_path, exist_ok=True)
+            for frame_id in frame_ids:
+                shutil.copy(
+                    os.path.join(
+                        scene_path,
+                        f"{scene_id}_frames",
+                        "lowres_wide",
+                        f"{video_id}_{frame_id}.png",
+                    ),
+                    img_path,
+                )
+                if parse_depths:
+                    shutil.copy(
+                        os.path.join(
+                            scene_path,
+                            f"{scene_id}_frames",
+                            "lowres_depth",
+                            f"{video_id}_{frame_id}.png",
+                        ),
+                        depth_path,
+                    )
+                intrinsic_fn = os.path.join(
+                    scene_path,
+                    f"{scene_id}_frames",
+                    "lowres_wide_intrinsics",
+                    f"{video_id}_{frame_id}.pincam",
+                )
+                intrinsic = st2_camera_intrinsics(intrinsic_fn)
+                intrinsics.append(intrinsic)
+                np.save(os.path.join(Ks_path, f"{video_id}_{frame_id}.npy"), intrinsic)
+                frame_id = f"{round(float(frame_id), 3):.3f}"
+                closest_key = min(
+                    poses_keys, key=lambda k: abs(float(k) - float(frame_id))
+                )
+                np.save(
+                    os.path.join(Ts_path, f"{video_id}_{frame_id}.npy"),
+                    poses[closest_key],
+                )
+                poses_list.append(poses[closest_key])
+            # match extrinsics
+            intrinsics = np.stack(intrinsics)
+            poses = np.stack(poses_list)
+
+            ply_path = os.path.join(scene_path, f"{scene_id}_3dod_mesh.ply")
+            if parse_pointclouds:
+                vertices, faces = read_plymesh(ply_path)
+                coords = vertices[:, :3]
+                colors = vertices[:, 3:6]
+                normals = vertex_normal(coords, faces)
+                data_dict = dict(coord=coords, color=colors, normal=normals)
+                os.makedirs(pc_output_path, exist_ok=True)
+                for key in data_dict.keys():
+                    np.save(os.path.join(pc_output_path, f"{key}.npy"), data_dict[key])
+            else:
+                coords = np.load(os.path.join(pc_input_path, "coord.npy"))
+
+            # Save img data
+            mesh = o3d.io.read_triangle_mesh(ply_path)
+            correspondenceSave(
+                mesh,
+                im_output_path,
+                coords,
+                os.path.join(im_output_path, "correspondence"),
+                (height, width),
+                intrinsics,
+                poses,
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ArkitScenes dataset containing 3dod folder",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--pointclouds_root",
+        default="data/arkitscenes",
+        type=str,
+        help="Input path where previous pointclouds folder located",
+    )
+    parser.add_argument(
+        "--frame_gap", default=50, type=int, help="Frame gap for processing"
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="thread_id",
+    )
+    parser.add_argument(
+        "--parse_pointclouds", action="store_true", help="Whether parse point clouds"
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    opt = parser.parse_args()
+    # Load scene paths
+    metadata_path = f"{opt.dataset_root}/3dod/metadata.csv"
+    metadata = pd.read_csv(metadata_path)
+    # Initialize the splits dictionary
+    splits = {}
+    scene_paths = []
+    # Populate the splits dictionary
+    for index, row in metadata.iterrows():
+        scene_id = str(row["video_id"])
+        split = row["fold"]
+        splits[scene_id] = split
+        scene_paths.append(os.path.join(opt.dataset_root, "3dod", split, scene_id))
+    # Preprocess data.
+    scene_paths_list = np.array_split(scene_paths, opt.num_workers)
+    scene_paths_ = scene_paths_list[opt.thread_id]
+    print("Processing scenes...")
+    for scene_path in scene_paths_:
+        handle_process(
+            scene_path,
+            opt.output_root,
+            opt.pointclouds_root,
+            opt.frame_gap,
+            opt.parse_pointclouds,
+            opt.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/arkitscenes/preprocess_arkitscenes.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=16  
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    p) parse_depths=true ;;
+    c) parse_pointclouds=true ;;
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/arkitscenes/preprocess_arkitscenes.py --thread_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root"
+
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/arkitscenes/rotation.py

@@ -0,0 +1,140 @@
+import cv2
+import math
+import numpy as np
+
+
+def eulerAnglesToRotationMatrix(theta):
+    """Euler rotation matrix with clockwise logic.
+    Rotation
+
+    Args:
+        theta: list of float
+            [theta_x, theta_y, theta_z]
+    Returns:
+        R: np.array (3, 3)
+            rotation matrix of Rz*Ry*Rx
+    """
+    R_x = np.array(
+        [
+            [1, 0, 0],
+            [0, math.cos(theta[0]), -math.sin(theta[0])],
+            [0, math.sin(theta[0]), math.cos(theta[0])],
+        ]
+    )
+
+    R_y = np.array(
+        [
+            [math.cos(theta[1]), 0, math.sin(theta[1])],
+            [0, 1, 0],
+            [-math.sin(theta[1]), 0, math.cos(theta[1])],
+        ]
+    )
+
+    R_z = np.array(
+        [
+            [math.cos(theta[2]), -math.sin(theta[2]), 0],
+            [math.sin(theta[2]), math.cos(theta[2]), 0],
+            [0, 0, 1],
+        ]
+    )
+
+    R = np.dot(R_z, np.dot(R_y, R_x))
+    return R
+
+
+def upright_camera_relative_transform(pose):
+    """Generate pose matrix with z-dim as height
+
+    Args:
+        pose: np.array (4, 4)
+    Returns:
+        urc: (4, 4)
+        urc_inv: (4, 4)
+    """
+
+    # take viewing direction in camera local coordiantes (which is simply unit vector along +z)
+    view_dir_camera = np.asarray([0, 0, 1])
+    R = pose[0:3, 0:3]
+    t = pose[0:3, 3]
+
+    # convert to world coordinates
+    view_dir_world = np.dot(R, view_dir_camera)
+
+    # compute heading
+    view_dir_xy = view_dir_world[0:2]
+    heading = math.atan2(view_dir_xy[1], view_dir_xy[0])
+
+    # compute rotation around Z to align heading with +Y
+    zRot = -heading + math.pi / 2
+
+    # translation first, back to camera point
+    urc_t = np.identity(4)
+    urc_t[0:2, 3] = -1 * t[0:2]
+
+    # compute rotation matrix
+    urc_r = np.identity(4)
+    urc_r[0:3, 0:3] = eulerAnglesToRotationMatrix([0, 0, zRot])
+
+    urc = np.dot(urc_r, urc_t)
+    urc_inv = np.linalg.inv(urc)
+
+    return urc, urc_inv
+
+
+def rotate_pc(pc, rotmat):
+    """Rotation points w.r.t. rotmat
+    Args:
+        pc: np.array (n, 3)
+        rotmat: np.array (4, 4)
+    Returns:
+        pc: (n, 3)
+    """
+    pc_4 = np.ones([pc.shape[0], 4])
+    pc_4[:, 0:3] = pc
+    pc_4 = np.dot(pc_4, np.transpose(rotmat))
+
+    return pc_4[:, 0:3]
+
+
+def rotate_points_along_z(points, angle):
+    """Rotation clockwise
+    Args:
+        points: np.array of np.array (B, N, 3 + C) or
+            (N, 3 + C) for single batch
+        angle: np.array of np.array (B, )
+            or (, ) for single batch
+            angle along z-axis, angle increases x ==> y
+    Returns:
+        points_rot:  (B, N, 3 + C) or (N, 3 + C)
+
+    """
+    single_batch = len(points.shape) == 2
+    if single_batch:
+        points = np.expand_dims(points, axis=0)
+        angle = np.expand_dims(angle, axis=0)
+    cosa = np.expand_dims(np.cos(angle), axis=1)
+    sina = np.expand_dims(np.sin(angle), axis=1)
+    zeros = np.zeros_like(cosa)  # angle.new_zeros(points.shape[0])
+    ones = np.ones_like(sina)  # angle.new_ones(points.shape[0])
+
+    rot_matrix = np.concatenate(
+        (cosa, -sina, zeros, sina, cosa, zeros, zeros, zeros, ones), axis=1
+    ).reshape(-1, 3, 3)
+
+    # print(rot_matrix.view(3, 3))
+    points_rot = np.matmul(points[:, :, :3], rot_matrix)
+    points_rot = np.concatenate((points_rot, points[:, :, 3:]), axis=-1)
+
+    if single_batch:
+        points_rot = points_rot.squeeze(0)
+
+    return points_rot
+
+
+def convert_angle_axis_to_matrix3(angle_axis):
+    """Return a Matrix3 for the angle axis.
+    Arguments:
+        angle_axis {Point3} -- a rotation in angle axis form.
+    """
+    matrix, jacobian = cv2.Rodrigues(angle_axis)
+    return matrix

pointcept/datasets/preprocessing/concerto/arkitscenes/splits.py

@@ -0,0 +1,64 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["Training", "Validation"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(
+            dataset_path, "data/arkitscenes"
+        )
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            im_split_name_path = os.path.join(im_split_path, name, "color")
+            co_split_name_path = os.path.join(im_split_path, name, "correspondence")
+            png_files = [
+                f for f in os.listdir(im_split_name_path) if f.endswith(".png")
+            ]
+            png_files = sorted(png_files, key=lambda x: float(x.split("_")[1][:-4]))
+            co_files = [f for f in os.listdir(co_split_name_path) if f.endswith(".npy")]
+            co_files = sorted(co_files, key=lambda x: float(x[:-4]))
+            # Get the full paths of the .png files
+            png_file_paths = [
+                os.path.join(im_split_name_path, f).replace(
+                    dataset_path, "data/arkitscenes"
+                )
+                for f in png_files
+            ]
+            co_file_paths = [
+                os.path.join(co_split_name_path, f).replace(
+                    dataset_path, "data/arkitscenes"
+                )
+                for f in co_files
+            ]
+            for i in range(0, len(png_file_paths), 4):
+                split_dict[f"{name}_{i//4}"] = {}
+                split_dict[f"{name}_{i//4}"]["pointclouds"] = os.path.join(
+                    pc_split_path, name
+                )
+                split_dict[f"{name}_{i//4}"]["images"] = png_file_paths[i : i + 4]
+                split_dict[f"{name}_{i//4}"]["correspondences"] = co_file_paths[
+                    i : i + 4
+                ]
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the arkitscenes dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/cap3d/preprocess_cap3d.py

@@ -0,0 +1,261 @@
+"""
+Preprocessing Script for Cap3D
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import argparse
+import torch
+import numpy as np
+from scipy.spatial import cKDTree
+from PIL import Image
+from pathlib import Path
+import multiprocessing as mp
+import zipfile
+import glob
+import json
+import math
+
+NUM_VIEWS = 4
+MAX_DEPTH = 5.0
+
+
+def recover_depth_from_image(depth_image_path, max_depth):
+    with Image.open(depth_image_path) as depth_pil:
+        depth_uint16 = np.array(depth_pil)
+    invalid_depth_mask = depth_uint16 == 65535
+    depth_normalized_float = depth_uint16.astype(np.float32) / 65535.0
+    depth_meters = depth_normalized_float * max_depth
+    return depth_meters, invalid_depth_mask
+
+
+def normalize_point_cloud(point_cloud: np.ndarray) -> np.ndarray:
+    bbox_min = np.min(point_cloud, axis=0)
+    bbox_max = np.max(point_cloud, axis=0)
+    bbox_size = bbox_max - bbox_min
+    max_dim = np.max(bbox_size)
+    if max_dim < 1e-8:
+        return point_cloud - bbox_min
+
+    scale = 1.0 / max_dim
+    scaled_point_cloud = point_cloud * scale
+    new_bbox_min = np.min(scaled_point_cloud, axis=0)
+    new_bbox_max = np.max(scaled_point_cloud, axis=0)
+    center = (new_bbox_min + new_bbox_max) / 2.0
+    offset = -center
+    normalized_point_cloud = scaled_point_cloud + offset
+
+    return normalized_point_cloud
+
+
+def get_cam_params(cam_params_file, width, height):
+    with open(cam_params_file, "r") as f:
+        data = json.load(f)
+
+    x_vector = data["x"]
+    y_vector = data["y"]
+    z_vector = data["z"]
+    origin = data["origin"]
+
+    rotation_matrix = np.array([x_vector, y_vector, z_vector]).T
+
+    translation_vector = np.array(origin)
+
+    rt_matrix = np.eye(4)
+    rt_matrix[:3, :3] = rotation_matrix
+    rt_matrix[:3, 3] = translation_vector
+
+    x_fov = data["x_fov"]
+    y_fov = data["y_fov"]
+
+    cx = width / 2.0
+    cy = height / 2.0
+
+    fx = (width / 2.0) / math.tan(x_fov / 2.0)
+    fy = (height / 2.0) / math.tan(y_fov / 2.0)
+
+    s = 0
+
+    K = np.array([[fx, s, cx], [0, fy, cy], [0, 0, 1]])
+
+    return rt_matrix, K
+
+
+def correspondenceGet(depth, invalid_depth_mask, T, K, width, height):
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    invalid_depth_mask = invalid_depth_mask.reshape((-1, 1))
+    valid = ~np.logical_or(np.isinf(depth), invalid_depth_mask).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict
+
+
+def correspondenceSave(
+    depth_paths, data_name, point_cloud_gt, cam_params_paths, output_root
+):
+    """
+    Calculates and saves the correspondences for all views of a given mesh and point cloud.
+    """
+    # Build a KDTree for the ground truth point cloud for fast querying.
+    tree = cKDTree(point_cloud_gt)
+
+    # Iterate over each predefined view.
+    for view_idx in range(len(depth_paths)):
+        # Construct file paths.
+        cam_params_path = cam_params_paths[view_idx]
+        depth_path = depth_paths[view_idx]
+
+        view_name = cam_params_path.split("/")[-1][:-5]
+
+        # Load camera parameters and image dimensions.
+        depth, invalid_depth_mask = recover_depth_from_image(depth_path, MAX_DEPTH)
+        height, width = depth.shape
+        T, K = get_cam_params(cam_params_path, width, height)
+
+        coord_dict = correspondenceGet(depth, invalid_depth_mask, T, K, width, height)
+        if coord_dict is None:
+            correspondences = -np.ones((1, 3))
+        else:
+            pixels_ = coord_dict["pixel"]
+            coords_ = coord_dict["coord"]
+            dis, idx = tree.query(coords_, k=1)
+            idx_valid = idx[dis < 0.01]
+            pixels_valid = pixels_[dis < 0.01]
+            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+        output_dir = output_root / data_name
+        os.makedirs(output_dir, exist_ok=True)
+        output_filename = output_dir / f"{view_name}.npy"
+        np.save(output_filename, correspondences)
+
+
+def handle_process(pt_path, output_root, cam_root):
+    """
+    The complete processing pipeline for a single GLB file.
+    """
+    data_name = pt_path.stem
+
+    if not pt_path.exists():
+        return
+    cam_path = cam_root / f"{data_name}"
+    cam_path_zip = cam_root / f"{data_name}.zip"
+
+    # Check if files exist, just in case.
+    if not cam_path.exists():
+        if not cam_path_zip.exists():
+            print(f"not exist {str(cam_path)} and its zip file")
+            return
+        with zipfile.ZipFile(cam_path_zip, "r") as zip_ref:
+            zip_ref.extractall(cam_path)
+
+    # Load the point cloud.
+    point_cloud = torch.load(pt_path)
+    point_cloud_np = (
+        point_cloud.cpu().numpy()
+        if hasattr(point_cloud, "numpy")
+        else np.array(point_cloud)
+    )
+    point_cloud_np = point_cloud_np[:3, :].T
+    depth_paths = glob.glob(os.path.join(cam_path, "*_depth.png"))
+    if len(depth_paths) >= NUM_VIEWS:
+        depth_paths = sorted(depth_paths)[:: len(depth_paths) // NUM_VIEWS]
+    else:
+        depth_paths = sorted(depth_paths)
+    cam_params_paths = glob.glob(os.path.join(cam_path, "*.json"))
+    if os.path.isfile(os.path.join(cam_path, "transforms_train.json")):
+        cam_params_paths.remove(os.path.join(cam_path, "transforms_train.json"))
+    if os.path.isfile(os.path.join(cam_path, "info.json")):
+        cam_params_paths.remove(os.path.join(cam_path, "info.json"))
+    if len(depth_paths) >= NUM_VIEWS:
+        cam_params_paths = sorted(cam_params_paths)[
+            :: len(cam_params_paths) // NUM_VIEWS
+        ]
+    else:
+        cam_params_paths = sorted(cam_params_paths)
+
+    valid_depth_paths = {
+        os.path.basename(p).replace("_depth.png", ""): p for p in depth_paths
+    }
+    valid_cam_params_paths = {
+        os.path.basename(p).replace(".json", ""): p for p in cam_params_paths
+    }
+    common_ids = set(valid_depth_paths.keys()) & set(valid_cam_params_paths.keys())
+    sorted_common_ids = sorted(list(common_ids))
+
+    valid_depth_paths_ = [valid_depth_paths[uid] for uid in sorted_common_ids]
+    valid_cam_params_paths_ = [valid_cam_params_paths[uid] for uid in sorted_common_ids]
+    # Calculate and save correspondences for all views.
+    correspondenceSave(
+        valid_depth_paths_,
+        data_name,
+        point_cloud_np,
+        valid_cam_params_paths_,
+        output_root,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Generate 2D-3D correspondences for GLB models and point clouds."
+    )
+    parser.add_argument(
+        "--cam_root",
+        type=Path,
+        required=True,
+        help="Root directory containing camera parameters.",
+    )
+    parser.add_argument(
+        "--point_cloud_root",
+        type=Path,
+        required=True,
+        help="Directory containing .pt point cloud files.",
+    )
+    parser.add_argument(
+        "--output_root",
+        type=Path,
+        required=True,
+        help="Output directory to save correspondence files.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=mp.cpu_count(),
+        help="Number of processes to use for processing.",
+    )
+    parser.add_argument(
+        "--thread_id", type=int, default=0, help="thread id to use for processing."
+    )
+    config = parser.parse_args()
+
+    # Ensure the output directory exists.
+    os.makedirs(config.output_root, exist_ok=True)
+
+    print("Scanning for .pt files...")
+    pt_files = list(config.point_cloud_root.glob("**/*.pt"))
+    print(f"Found {len(pt_files)} .pt files.")
+
+    # Load scene paths
+    pt_files_list = np.array_split(pt_files, config.num_workers)
+    pt_files_ = pt_files_list[config.thread_id]
+    # Preprocess data.
+    print("Processing scenes...")
+    for pt_files_i in pt_files_:
+        handle_process(
+            pt_files_i,
+            config.output_root,
+            config.cam_root,
+        )
+
+    print("\nAll files processed.")

pointcept/datasets/preprocessing/concerto/cap3d/preprocess_cap3d.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+cam_root=""
+point_cloud_root=""
+output_root=""
+num_workers=16
+
+while getopts "c:p:o:n" opt; do
+  case $opt in
+    c) cam_root=$OPTARG ;;
+    p) point_cloud_root=$OPTARG ;;
+    o) output_root=$OPTARG ;;
+    n) num_workers=$OPTARG ;;
+    *) echo "Usage: $0 -c <cam_root> -p <point_cloud_root> -o <output_root> [-n <num_workers>]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$cam_root" ] || [ -z "$point_cloud_root" ] || [ -z "$output_root" ]; then
+    echo "Error: Missing required arguments. Use -h for help."
+    echo "Usage: $0 -c <cam_root> -p <point_cloud_root> -o <output_root> [-n <num_workers>]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/cap3d/preprocess_cap3d.py \
+    --cam_root $cam_root \
+    --point_cloud_root $point_cloud_root \
+    --output_root $output_root \
+    --num_workers $num_workers \
+    --thread_id $i"
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/cap3d/splits.py

@@ -0,0 +1,76 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    co_path = os.path.join(dataset_path, "correspondences")
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        co_split_path = os.path.join(co_path, split)
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split)
+        split_names = [f.name for f in os.scandir(co_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            print(f"processing {name}")
+            co_name_split_path = os.path.join(co_split_path, name)
+            img_names = [
+                f.name.removesuffix(".npy")
+                for f in os.scandir(co_name_split_path)
+                if f.is_file() and f.name.endswith(".npy")
+            ]
+            img_names = sorted(img_names)
+            png_file_paths = [
+                os.path.join(im_split_path, name, img_name + ".png").replace(
+                    dataset_path, "data/cap3d"
+                )
+                for img_name in img_names
+            ]
+            # Get the full paths of the .png files
+            co_file_paths = [
+                os.path.join(co_name_split_path, img_name + ".npy").replace(
+                    dataset_path, "data/cap3d"
+                )
+                for img_name in img_names
+            ]
+            pc_file_path = os.path.join(pc_split_path, name + ".pt")
+            file_check = True
+            for png_file_path in png_file_paths:
+                if not os.path.isfile(png_file_path):
+                    file_check = False
+            for co_file_path in co_file_paths:
+                if not os.path.isfile(co_file_path):
+                    file_check = False
+            if not os.path.isfile(pc_file_path):
+                file_check = False
+                continue
+            if file_check:
+                split_dict[f"{name}"] = {}
+                split_dict[f"{name}"]["pointclouds"] = pc_file_path
+                split_dict[f"{name}"]["images"] = png_file_paths
+                split_dict[f"{name}"]["correspondences"] = co_file_paths
+            else:
+                split_dict[f"{name}"] = {}
+                split_dict[f"{name}"]["pointclouds"] = pc_file_path
+                split_dict[f"{name}"]["images"] = []
+                split_dict[f"{name}"]["correspondences"] = []
+    with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+        json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/graspnet/preprocess_graspnet_poses.py

@@ -0,0 +1,56 @@
+"""
+Preprocessing Script for GraspNet
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import numpy as np
+import argparse
+from tqdm import tqdm
+
+
+def split_poses(graspnet_root):
+    """
+    Splits the camera_poses.npy file into individual pose files for each scene.
+    """
+    scenes_path = os.path.join(graspnet_root, "scenes")
+    scene_names = sorted(
+        [
+            d
+            for d in os.listdir(scenes_path)
+            if os.path.isdir(os.path.join(scenes_path, d))
+        ]
+    )
+
+    print(f"Found {len(scene_names)} scenes. Starting preprocessing...")
+
+    for scene_name in tqdm(scene_names, desc="Processing Scenes"):
+        kinect_path = os.path.join(scenes_path, scene_name, "kinect")
+        poses_file = os.path.join(kinect_path, "camera_poses.npy")
+        if not os.path.exists(poses_file):
+            print(f"Warning: camera_poses.npy not found in {kinect_path}, skipping.")
+            continue
+
+        all_poses = np.load(poses_file)  # Shape: (256, 4, 4)
+        output_pose_dir = os.path.join(kinect_path, "pose")
+        os.makedirs(output_pose_dir, exist_ok=True)
+
+        for i in range(all_poses.shape[0]):
+            pose = all_poses[i]
+            output_file = os.path.join(output_pose_dir, f"{i:04d}.npy")
+            np.save(output_file, pose)
+
+    print("Preprocessing finished. Individual pose files have been created.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Preprocess GraspNet camera poses.")
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the root of the GraspNet dataset (the one containing the 'scenes' folder).",
+    )
+    args = parser.parse_args()
+    split_poses(args.dataset_root)

pointcept/datasets/preprocessing/concerto/graspnet/splits.py

@@ -0,0 +1,91 @@
+# GraspNet : train: 0-99, test_seen: 100-189
+# train: scene_0000 to scene_0100 (100 scenes)
+# val:   scene_0100 to scene_0189 (90 scenes)
+import os
+import json
+import argparse
+from tqdm import tqdm
+
+
+def generate_graspnet_splits(dataset_path, group_num=4):
+    """
+    Generates train/val/test JSON split files for the GraspNet dataset.
+    """
+    scenes_dir = os.path.join(dataset_path, "scenes")
+    all_scenes = sorted(
+        [f for f in os.scandir(scenes_dir) if f.is_dir()], key=lambda f: f.name
+    )
+    all_scene_names = [f.name for f in all_scenes]
+
+    split_definitions = {
+        "train": [s for s in all_scene_names if 100 > int(s.split("_")[1]) >= 0],
+        "val": [s for s in all_scene_names if int(s.split("_")[1]) >= 100],
+    }
+
+    output_split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(output_split_path, exist_ok=True)
+    print(f"JSON files will be saved in: {output_split_path}")
+
+    for split, scene_names in split_definitions.items():
+        print(f"Processing '{split}' split with {len(scene_names)} scenes...")
+        split_dict = {}
+        for scene_name in tqdm(scene_names, desc=f"Generating {split}.json"):
+            kinect_path = os.path.join(scenes_dir, scene_name, "kinect")
+
+            rgb_path = os.path.join(kinect_path, "rgb")
+            depth_path = os.path.join(kinect_path, "depth")
+            pose_path = os.path.join(kinect_path, "pose")
+            K_file_path_abs = os.path.join(kinect_path, "camK.npy")
+            K_file_path_rel = K_file_path_abs.replace(dataset_path, "data/graspnet")
+
+            rgb_files = [f for f in os.listdir(rgb_path) if f.endswith(".png")]
+            rgb_files = sorted(rgb_files, key=lambda x: int(os.path.splitext(x)[0]))
+
+            rgb_file_paths = [
+                os.path.join(rgb_path, f).replace(dataset_path, "data/graspnet")
+                for f in rgb_files
+            ]
+            depth_file_paths = [
+                os.path.join(depth_path, f).replace(dataset_path, "data/graspnet")
+                for f in rgb_files
+            ]
+            pose_file_paths = [
+                os.path.join(pose_path, f.replace(".png", ".npy")).replace(
+                    dataset_path, "data/graspnet"
+                )
+                for f in rgb_files
+            ]
+
+            for i in range(0, len(rgb_file_paths), group_num):
+                if i + 4 > len(rgb_file_paths):
+                    continue
+
+                chunk_key = f"{scene_name}_{i//group_num}"
+                split_dict[chunk_key] = {
+                    "images": rgb_file_paths[i : i + group_num],
+                    "depths": depth_file_paths[i : i + group_num],
+                    "Ts": pose_file_paths[i : i + group_num],
+                    "Ks": [K_file_path_rel] * group_num,
+                }
+
+        output_json_file = os.path.join(output_split_path, f"{split}.json")
+        with open(output_json_file, "w") as f:
+            json.dump(split_dict, f, indent=4)
+        print(f"Successfully created {output_json_file}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the GraspNet dataset root (containing the 'scenes' folder).",
+    )
+    parser.add_argument(
+        "--group_num",
+        default=4,
+        type=int,
+        help="Group Num.",
+    )
+    config = parser.parse_args()
+    generate_graspnet_splits(config.dataset_root, config.group_num)

pointcept/datasets/preprocessing/concerto/hk/hkremote_download.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# Export variables so they can be accessed by subshells spawned by xargs
+export BASE_URL="https://data11.map.gov.hk/api/3d-zip/OBJ/"
+# To get API_KEY, visit https://3d.map.gov.hk/download-api. After you agree the license, your API_KEY will show in Example URL: https://download.map.gov.hk/api/3d-zip/[format]/[tile_sheet_num].zip?key=[API_KEY]
+export API_KEY=${API_KEY}
+export OUTPUT_DIR="hk_3d_maps"
+export MAX_THREADS=8  # Set the number of concurrent downloads here
+
+mkdir -p "$OUTPUT_DIR"
+echo "Files will be downloaded to: $OUTPUT_DIR"
+echo "Using $MAX_THREADS concurrent threads..."
+
+# Define the download logic as a function
+download_map() {
+  local MAP_INDEX=$1
+  local ZIP_FILENAME="${MAP_INDEX}.zip"
+  local OUTPUT_PATH="${OUTPUT_DIR}/${ZIP_FILENAME}"
+  local DOWNLOAD_URL="${BASE_URL}${ZIP_FILENAME}?key=${API_KEY}"
+
+  # Check existence
+  if [ -f "$OUTPUT_PATH" ]; then
+    echo "[SKIP] File already exists: ${ZIP_FILENAME}"
+    return 0
+  fi
+
+  echo "[START] Downloading: ${MAP_INDEX}"
+  
+  # Use wget to download
+  # Added --timeout=10 and --tries=2 to prevent threads from hanging forever
+  wget --no-check-certificate -q --timeout=10 --tries=2 -O "$OUTPUT_PATH" "$DOWNLOAD_URL"
+
+  # Check wget exit code; $? == 0 means success
+  if [ $? -eq 0 ]; then
+    echo "[SUCCESS] Downloaded: ${ZIP_FILENAME}"
+  else
+    echo "[FAILED] Not found or failed: ${ZIP_FILENAME}"
+    # Delete the potentially empty file created by a failed download
+    rm -f "$OUTPUT_PATH"
+  fi
+}
+
+# Export the function so xargs subshells can use it
+export -f download_map
+
+# Generate the list of map indices and pipe them to xargs for parallel execution
+# We use standard space-separated lists for the string variables for cleaner execution
+(
+  for base in {1..15}; do
+    for quad_10k in NW NE SW SE; do
+      for grid_2k in {1..25}; do
+        for sub_grid_1k in A B C D; do
+          echo "${base}-${quad_10k}-${grid_2k}${sub_grid_1k}"
+        done
+      done
+    done
+  done
+) | xargs -P "$MAX_THREADS" -I {} bash -c 'download_map "$@"' _ {}
+
+echo "All download tasks completed!"

pointcept/datasets/preprocessing/concerto/hk/preprocess_hk.py

@@ -0,0 +1,110 @@
+"""
+Preprocessing Script for HK Remote
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import open3d as o3d
+import numpy as np
+from pathlib import Path
+import argparse
+import multiprocessing as mp
+
+
+def handle_process(obj_file_paths, output_directory):
+    num_points_to_sample = 10000000
+    for obj_file_path in obj_file_paths:
+        try:
+            print(f"Loading mesh from: {obj_file_path}")
+            parent_name = obj_file_path.parent.name
+            save_path = output_directory / parent_name
+            if save_path.exists():
+                print(f"exist {str(parent_name)}")
+                continue
+            mesh = o3d.io.read_triangle_mesh(obj_file_path, enable_post_processing=True)
+
+            if not mesh.has_vertices():
+                print("Error: Mesh could not be loaded.")
+                exit()
+
+            if not mesh.has_textures():
+                print(
+                    "Warning: Mesh does not have textures. Sampled colors will be black or default."
+                )
+
+            mesh.compute_vertex_normals()
+            print("Mesh loaded successfully. Proceeding to sample points...")
+
+            pcd = mesh.sample_points_uniformly(number_of_points=num_points_to_sample)
+            points = np.asarray(pcd.points)
+            colors = np.asarray(pcd.colors)
+            if colors.shape[0] == 0:
+                colors = np.zeros_like(points)
+            normals = np.asarray(pcd.normals)
+            if normals.shape[0] == 0:
+                normals = np.zeros_like(points)
+
+            print(f"\nSuccessfully sampled {len(points)} points.")
+            print(f"Points array shape: {points.shape}")  # (N, 3) for XYZ coordinates
+            print(
+                f"Colors array shape: {colors.shape}"
+            )  # (N, 3) for RGB values (0-1 range)
+            print(f"Normals array shape: {normals.shape}")  # (N, 3) for normal vectors
+
+            output_coord_path = output_directory / parent_name / "coord.npy"
+            output_color_path = output_directory / parent_name / "color.npy"
+            output_normal_path = output_directory / parent_name / "normal.npy"
+            (output_directory / parent_name).mkdir(exist_ok=True)
+
+            np.save(output_coord_path, points)
+            np.save(output_color_path, (colors * 255).astype(np.int32))
+            np.save(output_normal_path, normals)
+
+            pcd = o3d.geometry.PointCloud()
+            pcd.points = o3d.utility.Vector3dVector(points)
+            pcd.colors = o3d.utility.Vector3dVector(colors)
+            pcd.normals = o3d.utility.Vector3dVector(normals)
+            print(
+                f"\nSampled point cloud saved to: {str(output_directory/parent_name)}"
+            )
+        except:
+            print(f"fail {parent_name}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="Thread id for parallel processing",
+    )
+    config = parser.parse_args()
+    output_directory = Path(config.output_root) / "train"
+    root_directory = Path(config.dataset_root)
+    output_directory.mkdir(exist_ok=True)
+    obj_file_paths_generator = root_directory.rglob("*.obj")
+    obj_file_paths = list(obj_file_paths_generator)
+    print(f"Found {len(obj_file_paths)} .obj files.")
+
+    obj_file_paths_list = np.array_split(obj_file_paths, config.num_workers)
+    obj_file_paths_ = obj_file_paths_list[config.thread_id]
+
+    handle_process(obj_file_paths_, output_directory)

pointcept/datasets/preprocessing/concerto/hk/preprocess_hk.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=16  
+parse_depths=false
+
+while getopts "d:o:n:p" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/hk/preprocess_hk.py --thread_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root"
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/hk/splits.py

@@ -0,0 +1,35 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    pc_path = dataset_path
+    splits = ["train"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        pc_split_path = os.path.join(pc_path, split)
+        split_names = [f.name for f in os.scandir(pc_split_path) if f.is_dir()]
+        pc_split_path = pc_split_path.replace(dataset_path, "data/hk")
+        split_dict = {}
+        for name in split_names:
+            split_dict[f"{name}"] = {}
+            split_dict[f"{name}"]["pointclouds"] = os.path.join(pc_split_path, name)
+            split_dict[f"{name}"]["images"] = []
+            split_dict[f"{name}"]["correspondences"] = []
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/hm3d/common/utils.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import random
+from typing import List, Optional
+
+import habitat_sim
+import numpy as np
+import quaternion as qt
+from habitat_sim.utils.common import quat_from_angle_axis
+
+
+def get_random_quaternion() -> qt.quaternion:
+    return quat_from_angle_axis(
+        math.radians(random.uniform(-180.0, 180.0)), np.array([0, 1.0, 0])
+    )
+
+
+def convert_heading_to_quaternion(heading: float) -> qt.quaternion:
+    # heading angle in degrees
+    return quat_from_angle_axis(math.radians(heading), np.array([0, 1.0, 0]))
+
+
+def make_habitat_configuration(scene_path: str, use_sensor: bool = False):
+    # simulator configuration
+    backend_cfg = habitat_sim.SimulatorConfiguration()
+    backend_cfg.scene_id = scene_path
+
+    # agent configuration
+    sensor_cfg = habitat_sim.CameraSensorSpec()
+    sensor_cfg.resolution = [1080, 960]
+    sensor_cfg.sensor_type = habitat_sim.SensorType.DEPTH
+    agent_cfg = habitat_sim.agent.AgentConfiguration()
+    agent_cfg.sensor_specifications = [sensor_cfg] if use_sensor else []
+
+    return habitat_sim.Configuration(backend_cfg, [agent_cfg])
+
+
+def robust_load_sim(scene_path: str) -> habitat_sim.Simulator:
+    sim_cfg = make_habitat_configuration(scene_path, use_sensor=False)
+    hsim = habitat_sim.Simulator(sim_cfg)
+    if not hsim.pathfinder.is_loaded:
+        hsim.close()
+        sim_cfg = make_habitat_configuration(scene_path, use_sensor=True)
+        hsim = habitat_sim.Simulator(sim_cfg)
+        navmesh_settings = habitat_sim.NavMeshSettings()
+        navmesh_settings.set_defaults()
+        hsim.recompute_navmesh(hsim.pathfinder, navmesh_settings)
+    return hsim
+
+
+def get_filtered_scenes(scenes: List[str], filter_scenes_path: str) -> List[str]:
+    """
+    Filter scenes to only include valid scenes.
+    """
+    with open(filter_scenes_path, "r") as fp:
+        filter_scenes = fp.readlines()
+    filter_scenes = [f.strip("\n") for f in filter_scenes]
+    filtered_scenes = []
+    for scene in scenes:
+        scene_name = scene.split("/")[-1][: -len(".glb")]
+        if scene_name in filter_scenes:
+            filtered_scenes.append(scene)
+    return filtered_scenes
+
+
+def quaternion_to_list(q: qt.quaternion):
+    return q.imag.tolist() + [q.real]
+
+
+# -------------------------------------------------------------------------------
+# Functionality from habitat-lab
+# -------------------------------------------------------------------------------
+def calculate_meters_per_pixel(map_resolution: int, pathfinder=None):
+    r"""Calculate the meters_per_pixel for a given map resolution"""
+    lower_bound, upper_bound = pathfinder.get_bounds()
+    return min(
+        abs(upper_bound[coord] - lower_bound[coord]) / map_resolution
+        for coord in [0, 2]
+    )
+
+
+def get_topdown_map(
+    pathfinder,
+    height: float,
+    map_resolution: int = 1024,
+    meters_per_pixel: Optional[float] = None,
+) -> np.ndarray:
+    r"""Return a top-down occupancy map for a sim. Note, this only returns valid
+    values for whatever floor the agent is currently on.
+    :param pathfinder: A habitat-sim pathfinder instances to get the map from
+    :param height: The height in the environment to make the topdown map
+    :param map_resolution: Length of the longest side of the map.  Used to calculate :p:`meters_per_pixel`
+    :param draw_border: Whether or not to draw a border
+    :param meters_per_pixel: Overrides map_resolution an
+    :return: Image containing 0 if occupied, 1 if unoccupied, and 2 if border (if
+        the flag is set).
+    """
+
+    if meters_per_pixel is None:
+        meters_per_pixel = calculate_meters_per_pixel(
+            map_resolution, pathfinder=pathfinder
+        )
+
+    top_down_map = pathfinder.get_topdown_view(
+        meters_per_pixel=meters_per_pixel, height=height
+    ).astype(np.uint8)
+
+    return np.ascontiguousarray(top_down_map)

pointcept/datasets/preprocessing/concerto/hm3d/preprocessed_hm3d.py

@@ -0,0 +1,617 @@
+"""
+Preprocessing Script for hm3d using habitat-sim
+https://github.com/facebookresearch/habitat-sim/blob/main/examples/tutorials/notebooks/ECCV_2020_Navigation.ipynb
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import glob
+import itertools
+import multiprocessing as mp
+import os
+from pathlib import Path
+import random
+from typing import Any, Dict, List, Optional, Tuple
+import habitat_sim
+import imageio
+import numpy as np
+from sklearn.cluster import DBSCAN
+from scipy.spatial import cKDTree
+import trimesh
+import sys
+
+sys.path.append("../../hm3d/")
+from hm3d_constants import CLASS_LABELS_40
+
+
+from common.utils import (
+    convert_heading_to_quaternion,
+    get_topdown_map,
+)
+
+MPCAT2INDEX = dict([(CLASS_LABELS_40[i], i) for i in range(40)])
+MPCAT2INDEX["unlabeled"] = -1
+CAT2INDEX = dict()
+with open(Path("../../hm3d") / "hm3dsem_category_mappings.tsv") as f:
+    f.readline()  # raw_category	category	mpcat40
+    lines = f.readlines()
+for line in lines:
+    line = line.strip().split("\t")
+    if len(line) == 2:
+        # L2164: "\tunknown\tunlabeled"
+        line.insert(0, "")
+    CAT2INDEX[line[1]] = MPCAT2INDEX[line[2]]
+
+
+def correspondenceGet(depth, K, T, img_size):
+    height, width = img_size
+    if np.isnan(T).any() or np.isinf(T).any():
+        return None
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    valid = ~np.isinf(depth).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict
+
+
+def uv_to_texture_color(uv, texture):
+    width, height = texture.size
+    u = int(uv[0] * width)
+    v = int((1 - uv[1]) * height)  # Flip y-axis for image coordinates
+    if 0 <= u < width and 0 <= v < height:
+        return texture.getpixel((u, v))
+    else:
+        return 0, 0, 0
+
+
+def load_hex_mapping(mapping_path):
+    hex2label = {}
+    with open(mapping_path) as f:
+        f.readline()  # remove 'HM3D Semantic Annotations\n'
+        lines = f.readlines()  # get the left
+    for line in lines:
+        line = line.strip().split(",")
+        cat = line[2].strip('"')
+        if cat == "trashcan":
+            cat = "trash can"
+        elif cat == "fridge":
+            cat = "refrigerator"
+        hex2label[line[1]] = dict(instance=line[0], segment=CAT2INDEX[cat])
+    return hex2label
+
+
+def make_habitat_configuration(
+    scene_path: str,
+    hfov: int = 90,
+    resolution: Tuple[int] = (300, 300),
+    stage_json_path: Optional[str] = None,
+):
+    # simulator configuration
+    backend_cfg = habitat_sim.SimulatorConfiguration()
+    if stage_json_path is not None:
+        backend_cfg.scene_dataset_config_file = stage_json_path
+        backend_cfg.scene_id = "habitat/" + scene_path.split("/")[-1]
+    else:
+        backend_cfg.scene_id = scene_path
+
+    # agent configuration
+    rgb_sensor_cfg = habitat_sim.CameraSensorSpec()
+    rgb_sensor_cfg.uuid = "rgba"
+    rgb_sensor_cfg.sensor_type = habitat_sim.SensorType.COLOR
+    rgb_sensor_cfg.resolution = resolution
+    rgb_sensor_cfg.hfov = hfov
+    rgb_sensor_cfg.position = [0.0, 0.0, 0.0]
+    depth_sensor_cfg = habitat_sim.CameraSensorSpec()
+    depth_sensor_cfg.uuid = "depth"
+    depth_sensor_cfg.sensor_type = habitat_sim.SensorType.DEPTH
+    depth_sensor_cfg.resolution = resolution
+    depth_sensor_cfg.hfov = hfov
+    depth_sensor_cfg.position = [0.0, 0.0, 0.0]
+
+    agent_cfg = habitat_sim.agent.AgentConfiguration()
+    agent_cfg.sensor_specifications = [rgb_sensor_cfg, depth_sensor_cfg]
+
+    return habitat_sim.Configuration(backend_cfg, [agent_cfg])
+
+
+def robust_load_sim(scene_path: str, **kwargs: Any) -> habitat_sim.Simulator:
+    sim_cfg = make_habitat_configuration(scene_path, **kwargs)
+    hsim = habitat_sim.Simulator(sim_cfg)
+    if not hsim.pathfinder.is_loaded:
+        navmesh_settings = habitat_sim.NavMeshSettings()
+        navmesh_settings.set_defaults()
+        hsim.recompute_navmesh(hsim.pathfinder, navmesh_settings)
+    return hsim
+
+
+def get_floor_heights(
+    sim: habitat_sim.Simulator, max_points_to_sample: int = 20000
+) -> List[Dict[str, float]]:
+    """Get heights of different floors in a scene. This is done in two steps.
+    (1) Randomly samples navigable points in the scene.
+    (2) Cluster the points based on discretized y coordinates to get floors.
+
+    Args:
+        sim: habitat simulator instance
+        max_points_to_sample: number of navigable points to randomly sample
+    """
+    nav_points = []
+    for _ in range(max_points_to_sample):
+        nav_points.append(sim.pathfinder.get_random_navigable_point())
+    nav_points = np.stack(nav_points, axis=0)
+    y_coors = np.around(nav_points[:, 1], decimals=1)
+    # cluster Y coordinates
+    clustering = DBSCAN(eps=0.2, min_samples=2000).fit(y_coors[:, np.newaxis])
+    c_labels = clustering.labels_
+    n_clusters = len(set(c_labels)) - (1 if -1 in c_labels else 0)
+    # get floor extents in Y
+    # each cluster corresponds to points from 1 floor
+    floor_extents = []
+    core_sample_y = y_coors[clustering.core_sample_indices_]
+    core_sample_labels = c_labels[clustering.core_sample_indices_]
+    for i in range(n_clusters):
+        floor_min = core_sample_y[core_sample_labels == i].min().item()
+        floor_max = core_sample_y[core_sample_labels == i].max().item()
+        floor_mean = core_sample_y[core_sample_labels == i].mean().item()
+        floor_extents.append({"min": floor_min, "max": floor_max, "mean": floor_mean})
+    floor_extents = sorted(floor_extents, key=lambda x: x["mean"])
+
+    # reject floors that have too few points
+    max_points = 0
+    for fext in floor_extents:
+        top_down_map = get_topdown_map(sim.pathfinder, fext["mean"])
+        max_points = max(np.count_nonzero(top_down_map), max_points)
+    clean_floor_extents = []
+    for fext in floor_extents:
+        top_down_map = get_topdown_map(sim.pathfinder, fext["mean"])
+        num_points = np.count_nonzero(top_down_map)
+        if num_points < 0.2 * max_points:
+            continue
+        clean_floor_extents.append(fext)
+
+    return clean_floor_extents
+
+
+def get_navmesh_extents_at_y(
+    sim: habitat_sim.Simulator, y_bounds: Optional[Tuple[float]] = None
+) -> Tuple[float]:
+    if y_bounds is None:
+        lower_bound, upper_bound = sim.pathfinder.get_bounds()
+    else:
+        assert len(y_bounds) == 2
+        assert y_bounds[0] < y_bounds[1]
+        navmesh_vertices = np.array(sim.pathfinder.build_navmesh_vertices())
+        navmesh_vertices = navmesh_vertices[
+            (y_bounds[0] <= navmesh_vertices[:, 1])
+            & (navmesh_vertices[:, 1] <= y_bounds[1])
+        ]
+        lower_bound = navmesh_vertices.min(axis=0)
+        upper_bound = navmesh_vertices.max(axis=0)
+    return (lower_bound, upper_bound)
+
+
+def get_dense_navmesh_vertices(
+    sim: habitat_sim.Simulator, sampling_resolution: float = 0.5
+) -> np.ndarray:
+
+    navmesh_vertices = []
+    floor_extents = get_floor_heights(sim)
+    for fext in floor_extents:
+        l_bound, u_bound = get_navmesh_extents_at_y(
+            sim, y_bounds=(fext["min"] - 0.5, fext["max"] + 0.5)
+        )
+        x_range = np.arange(l_bound[0].item(), u_bound[0].item(), sampling_resolution)
+        y = fext["mean"]
+        z_range = np.arange(l_bound[2].item(), u_bound[2].item(), sampling_resolution)
+        for x, z in itertools.product(x_range, z_range):
+            if sim.pathfinder.is_navigable(np.array([x, y, z])):
+                navmesh_vertices.append((np.array([x, y, z])))
+    if len(navmesh_vertices) > 0:
+        navmesh_vertices = np.stack(navmesh_vertices, axis=0)
+    else:
+        navmesh_vertices = np.zeros((0, 3))
+    return navmesh_vertices
+
+
+def get_dense_mesh_vertices(
+    sim: habitat_sim.Simulator, room_bounds: dict, sampling_resolution: float = 0.5
+) -> np.ndarray:
+    navmesh_vertices = []
+    for room_id, room_bound in room_bounds.items():
+        x_range = np.arange(
+            room_bound["min_x"] + 0.5, room_bound["max_x"] - 0.5, sampling_resolution
+        )
+        y_range = np.arange(
+            room_bound["min_y"] + 0.5, room_bound["max_y"] - 0.5, sampling_resolution
+        )
+        z_range = np.arange(
+            room_bound["min_z"] + 0.5, room_bound["max_z"] - 0.5, sampling_resolution
+        )
+        for x, y, z in itertools.product(x_range, y_range, z_range):
+            navmesh_vertices.append((np.array([x, y, z])))
+    if len(navmesh_vertices) > 0:
+        navmesh_vertices = np.stack(navmesh_vertices, axis=0)
+    else:
+        navmesh_vertices = np.zeros((0, 3))
+    return navmesh_vertices
+
+
+def get_scene_name(scene_path, dataset):
+    if dataset == "replica":
+        scene_name = scene_path.split("/")[-2].split(".")[0]
+    else:
+        scene_name = scene_path.split("/")[-1].split(".")[0]
+    return scene_name
+
+
+def plySave(coords, colors, save_path):
+    ply_path = save_path / "point_cloud.ply"
+    trimesh.points.PointCloud(coords, colors=colors).export(ply_path)
+
+
+def _aux_fn(args):
+    (
+        density,
+        scene_path,
+        save_prefix,
+        hfov,
+        resolution,
+        sampling_resolution,
+        num_rotations,
+        sim_kwargs,
+        parse_pointclouds,
+        parse_depths,
+    ) = args
+    if sim_kwargs is None:
+        sim_kwargs = {}
+    sim = robust_load_sim(scene_path, hfov=hfov, resolution=resolution, **sim_kwargs)
+    save_prefix = Path(save_prefix)
+    rgb_save_prefix = save_prefix / "images"
+    pc_save_prefix = save_prefix
+    intrinsic_save_path = Path(rgb_save_prefix) / "intrinsic"
+    intrinsic_save_path.mkdir(parents=True, exist_ok=True)
+    np.save(intrinsic_save_path / "K.npy", K)
+    scene_path = Path(scene_path)
+    scene_label_path = scene_path.with_suffix(".semantic.glb")
+    scene_mapping_path = scene_path.with_suffix(".semantic.txt")
+    scene_name = scene_path.parent.name.replace("-", "_")
+    scene_id = scene_name.split("_")[0]
+    labeled = True if scene_label_path.is_file() else False
+
+    # test split (900-1000) is reserved by official
+    if 0 <= int(scene_id) < 800:
+        split = "train"
+    else:
+        split = "val"
+    print(f"Parsing scene {scene_name} in {split} split...")
+    scene = trimesh.load(scene_path)
+    if labeled:
+        labeled_scene_ = trimesh.load(scene_label_path)
+        labeled_scene = trimesh.Scene()
+        for name, mesh in labeled_scene_.geometry.items():
+            # some case, name in labeled scene and scene is not matched, so only use chunk id
+            labeled_scene.add_geometry(mesh, geom_name=name.split("_")[0])
+        del labeled_scene_
+        hex2label = load_hex_mapping(scene_mapping_path)
+
+    room_dict = {}
+    room_bound = {}
+    for name, mesh in scene.geometry.items():
+        room_id = "_".join(name.split("_")[1:3]).replace("group", "").replace("sub", "")
+        if room_id not in room_dict.keys():
+            room_dict[room_id] = trimesh.Scene()
+            room_bound[room_id] = {
+                "min_x": np.inf,
+                "max_x": -np.inf,
+                "min_y": np.inf,
+                "max_y": -np.inf,
+                "min_z": np.inf,
+                "max_z": -np.inf,
+            }
+        room_dict[room_id].add_geometry(mesh, geom_name=name.split("_")[0])
+        bounding_box = mesh.bounds
+        room_bound[room_id]["min_x"] = min(
+            room_bound[room_id]["min_x"], bounding_box[0][0]
+        )
+        room_bound[room_id]["max_x"] = max(
+            room_bound[room_id]["max_x"], bounding_box[1][0]
+        )
+        room_bound[room_id]["min_y"] = min(
+            room_bound[room_id]["min_y"], bounding_box[0][1]
+        )
+        room_bound[room_id]["max_y"] = max(
+            room_bound[room_id]["max_y"], bounding_box[1][1]
+        )
+        room_bound[room_id]["min_z"] = min(
+            room_bound[room_id]["min_z"], bounding_box[0][2]
+        )
+        room_bound[room_id]["max_z"] = max(
+            room_bound[room_id]["max_z"], bounding_box[1][2]
+        )
+
+    # Get dense navmesh vertices
+    navmesh_vertices = get_dense_mesh_vertices(
+        sim, room_bound, sampling_resolution=sampling_resolution
+    )
+    room_vertices = dict()
+    for navmesh_vertex in navmesh_vertices:
+        for room_id, bound in room_bound.items():
+            if (
+                bound["min_x"] < navmesh_vertex[0] < bound["max_x"]
+                and bound["min_z"] < navmesh_vertex[2] < bound["max_z"]
+                and bound["min_y"] < navmesh_vertex[1] < bound["max_y"]
+            ):
+                if room_id not in room_vertices.keys():
+                    room_vertices[room_id] = []
+                room_vertices[room_id].append(navmesh_vertex)
+                break
+
+    for room_id, scene in room_dict.items():
+        if room_id not in room_vertices.keys():
+            continue
+        np.random.seed(int(scene_id + room_id.replace("_", "")))
+        room_coord = []
+        room_color = []
+        room_normal = []
+        if labeled:
+            room_label_color = []
+
+        for name in scene.geometry.keys():
+            mesh = scene.geometry[name]
+            num_points = int(np.sum(mesh.area_faces) / density**2)
+            if num_points == 0:
+                continue
+            coords, face_indices = mesh.sample(num_points, return_index=True)
+            faces = mesh.faces[face_indices]
+            triangles = mesh.vertices[faces]
+            bary_coords = trimesh.triangles.points_to_barycentric(triangles, coords)
+            uv_coords = mesh.visual.uv[faces]
+            sampled_uvs = np.einsum("ijk,ij->ik", uv_coords, bary_coords)
+
+            pbr_material = mesh.visual.material
+            texture_image = pbr_material.baseColorTexture
+            if texture_image is None:
+                continue
+            colors = np.array(
+                [uv_to_texture_color(uv, texture_image) for uv in sampled_uvs]
+            )
+            normals = mesh.vertex_normals[faces]
+            normals = np.einsum("ijk,ij->ik", normals, bary_coords)
+            room_coord.append(coords)
+            room_color.append(colors)
+            room_normal.append(normals)
+
+            if labeled:
+                labeled_mash = labeled_scene.geometry[name]
+                label_texture_image = labeled_mash.visual.material.baseColorTexture
+                label_color = np.array(
+                    [uv_to_texture_color(uv, label_texture_image) for uv in sampled_uvs]
+                )
+                room_label_color.append(label_color)
+        if len(room_coord) == 0:
+            continue
+
+        room_coord = np.concatenate(room_coord, axis=0).astype("float32")
+        if parse_pointclouds:
+            room_color = np.concatenate(room_color, axis=0).astype("uint8")
+            room_normal = np.concatenate(room_normal, axis=0).astype("float32")
+            data_dict = dict(coord=room_coord, color=room_color, normal=room_normal)
+
+            if labeled:
+                room_label_color = np.concatenate(room_label_color, axis=0)
+                instance_label_color = np.unique(room_label_color, axis=0)
+                room_instance = -np.ones(len(room_label_color), dtype="int16")
+                room_segment = -np.ones(len(room_label_color), dtype="int16")
+                for i in range(len(instance_label_color)):
+                    label_color = instance_label_color[i]
+                    label_hex = "{c[0]:02x}{c[1]:02x}{c[2]:02x}".format(c=label_color)
+                    mask = np.all(room_label_color == label_color, axis=-1)
+                    room_instance[mask] = i
+                    if label_hex.upper() in hex2label.keys():
+                        room_segment[mask] = hex2label[label_hex.upper()]["segment"]
+                data_dict["instance"] = room_instance
+                data_dict["segment"] = room_segment
+            pc_save_path_single = (
+                Path(pc_save_prefix) / split / "_".join([scene_name, room_id])
+            )
+            os.makedirs(pc_save_path_single, exist_ok=True)
+            print(f"Saving data to {pc_save_path_single}...")
+            for key, value in data_dict.items():
+                np.save(pc_save_path_single / f"{key}.npy", value)
+
+        agent = sim.get_agent(0)
+        rgb_paths = []
+        depth_paths = []
+        if not sim.pathfinder.is_loaded:
+            sim.close()
+            return rgb_paths, depth_paths
+
+        vertex = sum(room_vertices[room_id]) / len(room_vertices[room_id])
+
+        T0 = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]])
+        Ts = []
+        loc_ = vertex + np.random.uniform(-0.25, 0.25, size=vertex.shape)
+        count = 0
+        loc = loc_ @ T0[:3, :3].T
+        init_angle = random.uniform(0, 90)
+        pose_save_path = (
+            Path(rgb_save_prefix) / split / "_".join([scene_name, room_id]) / "pose"
+        )
+        pose_save_path.mkdir(parents=True, exist_ok=True)
+        rgb_path = (
+            Path(rgb_save_prefix) / split / "_".join([scene_name, room_id]) / "color"
+        )
+        os.makedirs(rgb_path, exist_ok=True)
+        if parse_depths:
+            depth_path = (
+                Path(rgb_save_prefix)
+                / split
+                / "_".join([scene_name, room_id])
+                / "depth"
+            )
+            os.makedirs(depth_path, exist_ok=True)
+        for heading in np.linspace(init_angle, init_angle + 360, num_rotations + 1)[
+            :-1
+        ]:
+            rot = convert_heading_to_quaternion(heading)
+            angle = np.radians(-heading)  # Replace 45 with the desired angle in degrees
+            R_y = np.array(
+                [
+                    [np.cos(angle), 0, np.sin(angle)],
+                    [0, 1, 0],
+                    [-np.sin(angle), 0, np.cos(angle)],
+                ]
+            )
+            T = np.eye(4)
+            T[:3, :3] = T0[:3, :3] @ R_y
+            T[:3, 3] = loc_
+            Ts.append(T)
+            np.save(pose_save_path / f"{count}.npy", T)
+            agent_state = agent.get_state()
+            agent_state.position = loc
+            agent_state.rotation = rot
+            agent.set_state(agent_state, reset_sensors=True)
+            obs = sim.get_sensor_observations()
+            rgb = obs["rgba"][..., :3]
+            depth = obs["depth"]
+            rgb_path_ = rgb_path / f"{count}.png"
+            imageio.imwrite(rgb_path_, rgb)
+            if parse_depths:
+                depth_path_ = depth_path / f"{count}.png"
+                depth_img = (depth * 1000).astype(np.uint16)  # Convert to mm
+                imageio.imwrite(depth_path_, depth_img)
+            coord_dict = correspondenceGet(depth, K, T, RESOLUTION)
+            if coord_dict is None:
+                correspondences = np.ones((1, 3))
+            else:
+                pixels_ = coord_dict["pixel"]
+                coords_ = coord_dict["coord"]
+                tree = cKDTree(room_coord)
+                dis, idx = tree.query(coords_, k=1)
+                idx_valid = idx[dis < 0.01]
+                pixels_valid = pixels_[dis < 0.01]
+                co_save_path = (
+                    Path(rgb_save_prefix)
+                    / split
+                    / "_".join([scene_name, room_id])
+                    / "correspondence"
+                )
+                co_save_path.mkdir(parents=True, exist_ok=True)
+                correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+            np.save(co_save_path / "{}.npy".format(count), correspondences)
+            count += 1
+    sim.close()
+
+
+HFOV = 90
+HFOV_rad = np.deg2rad(HFOV)
+RESOLUTION = [720, 720]
+NUM_ROTATIONS = 4
+F = RESOLUTION[1] / (2 * np.tan(HFOV_rad / 2))
+CX, CY = RESOLUTION[1] / 2, RESOLUTION[0] / 2
+K = np.array([[F, 0, CX], [0, F, CY], [0, 0, 1]])
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--parse_pointclouds",
+        action="store_true",
+        help="Parse point clouds from the scenes",
+    )
+    parser.add_argument(
+        "--parse_depths",
+        action="store_true",
+        help="Parse depths from the scenes",
+    )
+    parser.add_argument(
+        "--density",
+        default=0.02,
+        type=float,
+        help="Sampling density on mesh surface (m)",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--worker_id",
+        default=0,
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    args = parser.parse_args()
+
+    scene_list = glob.glob(os.path.join(args.dataset_root, "*", "*.glb"))
+    scene_list = [scene for scene in scene_list if not scene.endswith("semantic.glb")]
+    scene_list = [scene for scene in scene_list if not scene.endswith("basis.glb")]
+    scene_list = sorted(scene_list, key=lambda x: int(x.split("/")[-2].split("-")[0]))
+    assert len(scene_list) == 900
+    print(
+        f"Found {len(scene_list)} scenes in {args.dataset_root},using {args.num_workers} workers."
+    )
+    scene_list = [scene_list[i :: args.num_workers] for i in range(args.num_workers)][
+        args.worker_id
+    ]
+
+    pc_train_dir = os.path.join(args.output_root, "pointclouds", "train")
+    pc_val_dir = os.path.join(args.output_root, "pointclouds", "val")
+    existing_scene_ids = set()
+
+    for pc_dir in [pc_train_dir, pc_val_dir]:
+        if os.path.exists(pc_dir):
+            for folder in os.listdir(pc_dir):
+                if "_" in folder:
+                    scene_id = folder.split("_")[0]
+                    existing_scene_ids.add(scene_id)
+
+    scene_list = [
+        scene
+        for scene in scene_list
+        if scene.split("/")[-2].split("-")[0] not in existing_scene_ids
+    ]
+    print(f"Processing {len(scene_list)} scenes...")
+
+    for scene_path in scene_list:
+        inputs = (
+            args.density,
+            scene_path,
+            args.output_root,
+            HFOV,
+            RESOLUTION,
+            0.5,
+            NUM_ROTATIONS,
+            None,
+            args.parse_pointclouds,
+            args.parse_depths,
+        )
+        _aux_fn(inputs)

pointcept/datasets/preprocessing/concerto/hm3d/preprocessed_hm3d.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=8 
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    p) parse_depths=true ;;
+    c) parse_pointclouds=true ;;
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/hm3d/preprocess_hm3d.py --worker_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root"
+
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/hm3d/splits.py

@@ -0,0 +1,60 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train", "val"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(dataset_path, "data/hm3d")
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            im_split_name_path = os.path.join(im_split_path, name, "color")
+            co_split_name_path = os.path.join(im_split_path, name, "correspondence")
+            png_files = [
+                f
+                for f in os.listdir(im_split_name_path)
+                if f.endswith(".png") and (not f.startswith("."))
+            ]
+            png_files = sorted(png_files, key=lambda x: int(x.split(".")[0]))
+            # Get the full paths of the .png files
+            png_file_paths = [
+                os.path.join(im_split_name_path, f).replace(dataset_path, "data/hm3d")
+                for f in png_files
+            ]
+            co_file_paths = [
+                os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace(
+                    dataset_path, "data/hm3d"
+                )
+                for f in png_files
+            ]
+            for i in range(0, len(png_file_paths), 4):
+                split_dict[f"{name}_{i//4}"] = {}
+                split_dict[f"{name}_{i//4}"]["pointclouds"] = os.path.join(
+                    pc_split_path, name
+                )
+                split_dict[f"{name}_{i//4}"]["images"] = png_file_paths[i : i + 4]
+                split_dict[f"{name}_{i//4}"]["correspondences"] = co_file_paths[
+                    i : i + 4
+                ]
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/nuscenes/preprocess_nuscenes_info.py

@@ -0,0 +1,687 @@
+"""
+Preprocessing Script for nuScenes Informantion
+modified from OpenPCDet (https://github.com/open-mmlab/OpenPCDet)
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+from pathlib import Path
+import numpy as np
+import argparse
+import tqdm
+import pickle
+from functools import reduce
+from pyquaternion import Quaternion
+from nuscenes.nuscenes import NuScenes
+from nuscenes.utils import splits
+from nuscenes.utils.geometry_utils import transform_matrix
+
+map_name_from_general_to_detection = {
+    "human.pedestrian.adult": "pedestrian",
+    "human.pedestrian.child": "pedestrian",
+    "human.pedestrian.wheelchair": "ignore",
+    "human.pedestrian.stroller": "ignore",
+    "human.pedestrian.personal_mobility": "ignore",
+    "human.pedestrian.police_officer": "pedestrian",
+    "human.pedestrian.construction_worker": "pedestrian",
+    "animal": "ignore",
+    "vehicle.car": "car",
+    "vehicle.motorcycle": "motorcycle",
+    "vehicle.bicycle": "bicycle",
+    "vehicle.bus.bendy": "bus",
+    "vehicle.bus.rigid": "bus",
+    "vehicle.truck": "truck",
+    "vehicle.construction": "construction_vehicle",
+    "vehicle.emergency.ambulance": "ignore",
+    "vehicle.emergency.police": "ignore",
+    "vehicle.trailer": "trailer",
+    "movable_object.barrier": "barrier",
+    "movable_object.trafficcone": "traffic_cone",
+    "movable_object.pushable_pullable": "ignore",
+    "movable_object.debris": "ignore",
+    "static_object.bicycle_rack": "ignore",
+}
+
+
+cls_attr_dist = {
+    "barrier": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 0,
+        "vehicle.parked": 0,
+        "vehicle.stopped": 0,
+    },
+    "bicycle": {
+        "cycle.with_rider": 2791,
+        "cycle.without_rider": 8946,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 0,
+        "vehicle.parked": 0,
+        "vehicle.stopped": 0,
+    },
+    "bus": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 9092,
+        "vehicle.parked": 3294,
+        "vehicle.stopped": 3881,
+    },
+    "car": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 114304,
+        "vehicle.parked": 330133,
+        "vehicle.stopped": 46898,
+    },
+    "construction_vehicle": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 882,
+        "vehicle.parked": 11549,
+        "vehicle.stopped": 2102,
+    },
+    "ignore": {
+        "cycle.with_rider": 307,
+        "cycle.without_rider": 73,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 165,
+        "vehicle.parked": 400,
+        "vehicle.stopped": 102,
+    },
+    "motorcycle": {
+        "cycle.with_rider": 4233,
+        "cycle.without_rider": 8326,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 0,
+        "vehicle.parked": 0,
+        "vehicle.stopped": 0,
+    },
+    "pedestrian": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 157444,
+        "pedestrian.sitting_lying_down": 13939,
+        "pedestrian.standing": 46530,
+        "vehicle.moving": 0,
+        "vehicle.parked": 0,
+        "vehicle.stopped": 0,
+    },
+    "traffic_cone": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 0,
+        "vehicle.parked": 0,
+        "vehicle.stopped": 0,
+    },
+    "trailer": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 3421,
+        "vehicle.parked": 19224,
+        "vehicle.stopped": 1895,
+    },
+    "truck": {
+        "cycle.with_rider": 0,
+        "cycle.without_rider": 0,
+        "pedestrian.moving": 0,
+        "pedestrian.sitting_lying_down": 0,
+        "pedestrian.standing": 0,
+        "vehicle.moving": 21339,
+        "vehicle.parked": 55626,
+        "vehicle.stopped": 11097,
+    },
+}
+
+
+def get_available_scenes(nusc):
+    available_scenes = []
+    for scene in nusc.scene:
+        scene_token = scene["token"]
+        scene_rec = nusc.get("scene", scene_token)
+        sample_rec = nusc.get("sample", scene_rec["first_sample_token"])
+        sd_rec = nusc.get("sample_data", sample_rec["data"]["LIDAR_TOP"])
+        has_more_frames = True
+        scene_not_exist = False
+        while has_more_frames:
+            lidar_path, boxes, _ = nusc.get_sample_data(sd_rec["token"])
+            if not Path(lidar_path).exists():
+                scene_not_exist = True
+                break
+            else:
+                break
+        if scene_not_exist:
+            continue
+        available_scenes.append(scene)
+    return available_scenes
+
+
+def get_sample_data(nusc, sample_data_token, selected_anntokens=None):
+    """
+    Returns the data path as well as all annotations related to that sample_data.
+    Note that the boxes are transformed into the current sensor"s coordinate frame.
+    Args:
+        nusc:
+        sample_data_token: Sample_data token.
+        selected_anntokens: If provided only return the selected annotation.
+
+    Returns:
+
+    """
+    # Retrieve sensor & pose records
+    sd_record = nusc.get("sample_data", sample_data_token)
+    cs_record = nusc.get("calibrated_sensor", sd_record["calibrated_sensor_token"])
+    sensor_record = nusc.get("sensor", cs_record["sensor_token"])
+    pose_record = nusc.get("ego_pose", sd_record["ego_pose_token"])
+
+    data_path = nusc.get_sample_data_path(sample_data_token)
+
+    if sensor_record["modality"] == "camera":
+        cam_intrinsic = np.array(cs_record["camera_intrinsic"])
+    else:
+        cam_intrinsic = None
+
+    # Retrieve all sample annotations and map to sensor coordinate system.
+    if selected_anntokens is not None:
+        boxes = list(map(nusc.get_box, selected_anntokens))
+    else:
+        boxes = nusc.get_boxes(sample_data_token)
+
+    # Make list of Box objects including coord system transforms.
+    box_list = []
+    for box in boxes:
+        box.velocity = nusc.box_velocity(box.token)
+        # Move box to ego vehicle coord system
+        box.translate(-np.array(pose_record["translation"]))
+        box.rotate(Quaternion(pose_record["rotation"]).inverse)
+
+        #  Move box to sensor coord system
+        box.translate(-np.array(cs_record["translation"]))
+        box.rotate(Quaternion(cs_record["rotation"]).inverse)
+
+        box_list.append(box)
+
+    return data_path, box_list, cam_intrinsic
+
+
+def quaternion_yaw(q: Quaternion) -> float:
+    """
+    Calculate the yaw angle from a quaternion.
+    Note that this only works for a quaternion that represents a box in lidar or global coordinate frame.
+    It does not work for a box in the camera frame.
+    :param q: Quaternion of interest.
+    :return: Yaw angle in radians.
+    """
+
+    # Project into xy plane.
+    v = np.dot(q.rotation_matrix, np.array([1, 0, 0]))
+
+    # Measure yaw using arctan.
+    yaw = np.arctan2(v[1], v[0])
+
+    return yaw
+
+
+def obtain_sensor2top(
+    nusc, sensor_token, l2e_t, l2e_r_mat, e2g_t, e2g_r_mat, sensor_type="lidar"
+):
+    """Obtain the info with RT matric from general sensor to Top LiDAR.
+
+    Args:
+        nusc (class): Dataset class in the nuScenes dataset.
+        sensor_token (str): Sample data token corresponding to the
+            specific sensor type.
+        l2e_t (np.ndarray): Translation from lidar to ego in shape (1, 3).
+        l2e_r_mat (np.ndarray): Rotation matrix from lidar to ego
+            in shape (3, 3).
+        e2g_t (np.ndarray): Translation from ego to global in shape (1, 3).
+        e2g_r_mat (np.ndarray): Rotation matrix from ego to global
+            in shape (3, 3).
+        sensor_type (str): Sensor to calibrate. Default: "lidar".
+
+    Returns:
+        sweep (dict): Sweep information after transformation.
+    """
+    sd_rec = nusc.get("sample_data", sensor_token)
+    cs_record = nusc.get("calibrated_sensor", sd_rec["calibrated_sensor_token"])
+    pose_record = nusc.get("ego_pose", sd_rec["ego_pose_token"])
+    data_path = str(nusc.get_sample_data_path(sd_rec["token"]))
+    # if os.getcwd() in data_path:  # path from lyftdataset is absolute path
+    #     data_path = data_path.split(f"{os.getcwd()}/")[-1]  # relative path
+    sweep = {
+        "data_path": data_path,
+        "type": sensor_type,
+        "sample_data_token": sd_rec["token"],
+        "sensor2ego_translation": cs_record["translation"],
+        "sensor2ego_rotation": cs_record["rotation"],
+        "ego2global_translation": pose_record["translation"],
+        "ego2global_rotation": pose_record["rotation"],
+        "timestamp": sd_rec["timestamp"],
+    }
+    l2e_r_s = sweep["sensor2ego_rotation"]
+    l2e_t_s = sweep["sensor2ego_translation"]
+    e2g_r_s = sweep["ego2global_rotation"]
+    e2g_t_s = sweep["ego2global_translation"]
+
+    # obtain the RT from sensor to Top LiDAR
+    # sweep->ego->global->ego'->lidar
+    l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix
+    e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix
+    R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ (
+        np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T
+    )
+    T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ (
+        np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T
+    )
+    T -= (
+        e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
+        + l2e_t @ np.linalg.inv(l2e_r_mat).T
+    ).squeeze(0)
+    sweep["sensor2lidar_rotation"] = R.T  # points @ R.T + T
+    sweep["sensor2lidar_translation"] = T
+    return sweep
+
+
+def find_closest_cam_sd_token(nusc, start_cam_sd_token, target_timestamp):
+    """Find the camera sample_data closest in time to a LiDAR sweep."""
+    start_cam_sd = nusc.get("sample_data", start_cam_sd_token)
+    closest_cam_sd = start_cam_sd
+    closest_time_diff = abs(start_cam_sd["timestamp"] - target_timestamp)
+
+    current_cam_sd = start_cam_sd
+    while current_cam_sd["next"]:
+        next_cam_sd = nusc.get("sample_data", current_cam_sd["next"])
+        time_diff = abs(next_cam_sd["timestamp"] - target_timestamp)
+        if time_diff >= closest_time_diff:
+            break
+        closest_cam_sd = next_cam_sd
+        closest_time_diff = time_diff
+        current_cam_sd = next_cam_sd
+
+    current_cam_sd = start_cam_sd
+    while current_cam_sd["prev"]:
+        prev_cam_sd = nusc.get("sample_data", current_cam_sd["prev"])
+        time_diff = abs(prev_cam_sd["timestamp"] - target_timestamp)
+        if time_diff >= closest_time_diff:
+            break
+        closest_cam_sd = prev_cam_sd
+        closest_time_diff = time_diff
+        current_cam_sd = prev_cam_sd
+
+    return closest_cam_sd["token"]
+
+
+def fill_trainval_infos(
+    data_path, nusc, train_scenes, test=False, max_sweeps=10, with_camera=False
+):
+    train_nusc_infos = []
+    val_nusc_infos = []
+    progress_bar = tqdm.tqdm(
+        total=len(nusc.sample), desc="create_info", dynamic_ncols=True
+    )
+
+    ref_chan = "LIDAR_TOP"  # The radar channel from which we track back n sweeps to aggregate the point cloud.
+    chan = "LIDAR_TOP"  # The reference channel of the current sample_rec that the point clouds are mapped to.
+
+    camera_types = [
+        "CAM_FRONT",
+        "CAM_FRONT_RIGHT",
+        "CAM_FRONT_LEFT",
+        "CAM_BACK",
+        "CAM_BACK_LEFT",
+        "CAM_BACK_RIGHT",
+    ]
+
+    for index, sample in enumerate(nusc.sample):
+        progress_bar.update()
+
+        ref_sd_token = sample["data"][ref_chan]
+        ref_sd_rec = nusc.get("sample_data", ref_sd_token)
+        ref_cs_rec = nusc.get(
+            "calibrated_sensor", ref_sd_rec["calibrated_sensor_token"]
+        )
+        ref_pose_rec = nusc.get("ego_pose", ref_sd_rec["ego_pose_token"])
+        ref_time = 1e-6 * ref_sd_rec["timestamp"]
+
+        ref_lidar_path, ref_boxes, _ = get_sample_data(nusc, ref_sd_token)
+
+        ref_cam_front_token = sample["data"]["CAM_FRONT"]
+        ref_cam_path, _, ref_cam_intrinsic = nusc.get_sample_data(ref_cam_front_token)
+
+        # Homogeneous transform from ego car frame to reference frame
+        ref_from_car = transform_matrix(
+            ref_cs_rec["translation"], Quaternion(ref_cs_rec["rotation"]), inverse=True
+        )
+
+        # Homogeneous transformation matrix from global to _current_ ego car frame
+        car_from_global = transform_matrix(
+            ref_pose_rec["translation"],
+            Quaternion(ref_pose_rec["rotation"]),
+            inverse=True,
+        )
+        info = {
+            "lidar_path": Path(ref_lidar_path).relative_to(data_path).__str__(),
+            "lidar_token": ref_sd_token,
+            "cam_front_path": Path(ref_cam_path).relative_to(data_path).__str__(),
+            "cam_intrinsic": ref_cam_intrinsic,
+            "token": sample["token"],
+            "sweeps": [],
+            "ref_from_car": ref_from_car,
+            "car_from_global": car_from_global,
+            "timestamp": ref_time,
+        }
+        if with_camera:
+            info["cams"] = dict()
+            l2e_r = ref_cs_rec["rotation"]
+            l2e_t = (ref_cs_rec["translation"],)
+            e2g_r = ref_pose_rec["rotation"]
+            e2g_t = ref_pose_rec["translation"]
+            l2e_r_mat = Quaternion(l2e_r).rotation_matrix
+            e2g_r_mat = Quaternion(e2g_r).rotation_matrix
+
+            for cam in camera_types:
+                cam_token = sample["data"][cam]
+                cam_path, _, camera_intrinsics = nusc.get_sample_data(cam_token)
+                cam_info = obtain_sensor2top(
+                    nusc, cam_token, l2e_t, l2e_r_mat, e2g_t, e2g_r_mat, cam
+                )
+                cam_info["data_path"] = (
+                    Path(cam_info["data_path"]).relative_to(data_path).__str__()
+                )
+                cam_info.update(camera_intrinsics=camera_intrinsics)
+                info["cams"].update({cam: cam_info})
+
+        sample_data_token = sample["data"][chan]
+        curr_sd_rec = nusc.get("sample_data", sample_data_token)
+
+        sweeps = []
+        while len(sweeps) < max_sweeps - 1:
+            if curr_sd_rec["prev"] == "":
+                if len(sweeps) == 0:
+                    sweep = {
+                        "lidar_path": Path(ref_lidar_path)
+                        .relative_to(data_path)
+                        .__str__(),
+                        "sample_data_token": curr_sd_rec["token"],
+                        "transform_matrix": None,
+                        "time_lag": curr_sd_rec["timestamp"] * 0,
+                        "is_key_frame": curr_sd_rec["is_key_frame"],
+                    }
+                    if with_camera:
+                        sweep_sample = nusc.get("sample", curr_sd_rec["sample_token"])
+                        cam_infos = dict()
+                        for cam in camera_types:
+                            cam_token = find_closest_cam_sd_token(
+                                nusc,
+                                sweep_sample["data"][cam],
+                                curr_sd_rec["timestamp"],
+                            )
+                            cam_path, _, camera_intrinsics = nusc.get_sample_data(
+                                cam_token
+                            )
+                            cam_info = obtain_sensor2top(
+                                nusc,
+                                cam_token,
+                                l2e_t,
+                                l2e_r_mat,
+                                e2g_t,
+                                e2g_r_mat,
+                                cam,
+                            )
+                            cam_info["data_path"] = (
+                                Path(cam_info["data_path"])
+                                .relative_to(data_path)
+                                .__str__()
+                            )
+                            cam_info.update(camera_intrinsics=camera_intrinsics)
+                            cam_infos.update({cam: cam_info})
+                        sweep.update({"cams": cam_infos})
+                    sweeps.append(sweep)
+                else:
+                    sweeps.append(sweeps[-1])
+            else:
+                curr_sd_rec = nusc.get("sample_data", curr_sd_rec["prev"])
+                if with_camera:
+                    sweep_sample = nusc.get("sample", curr_sd_rec["sample_token"])
+                    cam_infos = dict()
+                    for cam in camera_types:
+                        cam_token = find_closest_cam_sd_token(
+                            nusc,
+                            sweep_sample["data"][cam],
+                            curr_sd_rec["timestamp"],
+                        )
+                        cam_path, _, camera_intrinsics = nusc.get_sample_data(cam_token)
+                        cam_info = obtain_sensor2top(
+                            nusc, cam_token, l2e_t, l2e_r_mat, e2g_t, e2g_r_mat, cam
+                        )
+                        cam_info["data_path"] = (
+                            Path(cam_info["data_path"]).relative_to(data_path).__str__()
+                        )
+                        cam_info.update(camera_intrinsics=camera_intrinsics)
+                        cam_infos.update({cam: cam_info})
+                # Get past pose
+                current_pose_rec = nusc.get("ego_pose", curr_sd_rec["ego_pose_token"])
+                global_from_car = transform_matrix(
+                    current_pose_rec["translation"],
+                    Quaternion(current_pose_rec["rotation"]),
+                    inverse=False,
+                )
+
+                # Homogeneous transformation matrix from sensor coordinate frame to ego car frame.
+                current_cs_rec = nusc.get(
+                    "calibrated_sensor", curr_sd_rec["calibrated_sensor_token"]
+                )
+                car_from_current = transform_matrix(
+                    current_cs_rec["translation"],
+                    Quaternion(current_cs_rec["rotation"]),
+                    inverse=False,
+                )
+
+                tm = reduce(
+                    np.dot,
+                    [ref_from_car, car_from_global, global_from_car, car_from_current],
+                )
+
+                lidar_path = nusc.get_sample_data_path(curr_sd_rec["token"])
+
+                time_lag = ref_time - 1e-6 * curr_sd_rec["timestamp"]
+
+                sweep = {
+                    "lidar_path": Path(lidar_path).relative_to(data_path).__str__(),
+                    "sample_data_token": curr_sd_rec["token"],
+                    "transform_matrix": tm,
+                    "global_from_car": global_from_car,
+                    "car_from_current": car_from_current,
+                    "time_lag": time_lag,
+                    "is_key_frame": curr_sd_rec["is_key_frame"],
+                }
+                if with_camera:
+                    sweep.update({"cams": cam_infos})
+
+                sweeps.append(sweep)
+
+        info["sweeps"] = sweeps
+
+        assert len(info["sweeps"]) == max_sweeps - 1, (
+            f"sweep {curr_sd_rec['token']} only has {len(info['sweeps'])} sweeps, "
+            f"you should duplicate to sweep num {max_sweeps - 1}"
+        )
+
+        if not test:
+            # processing gt bbox
+            annotations = [
+                nusc.get("sample_annotation", token) for token in sample["anns"]
+            ]
+
+            # the filtering gives 0.5~1 map improvement
+            num_lidar_pts = np.array([anno["num_lidar_pts"] for anno in annotations])
+            num_radar_pts = np.array([anno["num_radar_pts"] for anno in annotations])
+            mask = num_lidar_pts + num_radar_pts > 0
+
+            locs = np.array([b.center for b in ref_boxes]).reshape(-1, 3)
+            dims = np.array([b.wlh for b in ref_boxes]).reshape(-1, 3)[
+                :, [1, 0, 2]
+            ]  # wlh == > dxdydz (lwh)
+            velocity = np.array([b.velocity for b in ref_boxes]).reshape(-1, 3)
+            rots = np.array([quaternion_yaw(b.orientation) for b in ref_boxes]).reshape(
+                -1, 1
+            )
+            names = np.array([b.name for b in ref_boxes])
+            tokens = np.array([b.token for b in ref_boxes])
+            gt_boxes = np.concatenate([locs, dims, rots, velocity[:, :2]], axis=1)
+
+            assert len(annotations) == len(gt_boxes) == len(velocity)
+
+            info["gt_boxes"] = gt_boxes[mask, :]
+            info["gt_boxes_velocity"] = velocity[mask, :]
+            info["gt_names"] = np.array(
+                [map_name_from_general_to_detection[name] for name in names]
+            )[mask]
+            info["gt_boxes_token"] = tokens[mask]
+            info["num_lidar_pts"] = num_lidar_pts[mask]
+            info["num_radar_pts"] = num_radar_pts[mask]
+
+            # processing gt segment
+            segment_path = nusc.get("lidarseg", ref_sd_token)["filename"]
+            info["gt_segment_path"] = segment_path
+
+        if sample["scene_token"] in train_scenes:
+            train_nusc_infos.append(info)
+        else:
+            val_nusc_infos.append(info)
+
+    progress_bar.close()
+    return train_nusc_infos, val_nusc_infos
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root", required=True, help="Path to the nuScenes dataset."
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where processed information located.",
+    )
+    parser.add_argument(
+        "--max_sweeps", default=10, type=int, help="Max number of sweeps. Default: 10."
+    )
+    parser.add_argument(
+        "--with_camera",
+        action="store_true",
+        default=False,
+        help="Whether use camera or not.",
+    )
+    config = parser.parse_args()
+
+    print(f"Loading nuScenes tables for version v1.0-trainval...")
+    nusc_trainval = NuScenes(
+        version="v1.0-trainval", dataroot=config.dataset_root, verbose=False
+    )
+    available_scenes_trainval = get_available_scenes(nusc_trainval)
+    available_scene_names_trainval = [s["name"] for s in available_scenes_trainval]
+    print("total scene num:", len(nusc_trainval.scene))
+    print("exist scene num:", len(available_scenes_trainval))
+    assert len(available_scenes_trainval) == len(nusc_trainval.scene) == 850
+
+    print(f"Loading nuScenes tables for version v1.0-test...")
+    nusc_test = NuScenes(
+        version="v1.0-test", dataroot=config.dataset_root, verbose=False
+    )
+    available_scenes_test = get_available_scenes(nusc_test)
+    available_scene_names_test = [s["name"] for s in available_scenes_test]
+    print("total scene num:", len(nusc_test.scene))
+    print("exist scene num:", len(available_scenes_test))
+    assert len(available_scenes_test) == len(nusc_test.scene) == 150
+
+    train_scenes = splits.train
+    train_scenes = set(
+        [
+            available_scenes_trainval[available_scene_names_trainval.index(s)]["token"]
+            for s in train_scenes
+        ]
+    )
+    test_scenes = splits.test
+    test_scenes = set(
+        [
+            available_scenes_test[available_scene_names_test.index(s)]["token"]
+            for s in test_scenes
+        ]
+    )
+    print(f"Filling trainval information...")
+    train_nusc_infos, val_nusc_infos = fill_trainval_infos(
+        config.dataset_root,
+        nusc_trainval,
+        train_scenes,
+        test=False,
+        max_sweeps=config.max_sweeps,
+        with_camera=config.with_camera,
+    )
+    print(f"Filling test information...")
+    test_nusc_infos, _ = fill_trainval_infos(
+        config.dataset_root,
+        nusc_test,
+        test_scenes,
+        test=True,
+        max_sweeps=config.max_sweeps,
+        with_camera=config.with_camera,
+    )
+
+    print(f"Saving nuScenes information...")
+    os.makedirs(os.path.join(config.output_root, "info"), exist_ok=True)
+    print(
+        f"train sample: {len(train_nusc_infos)}, val sample: {len(val_nusc_infos)}, test sample: {len(test_nusc_infos)}"
+    )
+    with open(
+        os.path.join(
+            config.output_root,
+            "info",
+            f"nuscenes_infos_{config.max_sweeps}sweeps_train.pkl",
+        ),
+        "wb",
+    ) as f:
+        pickle.dump(train_nusc_infos, f)
+    with open(
+        os.path.join(
+            config.output_root,
+            "info",
+            f"nuscenes_infos_{config.max_sweeps}sweeps_val.pkl",
+        ),
+        "wb",
+    ) as f:
+        pickle.dump(val_nusc_infos, f)
+    with open(
+        os.path.join(
+            config.output_root,
+            "info",
+            f"nuscenes_infos_{config.max_sweeps}sweeps_test.pkl",
+        ),
+        "wb",
+    ) as f:
+        pickle.dump(test_nusc_infos, f)

pointcept/datasets/preprocessing/concerto/re10k/preprocess_re10k.py

@@ -0,0 +1,556 @@
+"""
+Preprocessing Script for RE10K using VGGT(https://github.com/facebookresearch/vggt)
+
+Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+from pathlib import Path
+from torchvision.utils import save_image
+import torch
+from io import BytesIO
+from einops import rearrange, repeat
+from PIL import Image
+from tqdm import tqdm
+import torchvision.transforms as tf
+import argparse
+import os
+import cv2
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+import open3d as o3d
+from torchvision.utils import save_image
+
+from vggt.models.vggt import VGGT
+from vggt.utils.pose_enc import pose_encoding_to_extri_intri
+import camtools as ct
+from torchvision.transforms import transforms as T
+
+_alignment = 14
+target_size = 518
+
+generator = torch.Generator()
+generator.manual_seed(123)
+
+
+def extract_and_align_ground_plane(
+    pcd,
+    height_percentile=20,
+    ransac_distance_threshold=0.01,
+    ransac_n=3,
+    ransac_iterations=1000,
+    max_angle_degree=40,
+    max_trials=6,
+):
+    points = np.asarray(pcd.points)
+    z_vals = points[:, 2]
+    z_thresh = np.percentile(z_vals, height_percentile)
+    low_indices = np.where(z_vals <= z_thresh)[0]
+
+    remaining_indices = low_indices.copy()
+
+    for trial in range(max_trials):
+        if len(remaining_indices) < ransac_n:
+            raise ValueError("Not enough points left to fit a plane.")
+
+        low_pcd = pcd.select_by_index(remaining_indices)
+
+        plane_model, inliers = low_pcd.segment_plane(
+            distance_threshold=ransac_distance_threshold,
+            ransac_n=ransac_n,
+            num_iterations=ransac_iterations,
+        )
+        a, b, c, d = plane_model
+        normal = np.array([a, b, c])
+        normal /= np.linalg.norm(normal)
+
+        angle = np.arccos(np.clip(np.dot(normal, [0, 0, 1]), -1.0, 1.0)) * 180 / np.pi
+        if angle <= max_angle_degree:
+            inliers_global = remaining_indices[inliers]
+
+            target = np.array([0, 0, 1])
+            axis = np.cross(normal, target)
+            axis_norm = np.linalg.norm(axis)
+
+            if axis_norm < 1e-6:
+                rotation_matrix = np.eye(3)
+            else:
+                axis /= axis_norm
+                rot_angle = np.arccos(np.clip(np.dot(normal, target), -1.0, 1.0))
+                rotation = R.from_rotvec(axis * rot_angle)
+                rotation_matrix = rotation.as_matrix()
+
+            rotated_points = points @ rotation_matrix.T
+            ground_points_z = rotated_points[inliers_global, 2]
+            offset = np.mean(ground_points_z)
+            rotated_points[:, 2] -= offset
+
+            aligned_pcd = o3d.geometry.PointCloud()
+            aligned_pcd.points = o3d.utility.Vector3dVector(rotated_points)
+            if pcd.has_colors():
+                aligned_pcd.colors = pcd.colors
+            if pcd.has_normals():
+                rotated_normals = np.asarray(pcd.normals) @ rotation_matrix.T
+                aligned_pcd.normals = o3d.utility.Vector3dVector(rotated_normals)
+
+            return aligned_pcd, inliers_global, rotation_matrix, offset
+
+        else:
+            rejected_indices = remaining_indices[inliers]
+            remaining_indices = np.setdiff1d(remaining_indices, rejected_indices)
+
+    raise ValueError("Failed to find a valid ground plane within max trials.")
+
+
+def rotx(x, theta=90):
+    """
+    Rotate x by theta degrees around the x-axis
+    """
+    theta = np.deg2rad(theta)
+    rot_matrix = np.array(
+        [
+            [1, 0, 0, 0],
+            [0, np.cos(theta), -np.sin(theta), 0],
+            [0, np.sin(theta), np.cos(theta), 0],
+            [0, 0, 0, 1],
+        ]
+    )
+    return rot_matrix @ x
+
+
+def Coord2zup(points, extrinsics):
+    """
+    Convert the dust3r coordinate system to the z-up coordinate system
+    """
+    points = np.concatenate([points, np.ones([points.shape[0], 1])], axis=1).T
+    points = rotx(points, -90)[:3].T
+    t = np.min(points, axis=0)
+    points -= t
+    extrinsics = rotx(extrinsics, -90)
+    extrinsics[:, :3, 3] -= t.T
+    return points, extrinsics
+
+
+def resize_images_intrinsic(images, intrinsics, size, crop_size):
+    h, w = size
+    crop_h, crop_w = crop_size
+    intrinsics[:, 0, 0] = crop_w
+    intrinsics[:, 1, 1] = crop_h
+    cam_trans = torch.tensor(
+        [[w / crop_w, 0, 0], [0, h / crop_h, 0], [0, 0, 1]], dtype=torch.float32
+    )
+    intrinsics = torch.stack([cam_trans @ intrinsics_i for intrinsics_i in intrinsics])
+    return images, intrinsics
+
+
+def calDelta(ang1, ang2, dist1, dist2):
+    alpha = 20
+    angD = ang1 - ang2
+    distD = dist1 - dist2
+    ang_s = np.linalg.norm(angD)
+    dist_s = np.linalg.norm(distD)
+    return ang_s + alpha * dist_s, ang_s, dist_s
+
+
+def convert_poses(poses):
+    b, _ = poses.shape
+
+    # Convert the intrinsics to a 3x3 normalized K matrix.
+    intrinsics = torch.eye(3, dtype=torch.float32)
+    intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone()
+    fx, fy, cx, cy = poses[:, :4].T
+    intrinsics[:, 0, 0] = fx
+    intrinsics[:, 1, 1] = fy
+    intrinsics[:, 0, 2] = cx
+    intrinsics[:, 1, 2] = cy
+
+    # Convert the extrinsics to a 4x4 OpenCV-style W2C matrix.
+    w2c = repeat(torch.eye(4, dtype=torch.float32), "h w -> b h w", b=b).clone()
+    w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4)
+    return w2c.inverse(), intrinsics
+
+
+def convert_images(images):
+    torch_images = []
+    for image in images:
+        image = Image.open(BytesIO(image.numpy().tobytes()))
+        torch_images.append(tf.ToTensor()(image))
+    return torch.stack(torch_images)
+
+
+def plyGet(
+    images,
+    scene_dir,
+    pc_outputdir,
+    im_outputdir,
+    device,
+    dtype,
+    vggt_model,
+    depth_points="points",
+    conf=0.0,
+    parse_depths=False,
+):
+    # T0 = T0.numpy()
+    print(f"Processing scene: {scene_dir}")
+    height, width = images.shape[-2:]
+
+    # Skip every n frames
+    print(f"Processing {images.shape[0]} images")
+
+    image_num = images.shape[0]
+    # Load and preprocess images
+    images = images.to(device)
+
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=dtype):
+            # Predict attributes including cameras, depth maps, and point maps.
+            predictions = vggt_model(images)
+
+    # Convert pose encoding to extrinsic and intrinsic matrices
+    Ts, Ks = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
+    Ts = Ts[0].detach().cpu().numpy()
+    Ks = Ks[0].detach().cpu().numpy()
+    Ts = ct.convert.pad_0001(Ts)
+    Ts_inv = np.linalg.inv(Ts)
+    Cs = np.array([ct.convert.T_to_C(T) for T in Ts])  # (n, 3)
+
+    # correspondance
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel_ = []
+    for i in range(image_num):
+        pixel_id = np.hstack((pixel, i * np.ones((pixel.shape[0], 1))))
+        pixel_.append(pixel_id)
+    pixel = np.concatenate(pixel_, axis=0)
+
+    if depth_points == "points":
+        world_points = predictions["world_points"].detach().cpu().numpy()
+        world_points_conf = predictions["world_points_conf"].detach().cpu().numpy()
+        img_depths = predictions["depth"].detach().cpu().numpy().squeeze(-1)
+    else:
+        img_depths = predictions["depth"].detach().cpu().numpy().squeeze(-1)
+        img_depths_conf = predictions["depth_conf"].detach().cpu().numpy()
+        world_points = [
+            ct.project.im_depth_to_point_cloud(img_depth, K, T)
+            for img_depth, K, T in zip(img_depths[0], Ks, Ts)
+        ]
+        world_points = np.stack(world_points).reshape(
+            1, image_num, images.shape[-2], images.shape[-1], -1
+        )
+        world_points_conf = img_depths_conf
+
+    masks = world_points_conf > conf
+    points_masks = masks.reshape((image_num * images.shape[-2] * images.shape[-1],))
+
+    # Compute view direction for each pixel
+    # (b n h w c) - (n, 3)
+    view_dirs = world_points - rearrange(Cs, "n c -> 1 n 1 1 c")
+    view_dirs = rearrange(view_dirs, "b n h w c -> (b n h w) c")
+    view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True)
+
+    # Extract points and colors
+    # [1, 8, 3, 294, 518]
+    images = images.detach().cpu().unsqueeze(0).numpy()
+    points = rearrange(world_points, "b n h w c -> (b n h w) c")
+    colors = rearrange(images, "b n c h w -> (b n h w) c")
+
+    # masks
+    points = points[points_masks]
+    colors = colors[points_masks]
+    pixel = pixel[points_masks]
+    pixel = np.concatenate((pixel, np.arange(points.shape[0]).reshape(-1, 1)), axis=-1)
+
+    points, Ts_inv = Coord2zup(points, Ts_inv)
+    scale = 3 / (points[:, 2].max() - points[:, 2].min())
+    points *= scale
+    Ts_inv[:, :3, 3] *= scale
+
+    # Create a point cloud
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    pcd.colors = o3d.utility.Vector3dVector(colors)
+    pcd.estimate_normals()
+    try:
+        pcd, inliers, rotation_matrix, offset = extract_and_align_ground_plane(pcd)
+    except:
+        print("cannot find ground")
+        return
+    T_pcd = np.eye(4)
+    T_pcd[:3, :3] = rotation_matrix
+    T_pcd[2, 3] = -offset
+    Ts_inv = T_pcd @ Ts_inv
+    os.makedirs(pc_outputdir, exist_ok=True)
+    os.makedirs(im_outputdir, exist_ok=True)
+
+    # Filp normals such that normals always point to camera
+    # Compute the dot product between the normal and the view direction
+    # If the dot product is less than 0, flip the normal
+    normals = np.asarray(pcd.normals)
+    view_dirs = np.asarray(view_dirs)
+    view_dirs = view_dirs[points_masks]
+    dot_product = np.sum(normals * view_dirs, axis=-1)
+    flip_mask = dot_product > 0
+    normals[flip_mask] = -normals[flip_mask]
+
+    # Normalize normals a nd m
+    normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
+    colors = (colors * 255).astype(np.uint8)
+
+    np.save(Path(pc_outputdir) / "coord.npy", points)
+    np.save(Path(pc_outputdir) / "color.npy", colors)
+    np.save(Path(pc_outputdir) / "normal.npy", normals)
+
+    correspondence_path = os.path.join(im_outputdir, "correspondence")
+    img_output_path = os.path.join(im_outputdir, "color")
+    Ts_output_path = os.path.join(im_outputdir, "pose")
+    Ks_output_path = os.path.join(im_outputdir, "intrinsic")
+    os.makedirs(img_output_path, exist_ok=True)
+    os.makedirs(Ts_output_path, exist_ok=True)
+    os.makedirs(Ks_output_path, exist_ok=True)
+    os.makedirs(correspondence_path, exist_ok=True)
+    for id in range(image_num):
+        pixels_correspondence = pixel[pixel[:, 2] == id]
+        pixels_correspondence = pixels_correspondence[:, [0, 1, 3]]
+        np.save(
+            os.path.join(im_outputdir, "correspondence", f"{id}.npy"),
+            pixels_correspondence,
+        )
+    images = torch.tensor(images.squeeze(0))
+    if parse_depths:
+        depth_output_path = os.path.join(im_outputdir, "depth")
+        os.makedirs(depth_output_path, exist_ok=True)
+        img_depths = img_depths.squeeze(0)[..., np.newaxis] * scale * 1000
+    for id in range(images.shape[0]):
+        save_image(images[id], os.path.join(img_output_path, f"{id}.png"))
+        np.save(os.path.join(Ts_output_path, f"{id}.npy"), Ts_inv[id])
+        np.save(os.path.join(Ks_output_path, f"{id}.npy"), Ks[id])
+        if parse_depths:
+            cv2.imwrite(
+                os.path.join(depth_output_path, f"{id}.png"),
+                img_depths[id].astype(np.uint16),
+            )
+
+
+def parse_scene(
+    chunk_path,
+    vggt_model,
+    num_context_views,
+    outputdir,
+    device,
+    frame_gap,
+    overlap_range,
+    dtype,
+    conf,
+    parse_depths,
+):
+    print("loading from ", chunk_path)
+    split = chunk_path.parts[-2]
+    # Load the chunk.
+    chunk = torch.load(chunk_path)
+    for run_idx in range(len(chunk)):
+        example = chunk[run_idx]
+        extrinsics, intrinsics = convert_poses(example["cameras"])
+        scene = example["key"]
+
+        # Load the images.
+        context_images = example["images"]
+        context_images = convert_images(context_images)
+
+        v, _, _ = extrinsics.shape
+        ang_list = np.array(
+            [
+                R.from_matrix(ex_i[:3, :3].cpu().numpy()).as_euler("xyz", degrees=True)
+                for ex_i in extrinsics
+            ]
+        )
+        dist_list = np.array(extrinsics[:, :3, 3].reshape((-1, 3)).cpu())
+        context_indices = torch.randperm(v, generator=generator)
+        for context_index in tqdm(context_indices, "Finding context pair"):
+            choose_index = [context_index.item()]
+            while len(choose_index) < num_context_views:
+                # Step away from context view until the minimum overlap threshold is met.
+                valid_indices = []
+                for step in (1, -1):
+                    if step == 1:
+                        context_index = max(choose_index)
+                    else:
+                        context_index = min(choose_index)
+                    min_distance = frame_gap[0]
+                    max_distance = frame_gap[1]
+                    current_index = context_index + step * min_distance
+
+                    while 0 <= current_index < v:
+                        # Compute overlap.
+                        overlap, overlap_a, overlap_b = calDelta(
+                            ang_list[context_index],
+                            ang_list[current_index],
+                            dist_list[context_index],
+                            dist_list[current_index],
+                        )
+                        # print("similarity",overlap_a,overlap_b)
+                        min_overlap = overlap_range[0]
+                        max_overlap = overlap_range[1]
+                        if min_overlap <= overlap <= max_overlap:
+                            valid_indices.append((current_index, overlap_a, overlap_b))
+                        delta = np.abs(current_index - context_index)
+
+                        # Stop once the camera has panned away too much.
+                        if overlap < min_overlap or delta > max_distance:
+                            break
+
+                        current_index += step
+                if valid_indices:
+                    # Pick a random valid view. Index the resulting views.
+                    num_options = len(valid_indices)
+                    chosen = torch.randint(
+                        0, num_options, size=tuple(), generator=generator
+                    )
+                    chosen, overlap_a, overlap_b = valid_indices[chosen]
+                    choose_index.append(chosen)
+                else:
+                    break
+            if len(choose_index) < num_context_views:
+                continue
+            context_left = min(choose_index)
+            context_right = max(choose_index)
+            delta = context_right - context_left
+            choose_index.sort()
+            # Pick non-repeated random target views.
+
+            extrinsics = extrinsics[choose_index]
+            intrinsics = intrinsics[choose_index]
+            img_height, img_width = context_images.shape[-2:]
+            img_width_new = target_size
+            img_height_new = (
+                round(img_height * (img_width_new / img_width) / _alignment)
+                * _alignment
+            )
+            transform = T.Compose(
+                [
+                    T.Resize((img_height_new, img_width_new)),
+                ]
+            )
+            imgs_list = transform(context_images)
+            if img_height_new > target_size:
+                start_y = (img_height_new - target_size) // 2
+                context_images = context_images[
+                    :, :, start_y : start_y + target_size, :
+                ]
+            imgs_list, intrinsics_list = resize_images_intrinsic(
+                imgs_list,
+                intrinsics,
+                (img_height_new, img_width_new),
+                (target_size, target_size),
+            )
+            pc_outputdir = os.path.join(outputdir, split, scene)
+            im_outputdir = os.path.join(outputdir, "images", split, scene)
+            plyGet(
+                imgs_list[choose_index],
+                scene,
+                pc_outputdir,
+                im_outputdir,
+                device,
+                dtype,
+                vggt_model,
+                depth_points="points",
+                conf=conf,
+                parse_depths=parse_depths,
+            )
+            break
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--num_context_views",
+        default=4,
+        type=int,
+        help="num_context_views",
+    )
+    parser.add_argument(
+        "--frame_gap",
+        default=[15, 135],
+        type=int,
+        help="frame_gap",
+    )
+    parser.add_argument(
+        "--overlap_range",
+        default=[5, 40],
+        type=int,
+        help="overlap_range",
+    )
+    parser.add_argument(
+        "--conf",
+        default=0.0,
+        type=float,
+        help="overlap_range",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=1,
+        type=int,
+        help="num_workers",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="thread_id",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        type=str,
+        help="device",
+    )
+    parser.add_argument(
+        "--splits",
+        default=["train", "test"],
+        nargs="+",
+        choices=["train", "test"],
+        help="Splits need to process.",
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    cfg = parser.parse_args()
+    VGGT_model = VGGT().to(cfg.device)
+    _URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt"
+    state_dict = torch.hub.load_state_dict_from_url(_URL)
+    VGGT_model.load_state_dict(state_dict)
+    dtype = (
+        torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+    )
+
+    chunk_paths = []
+    for split in cfg.splits:
+        root = Path(cfg.dataset_root) / split
+        root_chunks = [entry for entry in root.iterdir() if ".json" not in str(entry)]
+        root_chunks = sorted(root_chunks)
+        chunk_paths.extend(root_chunks)
+    chunk_paths_list = np.array_split(chunk_paths, cfg.num_workers)
+    chunk_paths_ = chunk_paths_list[cfg.thread_id]
+
+    for chunk_path in chunk_paths_:
+        parse_scene(
+            chunk_path,
+            VGGT_model,
+            cfg.num_context_views,
+            cfg.output_root,
+            cfg.device,
+            cfg.frame_gap,
+            cfg.overlap_range,
+            dtype,
+            cfg.conf,
+            cfg.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/re10k/splits.py

@@ -0,0 +1,57 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train", "test"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(
+            dataset_path, "data/re10k_align"
+        )
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            im_split_name_path = os.path.join(im_split_path, name, "color")
+            co_split_name_path = os.path.join(im_split_path, name, "correspondence")
+            png_files = [
+                f for f in os.listdir(im_split_name_path) if f.endswith(".png")
+            ]
+            png_files = sorted(png_files, key=lambda x: int(x.split(".")[0]))
+            # Get the full paths of the .png files
+            png_file_paths = [
+                os.path.join(im_split_name_path, f).replace(
+                    dataset_path, "data/re10k_align"
+                )
+                for f in png_files
+            ]
+            co_file_paths = [
+                os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace(
+                    dataset_path, "data/re10k_align"
+                )
+                for f in png_files
+            ]
+            split_dict[f"{name}"] = {}
+            split_dict[f"{name}"]["pointclouds"] = os.path.join(pc_split_path, name)
+            split_dict[f"{name}"]["images"] = png_file_paths
+            split_dict[f"{name}"]["correspondences"] = co_file_paths
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the re10k dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/s3dis/preprocess_s3dis.py

@@ -0,0 +1,472 @@
+"""
+Preprocessing Script for S3DIS
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import os
+import argparse
+import glob
+import numpy as np
+import re
+import json
+from PIL import Image
+from scipy.spatial import cKDTree
+from pathlib import Path
+import camtools as ct
+import shutil
+
+try:
+    import open3d
+except ImportError:
+    import warnings
+
+    warnings.warn("Please install open3d for parsing normal")
+
+try:
+    import trimesh
+except ImportError:
+    import warnings
+
+    warnings.warn("Please install trimesh for parsing normal")
+
+area_mesh_dict = {}
+
+
+def correspondenceGet(mesh, K, T, img_size, coords_gt):
+    height, width = img_size
+    if np.isnan(T).any() or np.isinf(T).any():
+        return None
+    depth = ct.raycast.mesh_to_im_depth(
+        mesh=mesh, K=K, T=np.linalg.inv(T), height=height, width=width
+    )
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    valid = ~np.isinf(depth).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict
+
+
+def correspondenceSave(mesh, output_dir, img_size, coords_gt, Ks, Ts):
+    output_dir = Path(output_dir)
+    (output_dir / "correspondence").mkdir(exist_ok=True)
+    Ks = Ks[:, :3, :3]
+    coords_gt_ = coords_gt
+    pixels_ = []
+    coords_ = []
+    index_gt = [
+        img_name.split(".")[0]
+        for img_name in os.listdir(str(output_dir / "color"))
+        if img_name.endswith(".png")
+    ]
+    index_gt.sort(key=lambda x: int(x.split("_frame_")[1].split("_domain_")[0]))
+    for i, (K, T) in enumerate(zip(Ks, Ts)):
+        coord_dict = correspondenceGet(mesh, K, T, img_size, coords_gt)
+        if coord_dict is None:
+            correspondences = -np.ones((1, 3))
+        else:
+            pixels_ = coord_dict["pixel"]
+            coords_ = coord_dict["coord"]
+
+            tree = cKDTree(coords_gt_)
+            dis, idx = tree.query(coords_, k=1)
+            idx_valid = idx[dis < 0.01]
+            pixels_valid = pixels_[dis < 0.01]
+            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+        np.save(
+            output_dir / "correspondence" / "{}.npy".format(index_gt[i]),
+            correspondences,
+        )
+
+
+def parse_room(
+    room,
+    angle,
+    rgb_gap,
+    raw_root,
+    dataset_root,
+    output_root,
+    align_angle=True,
+    parse_normal=False,
+    parse_pointclouds=False,
+    parse_depths=False,
+):
+    print("Parsing: {}".format(room))
+    classes = [
+        "ceiling",
+        "floor",
+        "wall",
+        "beam",
+        "column",
+        "window",
+        "door",
+        "table",
+        "chair",
+        "sofa",
+        "bookcase",
+        "board",
+        "clutter",
+    ]
+    class2label = {cls: i for i, cls in enumerate(classes)}
+    source_dir = os.path.join(dataset_root, room)
+    output_root_pointclouds = output_root
+    output_root_imgs = os.path.join(output_root, "images")
+    os.makedirs(output_root_imgs, exist_ok=True)
+    save_path_imgs = os.path.join(output_root_imgs, room)
+    os.makedirs(save_path_imgs, exist_ok=True)
+    object_path_list = sorted(glob.glob(os.path.join(source_dir, "Annotations/*.txt")))
+
+    room_coords = []
+    room_colors = []
+    room_normals = []
+    room_semantic_gt = []
+    room_instance_gt = []
+
+    for object_id, object_path in enumerate(object_path_list):
+        object_name = os.path.basename(object_path).split("_")[0]
+        obj = np.loadtxt(object_path)
+        coords = obj[:, :3]
+        colors = obj[:, 3:6]
+        # note: in some room there is 'stairs' class
+        class_name = object_name if object_name in classes else "clutter"
+        semantic_gt = np.repeat(class2label[class_name], coords.shape[0])
+        semantic_gt = semantic_gt.reshape([-1, 1])
+        instance_gt = np.repeat(object_id, coords.shape[0])
+        instance_gt = instance_gt.reshape([-1, 1])
+
+        room_coords.append(coords)
+        room_colors.append(colors)
+        room_semantic_gt.append(semantic_gt)
+        room_instance_gt.append(instance_gt)
+
+    room_coords = np.ascontiguousarray(np.vstack(room_coords))
+    room_colors = np.ascontiguousarray(np.vstack(room_colors))
+
+    if parse_normal:
+        x_min, z_max, y_min = np.min(room_coords, axis=0)
+        x_max, z_min, y_max = np.max(room_coords, axis=0)
+        z_max = -z_max
+        z_min = -z_min
+        max_bound = np.array([x_max, y_max, z_max]) + 0.1
+        min_bound = np.array([x_min, y_min, z_min]) - 0.1
+        bbox = open3d.geometry.AxisAlignedBoundingBox(
+            min_bound=min_bound, max_bound=max_bound
+        )
+        # crop room
+        room_mesh = (
+            area_mesh_dict[os.path.dirname(room)]
+            .crop(bbox)
+            .transform(
+                np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+            )
+        )
+        vertices = np.array(room_mesh.vertices)
+        faces = np.array(room_mesh.triangles)
+        vertex_normals = np.array(room_mesh.vertex_normals)
+        o3d_room_mesh = room_mesh
+        room_mesh = trimesh.Trimesh(
+            vertices=vertices, faces=faces, vertex_normals=vertex_normals
+        )
+        closest_points, distances, face_id = room_mesh.nearest.on_surface(room_coords)
+        room_normals = room_mesh.face_normals[face_id]
+    Ks = []
+    Ts = []
+
+    area = room.split("/")[0].lower()
+    part = room.split("/")[1]
+    if area == "area_5":
+        camera_pose_paths = sorted(
+            glob.glob(
+                os.path.join(raw_root, f"{area}*", "data", "pose", f"*_{part}_*.json")
+            )
+        )
+        camera_rgb_paths = sorted(
+            glob.glob(
+                os.path.join(raw_root, f"{area}*", "data", "rgb", f"*_{part}_*.png")
+            )
+        )
+        camera_depth_paths = sorted(
+            glob.glob(
+                os.path.join(raw_root, f"{area}*", "data", "depth", f"*_{part}_*.png")
+            )
+        )
+    else:
+        camera_pose_paths = sorted(
+            glob.glob(os.path.join(raw_root, area, "data", "pose", f"*_{part}_*.json"))
+        )
+        camera_rgb_paths = sorted(
+            glob.glob(os.path.join(raw_root, area, "data", "rgb", f"*_{part}_*.png"))
+        )
+        camera_depth_paths = sorted(
+            glob.glob(os.path.join(raw_root, area, "data", "depth", f"*_{part}_*.png"))
+        )
+    camera_pose_names = set(
+        [i.split(f"_{part}_")[0].split("camera_")[1] for i in camera_pose_paths]
+    )
+    camera_rgb_names = set(
+        [i.split(f"_{part}_")[0].split("camera_")[1] for i in camera_rgb_paths]
+    )
+    camera_names = camera_pose_names.intersection(camera_rgb_names)
+    for camera_name in camera_names:
+        camera_pose_paths_withname = [
+            path for path in camera_pose_paths if camera_name in path
+        ]
+        camera_rgb_paths_withname = [
+            path for path in camera_rgb_paths if camera_name in path
+        ]
+        camera_depth_paths_withname = [
+            path for path in camera_depth_paths if camera_name in path
+        ]
+        pose_frame_ids = set(
+            [re.search(r"frame_(\d+)", s).group(1) for s in camera_pose_paths_withname]
+        )
+        rgb_frame_ids = set(
+            [re.search(r"frame_(\d+)", s).group(1) for s in camera_rgb_paths_withname]
+        )
+        frame_ids = pose_frame_ids.intersection(rgb_frame_ids)
+        frame_ids_selected = sorted(list(map(int, list(frame_ids))))[::rgb_gap]
+        camera_pose_prefix = camera_pose_paths_withname[0].split("frame_")[0] + "frame_"
+        camera_pose_postfix = (
+            "_domain" + camera_pose_paths_withname[0].split("_domain")[1]
+        )
+        camera_pose_path_selected = [
+            camera_pose_prefix + str(id) + camera_pose_postfix
+            for id in frame_ids_selected
+        ]
+        camera_rgb_prefix = camera_rgb_paths_withname[0].split("frame_")[0] + "frame_"
+        camera_rgb_postfix = (
+            "_domain" + camera_rgb_paths_withname[0].split("_domain")[1]
+        )
+        camera_rgb_path_selected = [
+            camera_rgb_prefix + str(id) + camera_rgb_postfix
+            for id in frame_ids_selected
+        ]
+        camera_depth_prefix = (
+            camera_depth_paths_withname[0].split("frame_")[0] + "frame_"
+        )
+        camera_depth_postfix = (
+            "_domain" + camera_depth_paths_withname[0].split("_domain")[1]
+        )
+        camera_depth_path_selected = [
+            camera_depth_prefix + str(id) + camera_depth_postfix
+            for id in frame_ids_selected
+        ]
+        image_shape = Image.open(camera_rgb_path_selected[0]).size
+        save_path_imgs_camera = os.path.join(save_path_imgs, camera_name)
+        imgs_output_path = os.path.join(save_path_imgs_camera, "color")
+        os.makedirs(imgs_output_path, exist_ok=True)
+        for img_path in camera_rgb_path_selected:
+            shutil.copy2(img_path, imgs_output_path)
+        if parse_depths:
+            depths_output_path = os.path.join(save_path_imgs_camera, "depth")
+            os.makedirs(depths_output_path, exist_ok=True)
+            for depth_path in camera_depth_path_selected:
+                shutil.copy2(depth_path, depths_output_path)
+        Ks = []
+        Ts = []
+        Ks_output_path = os.path.join(save_path_imgs_camera, "intrinsic")
+        Ts_output_path = os.path.join(save_path_imgs_camera, "pose")
+        os.makedirs(Ks_output_path, exist_ok=True)
+        os.makedirs(Ts_output_path, exist_ok=True)
+        for id, posepath in enumerate(camera_pose_path_selected):
+            with open(posepath, "r") as f:
+                data = json.load(f)
+            k_matrix = np.array(data["camera_k_matrix"])
+            rt_matrix = np.array(data["camera_rt_matrix"])
+            T_matrix = np.eye(4)
+            T_matrix[:3, :] = rt_matrix
+            Ks.append(k_matrix)
+            Ts.append(np.linalg.inv(T_matrix))
+            np.save(os.path.join(Ks_output_path, f"{id}.npy"), k_matrix)
+            np.save(os.path.join(Ts_output_path, f"{id}.npy"), T_matrix)
+        Ks = np.array(Ks)
+        Ts = np.array(Ts)
+        correspondenceSave(
+            o3d_room_mesh,
+            save_path_imgs_camera,
+            (image_shape[1], image_shape[0]),
+            room_coords,
+            Ks,
+            Ts,
+        )
+
+    if align_angle:
+        angle = (2 - angle / 180) * np.pi
+        rot_cos, rot_sin = np.cos(angle), np.sin(angle)
+        rot_t = np.array([[rot_cos, -rot_sin, 0], [rot_sin, rot_cos, 0], [0, 0, 1]])
+        room_center = (np.max(room_coords, axis=0) + np.min(room_coords, axis=0)) / 2
+        room_coords = (room_coords - room_center) @ np.transpose(rot_t) + room_center
+        if parse_normal:
+            room_normals = room_normals @ np.transpose(rot_t)
+
+    if parse_pointclouds:
+        os.makedirs(output_root_pointclouds, exist_ok=True)
+        save_path_pointclouds = os.path.join(output_root_pointclouds, room)
+        os.makedirs(save_path_pointclouds, exist_ok=True)
+        room_semantic_gt = np.ascontiguousarray(np.vstack(room_semantic_gt))
+        room_instance_gt = np.ascontiguousarray(np.vstack(room_instance_gt))
+        np.save(
+            os.path.join(save_path_pointclouds, "coord.npy"),
+            room_coords.astype(np.float32),
+        )
+        np.save(
+            os.path.join(save_path_pointclouds, "color.npy"),
+            room_colors.astype(np.uint8),
+        )
+        np.save(
+            os.path.join(save_path_pointclouds, "segment.npy"),
+            room_semantic_gt.astype(np.int16),
+        )
+        np.save(
+            os.path.join(save_path_pointclouds, "instance.npy"),
+            room_instance_gt.astype(np.int16),
+        )
+
+        if parse_normal:
+            np.save(
+                os.path.join(save_path_pointclouds, "normal.npy"),
+                room_normals.astype(np.float32),
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--splits",
+        required=True,
+        nargs="+",
+        choices=["Area_1", "Area_2", "Area_3", "Area_4", "Area_5", "Area_6"],
+        help="Splits need to process ([Area_1, Area_2, Area_3, Area_4, Area_5, Area_6]).",
+    )
+    parser.add_argument(
+        "--dataset_root", required=True, help="Path to Stanford3dDataset_v1.2 dataset"
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where area folders will be located",
+    )
+    parser.add_argument(
+        "--pointclouds_root",
+        default="data/s3dis",
+        type=str,
+        help="Input path where previous pointclouds folder located",
+    )
+    parser.add_argument(
+        "--raw_root",
+        default=None,
+        help="Path to Stanford2d3dDataset_noXYZ dataset (optional)",
+    )
+    parser.add_argument(
+        "--rgb_gap",
+        default=5,
+        help="gap of rgb (optional)",
+    )
+    parser.add_argument(
+        "--align_angle", action="store_true", help="Whether align room angles"
+    )
+    parser.add_argument(
+        "--parse_normal", action="store_true", help="Whether process normal"
+    )
+    parser.add_argument(
+        "--parse_pointclouds", action="store_true", help="Whether parse point clouds"
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    parser.add_argument(
+        "--num_workers", default=1, type=int, help="Num workers for preprocessing."
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="Thread id for parallel processing",
+    )
+    args = parser.parse_args()
+
+    if args.parse_normal:
+        assert args.raw_root is not None
+
+    room_list = []
+    angle_list = []
+
+    # Load room information
+    print("Loading room information ...")
+    for split in args.splits:
+        area_info = np.loadtxt(
+            os.path.join(
+                args.dataset_root,
+                split,
+                f"{split}_alignmentAngle.txt",
+            ),
+            dtype=str,
+        )
+        room_list += [os.path.join(split, room_info[0]) for room_info in area_info]
+        angle_list += [int(room_info[1]) for room_info in area_info]
+
+    if args.parse_normal:
+        # load raw mesh file to extract normal
+        print("Loading raw mesh file ...")
+        for split in args.splits:
+            if split != "Area_5":
+                mesh_dir = os.path.join(args.raw_root, split.lower(), "3d", "rgb.obj")
+                mesh = open3d.io.read_triangle_mesh(mesh_dir)
+                mesh.triangle_uvs.clear()
+                # trimesh_mesh = trimesh.load(mesh_dir, process=False)
+            else:
+                mesh_a_dir = os.path.join(
+                    args.raw_root, f"{split.lower()}a", "3d", "rgb.obj"
+                )
+                mesh_b_dir = os.path.join(
+                    args.raw_root, f"{split.lower()}b", "3d", "rgb.obj"
+                )
+                mesh_a = open3d.io.read_triangle_mesh(mesh_a_dir)
+                mesh_a.triangle_uvs.clear()
+                # trimesh_mesh_a = trimesh.load(mesh_a_dir, process=False)
+                mesh_b = open3d.io.read_triangle_mesh(mesh_b_dir)
+                mesh_b.triangle_uvs.clear()
+                # trimesh_mesh_b = trimesh.load(mesh_b_dir, process=False)
+                mesh_b = mesh_b.transform(
+                    np.array(
+                        [
+                            [0, 0, -1, -4.09703582],
+                            [0, 1, 0, 0],
+                            [1, 0, 0, -6.22617759],
+                            [0, 0, 0, 1],
+                        ]
+                    )
+                )
+                mesh = mesh_a + mesh_b
+                # trimesh_mesh = trimesh_mesh_a +trimesh_mesh_b
+            area_mesh_dict[split] = mesh
+            # trimesh_area_mesh_dict[split] = trimesh_mesh
+            print(f"{split} mesh is loaded")
+
+    room_list_list = np.array_split(room_list, args.num_workers)
+    room_list_ = room_list_list[args.thread_id]
+
+    for i in range(len(room_list_)):
+        parse_room(
+            room_list_[i],
+            angle_list[i],
+            args.rgb_gap,
+            args.raw_root,
+            args.dataset_root,
+            args.output_root,
+            args.align_angle,
+            args.parse_normal,
+            args.parse_pointclouds,
+            args.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/s3dis/preprocess_s3dis.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=6  
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:r:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    r) raw_root=$OPTARG ;;  
+    p) parse_depths=true ;; 
+    c) parse_pointclouds=true ;; 
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> -r <raw_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ] || [ -z "$raw_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> -r <raw_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 1 $((num_workers))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/s3dis/preprocess_s3dis.py --splits Area_$i \
+    --num_workers 1 \
+    --thread_id 0 \
+    --dataset_root $dataset_root \
+    --output_root $output_root \
+    --raw_root $raw_root \
+    --parse_normal"
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/s3dis/splits.py

@@ -0,0 +1,74 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["Area_1", "Area_2", "Area_3", "Area_4", "Area_5", "Area_6"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(dataset_path, "data/s3dis")
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            cam_split_name_path = os.path.join(im_split_path, name)
+            cam_names = [f.name for f in os.scandir(cam_split_name_path) if f.is_dir()]
+            for cam_name in cam_names:
+                im_split_name_path = os.path.join(
+                    cam_split_name_path, cam_name, "color"
+                )
+                co_split_name_path = os.path.join(
+                    cam_split_name_path, cam_name, "correspondence"
+                )
+                png_files = [
+                    f for f in os.listdir(im_split_name_path) if f.endswith(".png")
+                ]
+                png_files = sorted(
+                    png_files,
+                    key=lambda x: int(
+                        x.split(".")[0].split("_frame_")[1].split("_domain_")[0]
+                    ),
+                )
+                # Get the full paths of the .png files
+                png_file_paths = [
+                    os.path.join(im_split_name_path, f).replace(
+                        dataset_path, "data/s3dis"
+                    )
+                    for f in png_files
+                ]
+                co_file_paths = [
+                    os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace(
+                        dataset_path, "data/s3dis"
+                    )
+                    for f in png_files
+                ]
+                for i in range(0, len(png_file_paths), 4):
+                    split_dict[f"{name}_{cam_name}_{i//4}"] = {}
+                    split_dict[f"{name}_{cam_name}_{i//4}"]["pointclouds"] = (
+                        os.path.join(pc_split_path, name)
+                    )
+                    split_dict[f"{name}_{cam_name}_{i//4}"]["images"] = png_file_paths[
+                        i : i + 4
+                    ]
+                    split_dict[f"{name}_{cam_name}_{i//4}"]["correspondences"] = (
+                        co_file_paths[i : i + 4]
+                    )
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the S3DIS dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/scannet/preprocess_scannet.py

@@ -0,0 +1,574 @@
+"""
+Preprocessing Script for ScanNet 20/200
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import warnings
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+import os
+import camtools as ct
+import open3d as o3d
+from scipy.spatial import cKDTree
+import struct
+import zlib
+import imageio
+import cv2
+import argparse
+import glob
+import json
+import plyfile
+import numpy as np
+import pandas as pd
+import multiprocessing as mp
+from pathlib import Path
+
+# Load external constants
+import sys
+
+sys.path.append("pointcept/datasets/preprocessing/scannet/meta_data")
+from scannet200_constants import VALID_CLASS_IDS_200, VALID_CLASS_IDS_20
+
+CLOUD_FILE_PFIX = "_vh_clean_2"
+SEGMENTS_FILE_PFIX = ".0.010000.segs.json"
+AGGREGATIONS_FILE_PFIX = ".aggregation.json"
+CLASS_IDS200 = VALID_CLASS_IDS_200
+CLASS_IDS20 = VALID_CLASS_IDS_20
+IGNORE_INDEX = -1
+
+COMPRESSION_TYPE_COLOR = {-1: "unknown", 0: "raw", 1: "png", 2: "jpeg"}
+COMPRESSION_TYPE_DEPTH = {
+    -1: "unknown",
+    0: "raw_ushort",
+    1: "zlib_ushort",
+    2: "occi_ushort",
+}
+
+
+class RGBDFrame:
+    def load(self, file_handle):
+        self.camera_to_world = np.asarray(
+            struct.unpack("f" * 16, file_handle.read(16 * 4)), dtype=np.float32
+        ).reshape(4, 4)
+        self.timestamp_color = struct.unpack("Q", file_handle.read(8))[0]
+        self.timestamp_depth = struct.unpack("Q", file_handle.read(8))[0]
+        self.color_size_bytes = struct.unpack("Q", file_handle.read(8))[0]
+        self.depth_size_bytes = struct.unpack("Q", file_handle.read(8))[0]
+        self.color_data = b"".join(
+            struct.unpack(
+                "c" * self.color_size_bytes, file_handle.read(self.color_size_bytes)
+            )
+        )
+        self.depth_data = b"".join(
+            struct.unpack(
+                "c" * self.depth_size_bytes, file_handle.read(self.depth_size_bytes)
+            )
+        )
+
+    def decompress_depth(self, compression_type):
+        if compression_type == "zlib_ushort":
+            return self.decompress_depth_zlib()
+        else:
+            raise
+
+    def decompress_depth_zlib(self):
+        return zlib.decompress(self.depth_data)
+
+    def decompress_color(self, compression_type):
+        if compression_type == "jpeg":
+            return self.decompress_color_jpeg()
+        else:
+            raise
+
+    def decompress_color_jpeg(self):
+        return imageio.imread(self.color_data)
+
+
+class SensorData:
+    def __init__(self, filename):
+        self.version = 4
+        self.load(filename)
+
+    def load(self, filename):
+        with open(filename, "rb") as f:
+            version = struct.unpack("I", f.read(4))[0]
+            assert self.version == version
+            strlen = struct.unpack("Q", f.read(8))[0]
+            self.sensor_name = b"".join(struct.unpack("c" * strlen, f.read(strlen)))
+            self.intrinsic_color = np.asarray(
+                struct.unpack("f" * 16, f.read(16 * 4)), dtype=np.float32
+            ).reshape(4, 4)
+            self.extrinsic_color = np.asarray(
+                struct.unpack("f" * 16, f.read(16 * 4)), dtype=np.float32
+            ).reshape(4, 4)
+            self.intrinsic_depth = np.asarray(
+                struct.unpack("f" * 16, f.read(16 * 4)), dtype=np.float32
+            ).reshape(4, 4)
+            self.extrinsic_depth = np.asarray(
+                struct.unpack("f" * 16, f.read(16 * 4)), dtype=np.float32
+            ).reshape(4, 4)
+            self.color_compression_type = COMPRESSION_TYPE_COLOR[
+                struct.unpack("i", f.read(4))[0]
+            ]
+            self.depth_compression_type = COMPRESSION_TYPE_DEPTH[
+                struct.unpack("i", f.read(4))[0]
+            ]
+            self.color_width = struct.unpack("I", f.read(4))[0]
+            self.color_height = struct.unpack("I", f.read(4))[0]
+            self.depth_width = struct.unpack("I", f.read(4))[0]
+            self.depth_height = struct.unpack("I", f.read(4))[0]
+            self.depth_shift = struct.unpack("f", f.read(4))[0]
+            num_frames = struct.unpack("Q", f.read(8))[0]
+            self.frames = []
+            for i in range(num_frames):
+                frame = RGBDFrame()
+                frame.load(f)
+                self.frames.append(frame)
+
+    def export_depth_images(self, output_path, image_size=None, frame_skip=1):
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+        print(
+            "exporting", len(self.frames) // frame_skip, " depth frames to", output_path
+        )
+        for f in range(0, len(self.frames), frame_skip):
+            if os.path.exists((os.path.join(output_path, str(f) + ".png"))):
+                continue
+            if f % 100 == 0:
+                print(
+                    "exporting",
+                    f,
+                    "th depth frames to",
+                    os.path.join(output_path, str(f) + ".png"),
+                )
+
+            depth_data = self.frames[f].decompress_depth(self.depth_compression_type)
+            depth = np.fromstring(depth_data, dtype=np.uint16).reshape(
+                self.depth_height, self.depth_width
+            )
+            if image_size is not None:
+                depth = cv2.resize(
+                    depth,
+                    (image_size[1], image_size[0]),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            imageio.imwrite(os.path.join(output_path, str(f) + ".png"), depth)
+
+    def export_color_images(self, output_path, image_size=None, frame_skip=1):
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+        print(
+            "exporting", len(self.frames) // frame_skip, "color frames to", output_path
+        )
+        for f in range(0, len(self.frames), frame_skip):
+            if os.path.exists((os.path.join(output_path, str(f) + ".png"))):
+                continue
+            if f % 100 == 0:
+                print(
+                    "exporting",
+                    f,
+                    "th color frames to",
+                    os.path.join(output_path, str(f) + ".png"),
+                )
+            color = self.frames[f].decompress_color(self.color_compression_type)
+            if image_size is not None:
+                color = cv2.resize(
+                    color,
+                    (image_size[1], image_size[0]),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            # imageio.imwrite(os.path.join(output_path, str(f) + '.jpg'), color)
+            imageio.imwrite(os.path.join(output_path, str(f) + ".png"), color)
+
+    def save_mat_to_file(self, matrix, filename):
+        with open(filename, "w") as f:
+            for line in matrix:
+                np.savetxt(f, line[np.newaxis], fmt="%f")
+
+    def export_poses(self, output_path, frame_skip=1):
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+        print(
+            "exporting", len(self.frames) // frame_skip, "camera poses to", output_path
+        )
+        for f in range(0, len(self.frames), frame_skip):
+            np.save(
+                os.path.join(output_path, str(f) + ".npy"),
+                self.frames[f].camera_to_world,
+            )
+
+    def export_intrinsics(self, output_path):
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+        print("exporting camera intrinsics to", output_path)
+        np.save(os.path.join(output_path, "intrinsic.npy"), self.intrinsic_color)
+
+
+def reader(
+    filename,
+    output_path,
+    frame_skip,
+    export_color_images=False,
+    export_depth_images=False,
+    export_poses=False,
+    export_intrinsics=False,
+):
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+
+    # load the data
+    print("loading %s..." % filename)
+    sd = SensorData(filename)
+    if export_depth_images:
+        sd.export_depth_images(
+            os.path.join(output_path, "depth"), frame_skip=frame_skip
+        )
+    if export_color_images:
+        sd.export_color_images(
+            os.path.join(output_path, "color"), frame_skip=frame_skip
+        )
+    if export_poses:
+        sd.export_poses(os.path.join(output_path, "pose"), frame_skip=frame_skip)
+    if export_intrinsics:
+        sd.export_intrinsics(os.path.join(output_path, "intrinsic"))
+    return sd.color_height, sd.color_width
+
+
+def read_plymesh(filepath):
+    """Read ply file and return it as numpy array. Returns None if emtpy."""
+    with open(filepath, "rb") as f:
+        plydata = plyfile.PlyData.read(f)
+    if plydata.elements:
+        vertices = pd.DataFrame(plydata["vertex"].data).values
+        faces = np.stack(plydata["face"].data["vertex_indices"], axis=0)
+        return vertices, faces
+
+
+# Map the raw category id to the point cloud
+def point_indices_from_group(seg_indices, group, labels_pd):
+    group_segments = np.array(group["segments"])
+    label = group["label"]
+
+    # Map the category name to id
+    label_id20 = labels_pd[labels_pd["raw_category"] == label]["nyu40id"]
+    label_id20 = int(label_id20.iloc[0]) if len(label_id20) > 0 else 0
+    label_id200 = labels_pd[labels_pd["raw_category"] == label]["id"]
+    label_id200 = int(label_id200.iloc[0]) if len(label_id200) > 0 else 0
+
+    # Only store for the valid categories
+    if label_id20 in CLASS_IDS20:
+        label_id20 = CLASS_IDS20.index(label_id20)
+    else:
+        label_id20 = IGNORE_INDEX
+
+    if label_id200 in CLASS_IDS200:
+        label_id200 = CLASS_IDS200.index(label_id200)
+    else:
+        label_id200 = IGNORE_INDEX
+
+    # get points, where segment indices (points labelled with segment ids) are in the group segment list
+    point_idx = np.where(np.isin(seg_indices, group_segments))[0]
+    return point_idx, label_id20, label_id200
+
+
+def face_normal(vertex, face):
+    v01 = vertex[face[:, 1]] - vertex[face[:, 0]]
+    v02 = vertex[face[:, 2]] - vertex[face[:, 0]]
+    vec = np.cross(v01, v02)
+    length = np.sqrt(np.sum(vec**2, axis=1, keepdims=True)) + 1.0e-8
+    nf = vec / length
+    area = length * 0.5
+    return nf, area
+
+
+def vertex_normal(vertex, face):
+    nf, area = face_normal(vertex, face)
+    nf = nf * area
+
+    nv = np.zeros_like(vertex)
+    for i in range(face.shape[0]):
+        nv[face[i]] += nf[i]
+
+    length = np.sqrt(np.sum(nv**2, axis=1, keepdims=True)) + 1.0e-8
+    nv = nv / length
+    return nv
+
+
+def correspondenceGet(mesh, K, T, img_size, coords_gt):
+    height, width = img_size
+    if np.isnan(T).any() or np.isinf(T).any():
+        return None
+
+    depth = ct.raycast.mesh_to_im_depth(
+        mesh=mesh, K=K, T=np.linalg.inv(T), height=height, width=width
+    )
+
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    valid = ~np.isinf(depth).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict
+
+
+def correspondenceSave(mesh, scene_dir, coords_gt, output_dir, img_size):
+    os.makedirs(output_dir, exist_ok=True)
+    scene_dir = Path(scene_dir)
+    index_gt = [
+        img_name.split(".")[0]
+        for img_name in os.listdir(str(scene_dir / "color"))
+        if img_name.endswith(".png")
+    ]
+    index_gt = sorted(index_gt, key=lambda x: int(x))
+
+    Ks_path = str(scene_dir / "intrinsic" / "intrinsic.npy")
+    Ts_path = str(scene_dir / "pose")
+    Ts_files = sorted(
+        [f for f in os.listdir(Ts_path) if f.endswith(".npy")],
+        key=lambda x: int(x.split(".")[0]),
+    )
+
+    print(f"total pose num:{len(Ts_files)}")
+    Ts = []
+    for Ts_file in Ts_files:
+        file_path = os.path.join(Ts_path, Ts_file)
+        Ts_ = np.load(file_path)
+        Ts.append(Ts_)
+    Ts = np.stack(Ts)
+    Ks = np.load(Ks_path)
+
+    Ks = np.tile(Ks, (Ts.shape[0], 1, 1))
+    Ks = Ks[:, :3, :3]
+    coords_gt_ = coords_gt
+    pixels_ = []
+    coords_ = []
+
+    for i, (K, T) in enumerate(zip(Ks, Ts)):
+        coord_dict = correspondenceGet(mesh, K, T, img_size, coords_gt)
+        if coord_dict is None:
+            correspondences = -np.ones((1, 3))
+        else:
+            pixels_ = coord_dict["pixel"]
+            coords_ = coord_dict["coord"]
+            tree = cKDTree(coords_gt_)
+            dis, idx = tree.query(coords_, k=1)
+            idx_valid = idx[dis < 0.01]
+            pixels_valid = pixels_[dis < 0.01]
+            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+        np.save(Path(output_dir) / "{}.npy".format(index_gt[i]), correspondences)
+
+
+def handle_process(
+    scene_path,
+    output_path,
+    pointclouds_root,
+    labels_pd,
+    train_scenes,
+    val_scenes,
+    frame_gap=75,
+    parse_pointclouds=True,
+    parse_normals=True,
+    export_depth_images=True,
+):
+    pc_output_path = output_path
+    im_output_path = os.path.join(output_path, "images")
+    scene_id = os.path.basename(scene_path)
+    mesh_path = os.path.join(scene_path, f"{scene_id}{CLOUD_FILE_PFIX}.ply")
+    segments_file = os.path.join(
+        scene_path, f"{scene_id}{CLOUD_FILE_PFIX}{SEGMENTS_FILE_PFIX}"
+    )
+    aggregations_file = os.path.join(scene_path, f"{scene_id}{AGGREGATIONS_FILE_PFIX}")
+
+    if scene_id in train_scenes:
+        pc_output_path = os.path.join(pc_output_path, "train", f"{scene_id}")
+        pc_input_path = os.path.join(pointclouds_root, "train", f"{scene_id}")
+        im_output_path = os.path.join(im_output_path, "train", f"{scene_id}")
+        split_name = "train"
+    elif scene_id in val_scenes:
+        pc_output_path = os.path.join(pc_output_path, "val", f"{scene_id}")
+        pc_input_path = os.path.join(pointclouds_root, "val", f"{scene_id}")
+        im_output_path = os.path.join(im_output_path, "val", f"{scene_id}")
+        split_name = "val"
+    else:
+        pc_output_path = os.path.join(pc_output_path, "test", f"{scene_id}")
+        pc_input_path = os.path.join(pointclouds_root, "test", f"{scene_id}")
+        im_output_path = os.path.join(im_output_path, "test", f"{scene_id}")
+        split_name = "test"
+
+    print(f"Processing: {scene_id} in {split_name}")
+
+    if parse_pointclouds:
+        vertices, faces = read_plymesh(mesh_path)
+        coords = vertices[:, :3]
+        colors = vertices[:, 3:6]
+        save_dict = dict(
+            coord=coords.astype(np.float32),
+            color=colors.astype(np.uint8),
+        )
+        # Parse Normals
+        if parse_normals:
+            save_dict["normal"] = vertex_normal(coords, faces).astype(np.float32)
+
+        # Load segments file
+        if split_name != "test":
+            with open(segments_file) as f:
+                segments = json.load(f)
+                seg_indices = np.array(segments["segIndices"])
+
+            # Load Aggregations file
+            with open(aggregations_file) as f:
+                aggregation = json.load(f)
+                seg_groups = np.array(aggregation["segGroups"])
+
+            # Generate new labels
+            semantic_gt20 = np.ones((vertices.shape[0]), dtype=np.int16) * IGNORE_INDEX
+            semantic_gt200 = np.ones((vertices.shape[0]), dtype=np.int16) * IGNORE_INDEX
+            instance_ids = np.ones((vertices.shape[0]), dtype=np.int16) * IGNORE_INDEX
+            for group in seg_groups:
+                point_idx, label_id20, label_id200 = point_indices_from_group(
+                    seg_indices, group, labels_pd
+                )
+
+                semantic_gt20[point_idx] = label_id20
+                semantic_gt200[point_idx] = label_id200
+                instance_ids[point_idx] = group["id"]
+
+            semantic_gt20 = semantic_gt20.astype(int)
+            semantic_gt200 = semantic_gt200.astype(int)
+            instance_ids = instance_ids.astype(int)
+
+            save_dict["segment20"] = semantic_gt20
+            save_dict["segment200"] = semantic_gt200
+            save_dict["instance"] = instance_ids
+
+            # Concatenate with original cloud
+            processed_vertices = np.hstack((semantic_gt200, instance_ids))
+
+            if np.any(np.isnan(processed_vertices)) or not np.all(
+                np.isfinite(processed_vertices)
+            ):
+                raise ValueError(f"Find NaN in Scene: {scene_id}")
+
+        # Save pointcloud data
+        os.makedirs(pc_output_path, exist_ok=True)
+        for key in save_dict.keys():
+            np.save(os.path.join(pc_output_path, f"{key}.npy"), save_dict[key])
+    else:
+        coords = np.load(os.path.join(pc_input_path, "coord.npy"))
+        save_dict = dict(
+            coord=coords.astype(np.float32),
+        )
+
+    # Save img data
+    os.makedirs(im_output_path, exist_ok=True)
+    sens_dir = os.path.join(scene_path, scene_id + ".sens")
+    print(f"Parsing sens data{sens_dir}")
+    h, w = reader(
+        sens_dir,
+        im_output_path,
+        frame_gap,
+        export_color_images=True,
+        export_depth_images=export_depth_images,
+        export_poses=True,
+        export_intrinsics=True,
+    )
+    mesh = o3d.io.read_triangle_mesh(mesh_path)
+    correspondenceSave(
+        mesh,
+        im_output_path,
+        save_dict["coord"],
+        os.path.join(im_output_path, "correspondence"),
+        (h, w),
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--pointclouds_root",
+        default="data/scannet",
+        type=str,
+        help="Input path where previous pointclouds folder located",
+    )
+    parser.add_argument(
+        "--frame_gap", default=75, type=int, help="Frame gap for processing"
+    )
+    parser.add_argument(
+        "--parse_pointclouds", action="store_true", help="Whether parse point clouds"
+    )
+    parser.add_argument(
+        "--parse_normals", action="store_true", help="Whether parse point normals"
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="Thread id for parallel processing",
+    )
+    config = parser.parse_args()
+    meta_root = Path("pointcept/datasets/preprocessing/scannet") / "meta_data"
+
+    # Load label map
+    labels_pd = pd.read_csv(
+        meta_root / "scannetv2-labels.combined.tsv",
+        sep="\t",
+        header=0,
+    )
+
+    # Load train/val splits
+    with open(meta_root / "scannetv2_train.txt") as train_file:
+        train_scenes = train_file.read().splitlines()
+    with open(meta_root / "scannetv2_val.txt") as val_file:
+        val_scenes = val_file.read().splitlines()
+
+    # Load scene paths
+    scene_paths = sorted(glob.glob(config.dataset_root + "/scans*/scene*"))
+    scene_paths_list = np.array_split(scene_paths, config.num_workers)
+    scene_paths_ = scene_paths_list[config.thread_id]
+    # Preprocess data.
+    print("Processing scenes...")
+    for scene_paths_i in scene_paths_:
+        handle_process(
+            scene_paths_i,
+            config.output_root,
+            config.pointclouds_root,
+            labels_pd,
+            train_scenes,
+            val_scenes,
+            config.frame_gap,
+            config.parse_pointclouds,
+            config.parse_normals,
+            config.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/scannet/preprocess_scannet.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=16  
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    p) parse_depths=true ;;
+    c) parse_pointclouds=true ;;
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/scannet/preprocess_scannet.py --thread_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root \
+    --parse_normal"
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/scannet/splits.py

@@ -0,0 +1,62 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train", "val", "test"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(
+            dataset_path, "data/scannet"
+        )
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            im_split_name_path = os.path.join(im_split_path, name, "color")
+            co_split_name_path = os.path.join(im_split_path, name, "correspondence")
+            png_files = [
+                f for f in os.listdir(im_split_name_path) if f.endswith(".png")
+            ]
+            png_files = sorted(png_files, key=lambda x: int(x.split(".")[0]))
+            # Get the full paths of the .png files
+            png_file_paths = [
+                os.path.join(im_split_name_path, f).replace(
+                    dataset_path, "data/scannet"
+                )
+                for f in png_files
+            ]
+            co_file_paths = [
+                os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace(
+                    dataset_path, "data/scannet"
+                )
+                for f in png_files
+            ]
+            for i in range(0, len(png_file_paths), 4):
+                split_dict[f"{name}_{i//4}"] = {}
+                split_dict[f"{name}_{i//4}"]["pointclouds"] = os.path.join(
+                    pc_split_path, name
+                )
+                split_dict[f"{name}_{i//4}"]["images"] = png_file_paths[i : i + 4]
+                split_dict[f"{name}_{i//4}"]["correspondences"] = co_file_paths[
+                    i : i + 4
+                ]
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/scannetpp/preprocess_scannetpp.py

@@ -0,0 +1,557 @@
+"""
+Preprocessing Script for ScanNet 20/200
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import warnings
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+import argparse
+import json
+import torch
+import shutil
+import os
+import torch
+import camtools as ct
+import open3d as o3d
+from scipy.spatial import cKDTree
+import cv2
+import numpy as np
+import pandas as pd
+import multiprocessing as mp
+from collections import OrderedDict
+from scipy.spatial.transform import Rotation
+from pathlib import Path
+import re
+
+try:
+    import pointseg
+except:
+    # Pointseg is located in libs/pointseg
+    warnings.warn("Pointseg is not installed, superpoint segmentation will be skipped.")
+    pointseg = None
+
+REGEXPR_DSLR = re.compile(r".*DSC(?P<frameid>\d+).JPG$")
+REGEXPR_IPHONE = re.compile(r"frame_(?P<frameid>\d+).jpg$")
+
+
+def pose_from_qwxyz_txyz(elems):
+    qw, qx, qy, qz, tx, ty, tz = map(float, elems)
+    pose = np.eye(4)
+    pose[:3, :3] = Rotation.from_quat((qx, qy, qz, qw)).as_matrix()
+    pose[:3, 3] = (tx, ty, tz)
+    c2w = np.linalg.inv(pose)  # returns cam2world
+    return torch.tensor(c2w, dtype=torch.float32)
+
+
+def get_frame_number(name, cam_type="dslr"):
+    if cam_type == "dslr":
+        regex_expr = REGEXPR_DSLR
+    elif cam_type == "iphone":
+        regex_expr = REGEXPR_IPHONE
+    else:
+        raise NotImplementedError(f"wrong {cam_type=} for get_frame_number")
+    matches = re.match(regex_expr, name)
+    return matches["frameid"]
+
+
+def colmap_to_opencv_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] -= 0.5
+    K[1, 2] -= 0.5
+    return K
+
+
+def opencv_to_colmap_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] += 0.5
+    K[1, 2] += 0.5
+    return K
+
+
+def load_sfm(sfm_dir, cam_type="dslr"):
+    # load cameras
+    with open(os.path.join(sfm_dir, "cameras.txt"), "r") as f:
+        raw = f.read().splitlines()[3:]  # skip header
+
+    intrinsics = {}
+    for camera in raw:
+        camera = camera.split(" ")
+        intrinsics[int(camera[0])] = [camera[1]] + [float(cam) for cam in camera[2:]]
+        ins = intrinsics[int(camera[0])]
+    w, h, ins = undistort_ins(ins)
+    # load images
+    with open(os.path.join(sfm_dir, "images.txt"), "r") as f:
+        raw = f.read().splitlines()
+        raw = [line for line in raw if not line.startswith("#")]  # skip header
+
+    img_idx = {}
+    img_infos = {}
+    for image, points in zip(raw[0::2], raw[1::2]):
+        image = image.split(" ")
+        points = points.split(" ")
+
+        idx = image[0]
+        img_name = image[-1]
+        assert img_name not in img_idx, "duplicate db image: " + img_name
+        img_idx[img_name] = idx  # register image name
+
+        # current_points2D = {int(i): (float(x), float(y))
+        #                     for i, x, y in zip(points[2::3], points[0::3], points[1::3]) if i != '-1'}
+        img_infos[idx] = dict(
+            intrinsics=ins,
+            path=img_name,
+            frame_id=get_frame_number(img_name, cam_type),
+            cam_to_world=pose_from_qwxyz_txyz(image[1:-2]),
+            width=w,
+            height=h,
+        )
+
+    return img_idx, img_infos
+
+
+def undistort_ins(intrinsics):
+    camera_type = intrinsics[0]
+
+    width = int(intrinsics[1])
+    height = int(intrinsics[2])
+    fx = intrinsics[3]
+    fy = intrinsics[4]
+    cx = intrinsics[5]
+    cy = intrinsics[6]
+    distortion = np.array(intrinsics[7:])
+
+    K = np.zeros([3, 3])
+    K[0, 0] = fx
+    K[0, 2] = cx
+    K[1, 1] = fy
+    K[1, 2] = cy
+    K[2, 2] = 1
+
+    K = colmap_to_opencv_intrinsics(K)
+    if camera_type == "OPENCV_FISHEYE":
+        assert len(distortion) == 4
+
+        new_K = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
+            K,
+            distortion,
+            (width, height),
+            np.eye(3),
+            balance=0.0,
+        )
+        # Make the cx and cy to be the center of the image
+        new_K[0, 2] = width / 2.0
+        new_K[1, 2] = height / 2.0
+
+        map1, map2 = cv2.fisheye.initUndistortRectifyMap(
+            K, distortion, np.eye(3), new_K, (width, height), cv2.CV_32FC1
+        )
+    else:
+        new_K, _ = cv2.getOptimalNewCameraMatrix(
+            K, distortion, (width, height), 1, (width, height), True
+        )
+        map1, map2 = cv2.initUndistortRectifyMap(
+            K, distortion, np.eye(3), new_K, (width, height), cv2.CV_32FC1
+        )
+
+    # undistorted_image = cv2.remap(rgb, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+    # undistorted_mask = cv2.remap(mask, map1, map2, interpolation=cv2.INTER_LINEAR,
+    #                              borderMode=cv2.BORDER_CONSTANT, borderValue=255)
+    new_K = opencv_to_colmap_intrinsics(new_K)
+    new_K = torch.tensor(new_K, dtype=torch.float32)
+    return width, height, new_K
+
+
+def correspondenceGet(mesh, K, T, img_size, coords_gt):
+    height, width = img_size
+    if np.isnan(T).any() or np.isinf(T).any():
+        return None
+    depth = ct.raycast.mesh_to_im_depth(
+        mesh=mesh, K=K, T=np.linalg.inv(T), height=height, width=width
+    )
+    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+    pixel = pixel.reshape((-1, 2))
+    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+    depth = depth.reshape((-1, 1))
+    valid = ~np.isinf(depth).squeeze(-1)
+    coord = np.zeros_like(pixel, dtype=np.float32)
+    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
+    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row
+    pixel = pixel[valid]
+    coord = coord[valid]
+    if coord.shape[0] == 0:
+        return None
+    pixel = pixel[:, :2]
+    coord_dict = {"pixel": pixel, "coord": coord}
+    return coord_dict, depth
+
+
+def correspondenceSave(
+    mesh, scene_dir, coords_gt, Ks, Ts, output_dir, img_size, parse_depths
+):
+    os.makedirs(output_dir / "correspondence", exist_ok=True)
+    scene_dir = Path(scene_dir)
+    index_gt = [
+        img_name.split(".")[0]
+        for img_name in os.listdir(str(scene_dir / "color"))
+        if img_name.endswith(".JPG")
+    ]
+    index_gt = sorted(index_gt, key=lambda x: int(x.split("DSC")[1]))
+
+    Ks = Ks[:, :3, :3]
+    coords_gt_ = coords_gt
+    pixels_ = []
+    coords_ = []
+
+    for i, (K, T) in enumerate(zip(Ks, Ts)):
+        try:
+            coord_dict, depth = correspondenceGet(mesh, K, T, img_size, coords_gt)
+        except:
+            print(output_dir)
+        if coord_dict is None:
+            correspondences = -np.ones((1, 3))
+        else:
+            pixels_ = coord_dict["pixel"]
+            coords_ = coord_dict["coord"]
+            tree = cKDTree(coords_gt_)
+            dis, idx = tree.query(coords_, k=1)
+            idx_valid = idx[dis < 0.01]
+            pixels_valid = pixels_[dis < 0.01]
+            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
+        np.save(
+            output_dir / "correspondence" / "{}.npy".format(index_gt[i]),
+            correspondences,
+        )
+        if parse_depths:
+            os.makedirs(output_dir / "depth", exist_ok=True)
+            depth = depth * 1000
+            depth = np.where(np.isinf(depth), 65535, depth).astype(np.uint16)
+            depth = depth.reshape(img_size)
+            cv2.imwrite(output_dir / "depth" / "{}.png".format(index_gt[i]), depth)
+
+
+def handle_process(
+    name,
+    split,
+    dataset_path,
+    output_path,
+    pointclouds_root,
+    label_mapping,
+    class2idx,
+    ignore_index=-1,
+    frame_gap=20,
+    parse_pointclouds=False,
+    parse_depths=False,
+):
+    pc_output_path = output_path / split
+    im_output_path = output_path / "images" / split
+    scene_path = dataset_path / "data" / name
+    mesh_path = scene_path / "scans" / "mesh_aligned_0.05.ply"
+    segs_path = scene_path / "scans" / "segments.json"
+    anno_path = scene_path / "scans" / "segments_anno.json"
+    # load mesh vertices and colors
+    mesh = o3d.io.read_triangle_mesh(str(mesh_path))
+
+    if parse_pointclouds:
+        # extract mesh information
+        mesh.compute_vertex_normals(normalized=True)
+        coord = np.array(mesh.vertices).astype(np.float32)
+        color = (np.array(mesh.vertex_colors) * 255).astype(np.uint8)
+        normal = np.array(mesh.vertex_normals).astype(np.float32)
+        # extract superpoint information
+        if pointseg is not None:
+            vertices_sp = torch.from_numpy(np.array(mesh.vertices).astype(np.float32))
+            faces_sp = torch.from_numpy(np.array(mesh.triangles).astype(np.int64))
+            superpoint = pointseg.segment_mesh(vertices_sp, faces_sp).numpy()
+        else:
+            superpoint = None
+
+        # Save pointcloud data
+        os.makedirs(pc_output_path / name, exist_ok=True)
+        np.save(pc_output_path / name / "coord.npy", coord)
+        np.save(pc_output_path / name / "color.npy", color)
+        np.save(pc_output_path / name / "normal.npy", normal)
+        if superpoint is not None:
+            np.save(pc_output_path / "superpoint.npy", superpoint)
+
+        if split != "test":
+            # get label on vertices
+            # load segments = vertices per segment ID
+            with open(segs_path) as f:
+                segments = json.load(f)
+            # load anno = (instance, groups of segments)
+            with open(anno_path) as f:
+                anno = json.load(f)
+            seg_indices = np.array(segments["segIndices"], dtype=np.uint32)
+            num_vertices = len(seg_indices)
+            assert num_vertices == len(coord)
+            semantic_gt = np.ones((num_vertices, 3), dtype=np.int16) * ignore_index
+            instance_gt = np.ones((num_vertices, 3), dtype=np.int16) * ignore_index
+
+            # number of labels are used per vertex. initially 0
+            # increment each time a new label is added
+            instance_size = np.ones((num_vertices, 3), dtype=np.int16) * np.inf
+
+            # keep track of the size of the instance (#vertices) assigned to each vertex
+            # later, keep the label of the smallest instance for major label of vertices
+            # store inf initially so that we can pick the smallest instance
+            labels_used = np.zeros(num_vertices, dtype=np.int16)
+
+            for idx, instance in enumerate(anno["segGroups"]):
+                label = instance["label"]
+                instance["label_orig"] = label
+                # remap label
+                instance["label"] = label_mapping.get(label, None)
+                instance["label_index"] = class2idx.get(instance["label"], ignore_index)
+
+                if instance["label_index"] == ignore_index:
+                    continue
+                # get all the vertices with segment index in this instance
+                # and max number of labels not yet applied
+                mask = np.isin(seg_indices, instance["segments"]) & (labels_used < 3)
+                size = mask.sum()
+                if size == 0:
+                    continue
+
+                # get the position to add the label - 0, 1, 2
+                label_position = labels_used[mask]
+                semantic_gt[mask, label_position] = instance["label_index"]
+                # store all valid instance (include ignored instance)
+                instance_gt[mask, label_position] = instance["objectId"]
+                instance_size[mask, label_position] = size
+                labels_used[mask] += 1
+
+            # major label is the label of smallest instance for each vertex
+            # use major label for single class segmentation
+            # shift major label to the first column
+            mask = labels_used > 1
+            if mask.sum() > 0:
+                major_label_position = np.argmin(instance_size[mask], axis=1)
+
+                major_semantic_label = semantic_gt[mask, major_label_position]
+                semantic_gt[mask, major_label_position] = semantic_gt[:, 0][mask]
+                semantic_gt[:, 0][mask] = major_semantic_label
+
+                major_instance_label = instance_gt[mask, major_label_position]
+                instance_gt[mask, major_label_position] = instance_gt[:, 0][mask]
+                instance_gt[:, 0][mask] = major_instance_label
+
+            np.save(pc_output_path / name / "segment.npy", semantic_gt)
+            np.save(pc_output_path / name / "instance.npy", instance_gt)
+    else:
+        coord = np.load(pointclouds_root / split / name / "coord.npy")
+    rgb_path = scene_path / "dslr" / "undistorted_images"
+    colmap_path = scene_path / "dslr" / "colmap"
+    img_idxs, img_infos = load_sfm(str(colmap_path), cam_type="dslr")
+    rgb_path = sorted(
+        rgb_path.glob("*.JPG"), key=lambda x: int(str(x)[:-4].split("DSC")[1])
+    )[::frame_gap]
+    rgb_path = [str(path) for path in rgb_path]
+    color_output_path = im_output_path / name / "color"
+    os.makedirs(color_output_path, exist_ok=True)
+    for rgb_path_i in rgb_path:
+        shutil.copy(rgb_path_i, color_output_path)
+    imgnames = [path.split("/")[-1] for path in rgb_path]
+    img_idxs = [img_idxs[imgname] for imgname in imgnames]
+    img_infos = [img_infos[img_idx] for img_idx in img_idxs]
+    intrinsic_save_path = im_output_path / name / "intrinsic"
+    pose_save_path = im_output_path / name / "pose"
+    os.makedirs(intrinsic_save_path, exist_ok=True)
+    os.makedirs(pose_save_path, exist_ok=True)
+    Ks = []
+    Ts = []
+    for id, img_info in enumerate(img_infos):
+        np.save(intrinsic_save_path / f"{id}.npy", img_info["intrinsics"])
+        np.save(pose_save_path / f"{id}.npy", img_info["cam_to_world"])
+        Ks.append(img_info["intrinsics"])
+        Ts.append(img_info["cam_to_world"])
+    Ks = np.stack(Ks)
+    Ts = np.stack(Ts)
+
+    if len(rgb_path) == 0 or len(img_infos) == 0:
+        return
+    h = img_infos[0]["height"]
+    w = img_infos[0]["width"]
+    correspondenceSave(
+        mesh,
+        im_output_path / name,
+        coord,
+        Ks,
+        Ts,
+        im_output_path / name,
+        (h, w),
+        parse_depths,
+    )
+
+
+def filter_map_classes(mapping, mapping_type):
+    if mapping_type == "semantic":
+        map_key = "semantic_map_to"
+    elif mapping_type == "instance":
+        map_key = "instance_map_to"
+    else:
+        raise NotImplementedError
+    # create a dict with classes to be mapped
+    # classes that don't have mapping are entered as x->x
+    # otherwise x->y
+    map_dict = OrderedDict()
+
+    for i in range(mapping.shape[0]):
+        row = mapping.iloc[i]
+        class_name = row["class"]
+        map_target = row[map_key]
+
+        # map to None or some other label -> don't add this class to the label list
+        try:
+            if len(map_target) > 0:
+                # map to None -> don't use this class
+                if map_target == "None":
+                    pass
+                else:
+                    # map to something else -> use this class
+                    map_dict[class_name] = map_target
+        except TypeError:
+            # nan values -> no mapping, keep label as is
+            if class_name not in map_dict:
+                map_dict[class_name] = class_name
+
+    return map_dict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located",
+    )
+    parser.add_argument(
+        "--pointclouds_root",
+        default="data/scannetpp",
+        type=str,
+        help="Input path where previous pointclouds folder located",
+    )
+    parser.add_argument(
+        "--ignore_index",
+        default=-1,
+        type=int,
+        help="Default ignore index.",
+    )
+    parser.add_argument(
+        "--frame_gap", default=20, type=int, help="Frame gap for processing"
+    )
+    parser.add_argument(
+        "--parse_pointclouds", action="store_true", help="Whether parse point clouds"
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="Thread id for parallel processing",
+    )
+    config = parser.parse_args()
+
+    print("Loading meta data...")
+    config.dataset_root = Path(config.dataset_root)
+    config.output_root = Path(config.output_root)
+    config.pointclouds_root = Path(config.pointclouds_root)
+
+    train_list = np.loadtxt(
+        config.dataset_root / "splits" / "nvs_sem_train.txt",
+        dtype=str,
+    )
+    print("Num samples in training split:", len(train_list))
+
+    val_list = np.loadtxt(
+        config.dataset_root / "splits" / "nvs_sem_val.txt",
+        dtype=str,
+    )
+    print("Num samples in validation split:", len(val_list))
+
+    test_list = np.loadtxt(
+        config.dataset_root / "splits" / "sem_test.txt",
+        dtype=str,
+    )
+    print("Num samples in testing split:", len(test_list))
+
+    data_list = np.concatenate([train_list, val_list, test_list])
+    split_list = np.concatenate(
+        [
+            np.full_like(train_list, "train"),
+            np.full_like(val_list, "val"),
+            np.full_like(test_list, "test"),
+        ]
+    )
+
+    # Parsing label information and mapping
+    segment_class_names = np.loadtxt(
+        config.dataset_root / "metadata" / "semantic_benchmark" / "top100.txt",
+        dtype=str,
+        delimiter=".",  # dummy delimiter to replace " "
+    )
+    print("Num classes in segment class list:", len(segment_class_names))
+
+    instance_class_names = np.loadtxt(
+        config.dataset_root / "metadata" / "semantic_benchmark" / "top100_instance.txt",
+        dtype=str,
+        delimiter=".",  # dummy delimiter to replace " "
+    )
+    print("Num classes in instance class list:", len(instance_class_names))
+
+    label_mapping = pd.read_csv(
+        config.dataset_root / "metadata" / "semantic_benchmark" / "map_benchmark.csv"
+    )
+    label_mapping = filter_map_classes(label_mapping, mapping_type="semantic")
+    class2idx = {
+        class_name: idx for (idx, class_name) in enumerate(segment_class_names)
+    }
+    # Load scene paths
+
+    data_list_list = np.array_split(data_list, config.num_workers)
+    data_list_ = data_list_list[config.thread_id]
+    split_list_list = np.array_split(split_list, config.num_workers)
+    split_list_ = split_list_list[config.thread_id]
+    # Preprocess data.
+    print("Processing scenes...")
+    for data_list_i, split_list_i in zip(data_list_, split_list_):
+        handle_process(
+            data_list_i,
+            split_list_i,
+            config.dataset_root,
+            config.output_root,
+            config.pointclouds_root,
+            label_mapping,
+            class2idx,
+            config.ignore_index,
+            config.frame_gap,
+            config.parse_pointclouds,
+            config.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/scannetpp/preprocess_scannetpp.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=16  
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    p) parse_depths=true ;;
+    c) parse_pointclouds=true ;;
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/scannetpp/preprocess_scannetpp.py --thread_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root"
+
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/scannetpp/splits.py

@@ -0,0 +1,64 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train", "val", "test"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(
+            dataset_path, "data/scannetpp"
+        )
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            im_split_name_path = os.path.join(im_split_path, name, "color")
+            co_split_name_path = os.path.join(im_split_path, name, "correspondence")
+            png_files = [
+                f for f in os.listdir(im_split_name_path) if f.endswith(".JPG")
+            ]
+            png_files = sorted(
+                png_files, key=lambda x: int(x.split(".")[0].split("DSC")[-1])
+            )
+            # Get the full paths of the .png files
+            png_file_paths = [
+                os.path.join(im_split_name_path, f).replace(
+                    dataset_path, "data/scannetpp"
+                )
+                for f in png_files
+            ]
+            co_file_paths = [
+                os.path.join(co_split_name_path, f.replace(".JPG", ".npy")).replace(
+                    dataset_path, "data/scannetpp"
+                )
+                for f in png_files
+            ]
+            for i in range(0, len(png_file_paths), 4):
+                split_dict[f"{name}_{i//4}"] = {}
+                split_dict[f"{name}_{i//4}"]["pointclouds"] = os.path.join(
+                    pc_split_path, name
+                )
+                split_dict[f"{name}_{i//4}"]["images"] = png_file_paths[i : i + 4]
+                split_dict[f"{name}_{i//4}"]["correspondences"] = co_file_paths[
+                    i : i + 4
+                ]
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the ScanNet dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)

pointcept/datasets/preprocessing/concerto/structured3d/preprocess_structured3d.py

@@ -0,0 +1,557 @@
+"""
+Preprocessing Script for Structured3D
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import argparse
+import io
+import os
+from PIL import Image
+import cv2
+import zipfile
+import numpy as np
+import multiprocessing as mp
+
+VALID_CLASS_IDS_25 = (
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    11,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    22,
+    24,
+    25,
+    32,
+    34,
+    35,
+    38,
+    39,
+    40,
+)
+CLASS_LABELS_25 = (
+    "wall",
+    "floor",
+    "cabinet",
+    "bed",
+    "chair",
+    "sofa",
+    "table",
+    "door",
+    "window",
+    "picture",
+    "desk",
+    "shelves",
+    "curtain",
+    "dresser",
+    "pillow",
+    "mirror",
+    "ceiling",
+    "refrigerator",
+    "television",
+    "nightstand",
+    "sink",
+    "lamp",
+    "otherstructure",
+    "otherfurniture",
+    "otherprop",
+)
+
+
+def normal_from_cross_product(points_2d: np.ndarray) -> np.ndarray:
+    xyz_points_pad = np.pad(points_2d, ((0, 1), (0, 1), (0, 0)), mode="symmetric")
+    xyz_points_ver = (xyz_points_pad[:, :-1, :] - xyz_points_pad[:, 1:, :])[:-1, :, :]
+    xyz_points_hor = (xyz_points_pad[:-1, :, :] - xyz_points_pad[1:, :, :])[:, :-1, :]
+    xyz_normal = np.cross(xyz_points_hor, xyz_points_ver)
+    xyz_dist = np.linalg.norm(xyz_normal, axis=-1, keepdims=True)
+    xyz_normal = np.divide(
+        xyz_normal, xyz_dist, out=np.zeros_like(xyz_normal), where=xyz_dist != 0
+    )
+    return xyz_normal
+
+
+class Structured3DReader:
+    def __init__(self, files):
+        super().__init__()
+        if isinstance(files, str):
+            files = [files]
+        self.readers = [zipfile.ZipFile(f, "r") for f in files]
+        self.names_mapper = dict()
+        for idx, reader in enumerate(self.readers):
+            for name in reader.namelist():
+                self.names_mapper[name] = idx
+
+    def filelist(self):
+        return list(self.names_mapper.keys())
+
+    def listdir(self, dir_name):
+        dir_name = dir_name.lstrip(os.path.sep).rstrip(os.path.sep)
+        file_list = list(
+            np.unique(
+                [
+                    f.replace(dir_name + os.path.sep, "", 1).split(os.path.sep)[0]
+                    for f in self.filelist()
+                    if f.startswith(dir_name + os.path.sep)
+                ]
+            )
+        )
+        if "" in file_list:
+            file_list.remove("")
+        return file_list
+
+    def read(self, file_name):
+        split = self.names_mapper[file_name]
+        return self.readers[split].read(file_name)
+
+    def read_camera(self, camera_path):
+        z2y_top_m = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=np.float32)
+        cam_extr = np.fromstring(self.read(camera_path), dtype=np.float32, sep=" ")
+        cam_t = np.matmul(z2y_top_m, cam_extr[:3] / 1000)
+        if cam_extr.shape[0] > 3:
+            cam_front, cam_up = cam_extr[3:6], cam_extr[6:9]
+            cam_n = np.cross(cam_front, cam_up)
+            cam_r = np.stack((cam_front, cam_up, cam_n), axis=1).astype(np.float32)
+            cam_r = np.matmul(z2y_top_m, cam_r)
+            cam_f = cam_extr[9:11]
+        else:
+            cam_r = np.eye(3, dtype=np.float32)
+            cam_f = None
+        return cam_r, cam_t, cam_f
+
+    def read_depth(self, depth_path):
+        depth = cv2.imdecode(
+            np.frombuffer(self.read(depth_path), np.uint8), cv2.IMREAD_UNCHANGED
+        )[..., np.newaxis]
+        depth[depth == 0] = 65535
+        return depth
+
+    def read_color(self, color_path):
+        color = cv2.imdecode(
+            np.frombuffer(self.read(color_path), np.uint8), cv2.IMREAD_UNCHANGED
+        )[..., :3][..., ::-1]
+        return color
+
+    def read_segment(self, segment_path):
+        segment = np.array(Image.open(io.BytesIO(self.read(segment_path))))[
+            ..., np.newaxis
+        ]
+        return segment
+
+
+def parse_scene(
+    scene,
+    dataset_root,
+    output_root,
+    ignore_index=-1,
+    grid_size=None,
+    fuse_prsp=True,
+    fuse_pano=True,
+    parse_pointclouds=True,
+    parse_depths=True,
+    vis=False,
+):
+    assert fuse_prsp or fuse_pano
+    pc_output_root = output_root
+    im_output_root = os.path.join(output_root, "images")
+    reader = Structured3DReader(
+        [
+            os.path.join(dataset_root, f)
+            for f in os.listdir(dataset_root)
+            if f.endswith(".zip")
+        ]
+    )
+    scene_id = int(os.path.basename(scene).split("_")[-1])
+    if scene_id < 3000:
+        split = "train"
+    elif 3000 <= scene_id < 3250:
+        split = "val"
+    else:
+        split = "test"
+
+    print(f"Processing: {scene} in {split}")
+
+    rooms = reader.listdir(os.path.join("Structured3D", scene, "2D_rendering"))
+    for room in rooms:
+        im_save_path = os.path.join(
+            im_output_root, split, os.path.basename(scene), f"room_{room}"
+        )
+        pc_save_path = os.path.join(
+            pc_output_root, split, os.path.basename(scene), f"room_{room}"
+        )
+        if os.path.exists(im_save_path) and os.path.exists(pc_save_path):
+            print(f"exist {split}/{os.path.basename(scene)}/room_{room}")
+            continue
+
+        room_path = os.path.join("Structured3D", scene, "2D_rendering", room)
+        coord_list = list()
+        color_list = list()
+        normal_list = list()
+        segment_list = list()
+        prsp_list = list()
+        pano_list = list()
+        prsp_depth_list = list()
+        pano_depth_list = list()
+        prsp_correspondence_list = list()
+        pano_correspondence_list = list()
+        Ks_list = list()
+        Ts_list = list()
+        if fuse_prsp and scene:
+            prsp_path = os.path.join(room_path, "perspective", "full")
+            frames = reader.listdir(prsp_path)
+
+            for frame_id, frame in enumerate(frames):
+                try:
+                    cam_r, cam_t, cam_f = reader.read_camera(
+                        os.path.join(prsp_path, frame, "camera_pose.txt")
+                    )
+                    depth = reader.read_depth(
+                        os.path.join(prsp_path, frame, "depth.png")
+                    )
+                    color = reader.read_color(
+                        os.path.join(prsp_path, frame, "rgb_rawlight.png")
+                    )
+                    segment = reader.read_segment(
+                        os.path.join(prsp_path, frame, "semantic.png")
+                    )
+                except:
+                    print(
+                        f"Skipping {scene}_room{room}_frame{frame} perspective view due to loading error"
+                    )
+                else:
+                    fx, fy = cam_f
+                    height, width = depth.shape[0], depth.shape[1]
+                    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
+                    pixel = pixel.reshape((-1, 2))
+                    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
+                    k = np.diag([1.0, 1.0, 1.0])
+
+                    k[0, 2] = width / 2
+                    k[1, 2] = height / 2
+
+                    k[0, 0] = k[0, 2] / np.tan(fx)
+                    k[1, 1] = k[1, 2] / np.tan(fy)
+                    coord = (
+                        depth.reshape((-1, 1)) * (np.linalg.inv(k) @ pixel.T).T
+                    ).reshape(height, width, 3)
+                    coord = coord @ np.array([[0, 0, 1], [0, -1, 0], [1, 0, 0]])
+                    normal = normal_from_cross_product(coord)
+
+                    # Filtering invalid points
+                    view_dist = np.maximum(
+                        np.linalg.norm(coord, axis=-1, keepdims=True), float(10e-5)
+                    )
+                    cosine_dist = np.sum(
+                        (coord * normal / view_dist), axis=-1, keepdims=True
+                    )
+                    cosine_dist = np.abs(cosine_dist)
+                    mask = ((cosine_dist > 0.15) & (depth < 65535) & (segment > 0))[
+                        ..., 0
+                    ].reshape(-1)
+
+                    coord = np.matmul(coord / 1000, cam_r.T) + cam_t
+                    normal = normal_from_cross_product(coord)
+
+                    T = np.eye(4)
+                    T[:3, :3] = cam_r
+                    T[:3, 3] = cam_t
+
+                    Ks_list.append(k)
+                    Ts_list.append(T)
+
+                    pixel[:, 2] = pixel[:, 2] * frame_id
+                    pixel_valid = pixel[mask]
+
+                    if sum(mask) > 0:
+                        prsp_depth_list.append(depth)
+                        coord_list.append(coord.reshape(-1, 3)[mask])
+                        color_list.append(color.reshape(-1, 3)[mask])
+                        normal_list.append(normal.reshape(-1, 3)[mask])
+                        segment_list.append(segment.reshape(-1, 1)[mask])
+                        prsp_list.append(color)
+                        prsp_correspondence_list.append(pixel_valid)
+                    else:
+                        print(
+                            f"Skipping {scene}_room{room}_frame{frame} perspective view due to all points are filtered out"
+                        )
+
+        if fuse_pano:
+            pano_path = os.path.join(room_path, "panorama")
+            try:
+                _, cam_t, _ = reader.read_camera(
+                    os.path.join(pano_path, "camera_xyz.txt")
+                )
+                depth = reader.read_depth(os.path.join(pano_path, "full", "depth.png"))
+                color = reader.read_color(
+                    os.path.join(pano_path, "full", "rgb_rawlight.png")
+                )
+                segment = reader.read_segment(
+                    os.path.join(pano_path, "full", "semantic.png")
+                )
+            except:
+                print(f"Skipping {scene}_room{room} panorama view due to loading error")
+            else:
+                p_h, p_w = depth.shape[:2]
+                pixel = np.transpose(np.indices((p_w, p_h)), (2, 1, 0))
+                pixel = pixel.reshape((-1, 2))
+                p_a = np.arange(p_w, dtype=np.float32) / p_w * 2 * np.pi - np.pi
+                p_b = np.arange(p_h, dtype=np.float32) / p_h * np.pi * -1 + np.pi / 2
+                p_a = np.tile(p_a[None], [p_h, 1])[..., np.newaxis]
+                p_b = np.tile(p_b[:, None], [1, p_w])[..., np.newaxis]
+                p_a_sin, p_a_cos, p_b_sin, p_b_cos = (
+                    np.sin(p_a),
+                    np.cos(p_a),
+                    np.sin(p_b),
+                    np.cos(p_b),
+                )
+                x = depth * p_a_cos * p_b_cos
+                y = depth * p_b_sin
+                z = depth * p_a_sin * p_b_cos
+                coord = np.concatenate([x, y, z], axis=-1) / 1000
+                normal = normal_from_cross_product(coord)
+
+                # Filtering invalid points
+                view_dist = np.maximum(
+                    np.linalg.norm(coord, axis=-1, keepdims=True), float(10e-5)
+                )
+                cosine_dist = np.sum(
+                    (coord * normal / view_dist), axis=-1, keepdims=True
+                )
+                cosine_dist = np.abs(cosine_dist)
+                mask = ((cosine_dist > 0.15) & (depth < 65535) & (segment > 0))[
+                    ..., 0
+                ].reshape(-1)
+                coord = coord + cam_t
+
+                pixel = np.hstack((pixel, -np.ones((pixel.shape[0], 1))))
+                pixel_valid = pixel[mask]
+
+                if sum(mask) > 0:
+                    pano_depth_list.append(depth)
+                    coord_list.append(coord.reshape(-1, 3)[mask])
+                    color_list.append(color.reshape(-1, 3)[mask])
+                    normal_list.append(normal.reshape(-1, 3)[mask])
+                    segment_list.append(segment.reshape(-1, 1)[mask])
+                    pano_list.append(color)
+                    pano_correspondence_list.append(pixel_valid)
+                else:
+                    print(
+                        f"Skipping {scene}_room{room} panorama view due to all points are filtered out"
+                    )
+        if len(prsp_correspondence_list) > 0 and len(pano_correspondence_list) > 0:
+            coord = np.concatenate(coord_list, axis=0)
+            coord = coord @ np.array([[1, 0, 0], [0, 0, 1], [0, 1, 0]])
+            color = np.concatenate(color_list, axis=0)
+            normal = np.concatenate(normal_list, axis=0)
+            normal = normal @ np.array([[1, 0, 0], [0, 0, 1], [0, 1, 0]])
+            segment = np.concatenate(segment_list, axis=0)
+            segment25 = np.ones_like(segment, dtype=np.int64) * ignore_index
+            Ks_list = np.stack(Ks_list, axis=0)
+            Ts_list = np.stack(Ts_list, axis=0)
+            Ts_list = Ts_list @ np.array(
+                [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]
+            )
+            for idx, value in enumerate(VALID_CLASS_IDS_25):
+                mask = np.all(segment == value, axis=-1)
+                segment25[mask] = idx
+            correspondence = np.concatenate(prsp_correspondence_list, axis=0)
+            correspondence = np.concatenate(
+                [correspondence, pano_correspondence_list[0]], axis=0
+            )
+            data_dict = dict(
+                coord=coord.astype(np.float32),
+                color=color.astype(np.uint8),
+                normal=normal.astype(np.float32),
+                segment=segment25.astype(np.int16),
+                correspondence=correspondence.astype(np.int32),
+            )
+
+            # exclude ignore, wall, floor, ceiling
+            valid = np.sum(~np.isin(data_dict["segment"], [-1, 0, 1, 16]))
+            if valid == 0:
+                print(
+                    f"Skip {scene}_room{room} due to no effective points (exclude wall, floor, ceiling)"
+                )
+                continue
+
+            # Grid sampling data
+            if grid_size is not None:
+                grid_coord = np.floor(coord / grid_size).astype(int)
+                _, idx = np.unique(grid_coord, axis=0, return_index=True)
+                # coord = coord[idx]
+                for key in data_dict.keys():
+                    data_dict[key] = data_dict[key][idx]
+
+            correspondence = data_dict["correspondence"]
+            correspondence = np.concatenate(
+                [correspondence, np.arange(correspondence.shape[0])[:, None]], axis=1
+            )
+            frame_id_list = np.unique(correspondence[:, 2])
+            frame_id_list = frame_id_list[frame_id_list != -1]
+            pano_correspondence_list = [
+                correspondence[correspondence[:, 2] == -1][:, [0, 1, 3]]
+            ]
+            prsp_correspondence_list = [
+                correspondence[correspondence[:, 2] == i][:, [0, 1, 3]]
+                for i in frame_id_list
+            ]
+            if parse_pointclouds:
+                os.makedirs(pc_save_path, exist_ok=True)
+                # Save data
+                for key in data_dict.keys():
+                    if key in ["correspondence"]:
+                        continue
+                    np.save(os.path.join(pc_save_path, f"{key}.npy"), data_dict[key])
+
+            os.makedirs(im_save_path, exist_ok=True)
+            if fuse_prsp:
+                prsp_save_path = os.path.join(im_save_path, "color", "prsp")
+                prsp_correspondence_save_path = os.path.join(
+                    im_save_path, "correspondence", "prsp_correspondence"
+                )
+                Ks_save_path = os.path.join(im_save_path, "intrinsic")
+                Ts_save_path = os.path.join(im_save_path, "pose")
+                os.makedirs(prsp_correspondence_save_path, exist_ok=True)
+                os.makedirs(prsp_save_path, exist_ok=True)
+                os.makedirs(Ks_save_path, exist_ok=True)
+                os.makedirs(Ts_save_path, exist_ok=True)
+                if parse_depths:
+                    prsp_depth_save_path = os.path.join(im_save_path, "depth", "prsp")
+                    os.makedirs(prsp_depth_save_path, exist_ok=True)
+                for idx in range(len(prsp_list)):
+                    prsp_img_path = os.path.join(prsp_save_path, f"{idx}.png")
+                    cv2.imwrite(prsp_img_path, prsp_list[idx][..., ::-1])
+                    if parse_depths:
+                        prsp_depth_img_path = os.path.join(
+                            prsp_depth_save_path, f"{idx}.png"
+                        )
+                        cv2.imwrite(prsp_depth_img_path, prsp_depth_list[idx])
+                    np.save(
+                        os.path.join(prsp_correspondence_save_path, f"{idx}.npy"),
+                        prsp_correspondence_list[idx],
+                    )
+                    np.save(
+                        os.path.join(Ks_save_path, f"{idx}.npy"),
+                        Ks_list[idx],
+                    )
+                    np.save(
+                        os.path.join(Ts_save_path, f"{idx}.npy"),
+                        Ts_list[idx],
+                    )
+
+            if fuse_pano:
+                pano_save_path = os.path.join(im_save_path, "color", "pano")
+                pano_correspondence_save_path = os.path.join(
+                    im_save_path, "correspondence", "pano_correspondence"
+                )
+                os.makedirs(pano_save_path, exist_ok=True)
+                os.makedirs(pano_correspondence_save_path, exist_ok=True)
+                if parse_depths:
+                    pano_depth_save_path = os.path.join(im_save_path, "depth", "pano")
+                    os.makedirs(pano_depth_save_path, exist_ok=True)
+                for idx in range(len(pano_list)):
+                    pano_img_path = os.path.join(pano_save_path, f"{idx}.png")
+                    cv2.imwrite(pano_img_path, pano_list[idx][..., ::-1])
+                    if parse_depths:
+                        pano_depth_img_path = os.path.join(
+                            pano_depth_save_path, f"{idx}.png"
+                        )
+                        cv2.imwrite(pano_depth_img_path, pano_depth_list[idx])
+                for idx, pano_correspondence in enumerate(pano_correspondence_list):
+                    np.save(
+                        os.path.join(pano_correspondence_save_path, f"{idx}.npy"),
+                        pano_correspondence,
+                    )
+
+            if vis:
+                from pointcept.utils.visualization import save_point_cloud
+
+                os.makedirs("./vis", exist_ok=True)
+                save_point_cloud(
+                    coord, color / 255, f"./vis/{scene}_room{room}_color.ply"
+                )
+                save_point_cloud(
+                    coord, (normal + 1) / 2, f"./vis/{scene}_room{room}_normal.ply"
+                )
+        else:
+            print(f"Skipping {scene}_room{room} due to no valid points")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the Structured3D dataset containing scene folders.",
+    )
+    parser.add_argument(
+        "--output_root",
+        required=True,
+        help="Output path where train/val folders will be located.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=mp.cpu_count(),
+        type=int,
+        help="Num workers for preprocessing.",
+    )
+    parser.add_argument(
+        "--thread_id",
+        default=0,
+        type=int,
+        help="thread_id.",
+    )
+    parser.add_argument(
+        "--grid_size", default=None, type=float, help="Grid size for grid sampling."
+    )
+    parser.add_argument("--ignore_index", default=-1, type=float, help="Ignore index.")
+    parser.add_argument(
+        "--fuse_prsp", action="store_true", help="Whether fuse perspective view."
+    )
+    parser.add_argument(
+        "--fuse_pano", action="store_true", help="Whether fuse panorama view."
+    )
+    parser.add_argument(
+        "--parse_pointclouds", action="store_true", help="Whether parse point clouds"
+    )
+    parser.add_argument(
+        "--parse_depths", action="store_true", help="Whether parse depths"
+    )
+    config = parser.parse_args()
+
+    reader = Structured3DReader(
+        [
+            os.path.join(config.dataset_root, f)
+            for f in os.listdir(config.dataset_root)
+            if f.endswith(".zip")
+        ]
+    )
+
+    scenes_list = reader.listdir("Structured3D")
+    scenes_list = sorted(scenes_list)
+    split_scenes_list = np.array_split(scenes_list, config.num_workers)
+    split_scenes_list_ = split_scenes_list[config.thread_id]
+
+    for scenes_list_i in split_scenes_list_:
+        parse_scene(
+            scenes_list_i,
+            config.dataset_root,
+            config.output_root,
+            config.ignore_index,
+            config.grid_size,
+            config.fuse_prsp,
+            config.fuse_pano,
+            config.parse_pointclouds,
+            config.parse_depths,
+        )

pointcept/datasets/preprocessing/concerto/structured3d/preprocess_structured3d.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+dataset_root=""
+output_root=""
+num_workers=16  
+parse_depths=false
+parse_pointclouds=false
+
+while getopts "d:o:n:pc" opt; do
+  case $opt in
+    d) dataset_root=$OPTARG ;;  
+    o) output_root=$OPTARG ;;   
+    n) num_workers=$OPTARG ;;   
+    p) parse_depths=true ;;
+    c) parse_pointclouds=true ;;
+    *) echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"; exit 1 ;;
+  esac
+done
+
+if [ -z "$dataset_root" ] || [ -z "$output_root" ]; then
+    echo "Usage: $0 -d <dataset_root> -o <output_root> [-n <num_workers>] [-p] [-c]"
+    exit 1
+fi
+
+for i in $(seq 0 $((num_workers - 1))); do
+    cmd="python pointcept/datasets/preprocessing/concerto/structured3d/preprocess_structured3d.py --thread_id $i \
+    --num_workers $num_workers \
+    --dataset_root $dataset_root \
+    --output_root $output_root \
+    --fuse_prsp --fuse_pano"
+    
+    if $parse_depths; then
+        cmd="$cmd --parse_depths"
+    fi
+
+    if $parse_pointclouds; then
+        cmd="$cmd --parse_pointclouds"
+    fi
+
+    eval "$cmd &"
+done
+
+wait

pointcept/datasets/preprocessing/concerto/structured3d/splits.py

@@ -0,0 +1,69 @@
+import os
+import json
+import argparse
+
+
+def get_splits_paths(dataset_path):
+    # Get the names of all subfolders in the given folder
+    im_path = os.path.join(dataset_path, "images")
+    pc_path = dataset_path
+    splits = ["train", "val"]
+    split_path = os.path.join(dataset_path, "splits")
+    os.makedirs(split_path, exist_ok=True)
+    for split in splits:
+        im_split_path = os.path.join(im_path, split)
+        pc_split_path = os.path.join(pc_path, split).replace(
+            dataset_path, "data/structured3d"
+        )
+        split_names = [f.name for f in os.scandir(im_split_path) if f.is_dir()]
+        split_dict = {}
+        for name in split_names:
+            cam_split_name_path = os.path.join(im_split_path, name)
+            cam_names = [f.name for f in os.scandir(cam_split_name_path) if f.is_dir()]
+            for cam_name in cam_names:
+                im_split_name_path = os.path.join(
+                    cam_split_name_path, cam_name, "color", "prsp"
+                )
+                co_split_name_path = os.path.join(
+                    cam_split_name_path,
+                    cam_name,
+                    "correspondence",
+                    "prsp_correspondence",
+                )
+                png_files = [
+                    f for f in os.listdir(im_split_name_path) if f.endswith(".png")
+                ]
+                png_files = sorted(png_files, key=lambda x: int(x.split(".")[0]))
+                # Get the full paths of the .png files
+                png_file_paths = [
+                    os.path.join(im_split_name_path, f).replace(
+                        dataset_path, "data/structured3d"
+                    )
+                    for f in png_files
+                ]
+                co_file_paths = [
+                    os.path.join(co_split_name_path, f.replace(".png", ".npy")).replace(
+                        dataset_path, "data/structured3d"
+                    )
+                    for f in png_files
+                ]
+                split_dict[f"{name}_{cam_name}"] = {}
+                split_dict[f"{name}_{cam_name}"]["pointclouds"] = os.path.join(
+                    pc_split_path, name, cam_name
+                )
+                split_dict[f"{name}_{cam_name}"]["images"] = png_file_paths
+                split_dict[f"{name}_{cam_name}"]["correspondences"] = co_file_paths
+        with open(os.path.join(split_path, f"{split}.json"), "w") as f:
+            json.dump(split_dict, f, indent=4)
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset_root",
+        required=True,
+        help="Path to the Structured3D dataset containing scene folders",
+    )
+    config = parser.parse_args()
+    get_splits_paths(config.dataset_root)