code/pointcept_framework/pointcept/datasets/nuscenes.py

"""
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

os.environ["OMP_NUM_THREADS"] = "1"


@DATASETS.register_module()
class NuScenesDataset(DefaultDataset):
    OCCUPANCY_LABEL_MAP = {
        0: -1,
        1: 0,
        2: 1,
        3: 2,
        4: 3,
        5: 4,
        6: 5,
        7: 6,
        8: 7,
        9: 8,
        10: 9,
        11: 10,
        12: 11,
        13: 12,
        14: 13,
        15: 14,
        16: 15,
        17: -1,
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.occ_voxel_size = np.array([0.4, 0.4, 0.4], dtype=np.float32)
        self.occ_origin = np.array([-40.0, -40.0, -1.0], dtype=np.float32)
        self.occ_shape = np.array([200, 200, 16], dtype=np.int32)
        self._scene_of_token = None

    def get_info_path(self, split):
        assert split in ["train", "val", "test"]
        if split == "train":
            return os.path.join(self.data_root, "nuscenes_infos_train.pkl")
        elif split == "val":
            return os.path.join(self.data_root, "nuscenes_infos_val.pkl")
        elif split == "test":
            return os.path.join(self.data_root, "nuscenes_infos_test.pkl")

    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)
                if isinstance(info, dict) and "data_list" in info:
                    info = info["data_list"]
                data_list.extend(info)
        return data_list

    def get_data_name(self, idx):
        return self.data_list[idx % len(self.data_list)]["token"]

    def _build_scene_lookup(self):
        if self._scene_of_token is not None:
            return
        scene_of_token = {}
        gts_root = os.path.join(self.data_root, "gts")
        if os.path.isdir(gts_root):
            for scene_name in os.listdir(gts_root):
                scene_path = os.path.join(gts_root, scene_name)
                if not os.path.isdir(scene_path):
                    continue
                for token in os.listdir(scene_path):
                    token_path = os.path.join(scene_path, token)
                    if os.path.isdir(token_path):
                        scene_of_token[token] = scene_name
        self._scene_of_token = scene_of_token

    def _resolve_lidar_path(self, data):
        lidar_path = data["lidar_points"]["lidar_path"] if "lidar_points" in data else data["lidar_path"]
        if not lidar_path.startswith(("samples/", "sweeps/")) and "LIDAR_TOP" in lidar_path:
            lidar_path = "samples/LIDAR_TOP/" + lidar_path
        full_lidar_path = os.path.join(self.data_root, lidar_path)
        if not os.path.exists(full_lidar_path):
            full_lidar_path = os.path.join(self.data_root, "raw", lidar_path)
        return full_lidar_path

    def _load_points(self, data):
        full_lidar_path = self._resolve_lidar_path(data)
        points = np.fromfile(str(full_lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
        coord = points[:, :3]
        strength = points[:, 3].reshape([-1, 1]) / 255
        return points, coord, strength

    def _load_lidarseg(self, token, point_count):
        lidarseg_path = os.path.join(self.data_root, "lidarseg", "v1.0-trainval", token + "_lidarseg.bin")
        if not os.path.exists(lidarseg_path):
            return None
        segment_data = np.fromfile(lidarseg_path, dtype=np.uint8, count=-1)
        if segment_data.shape[0] != point_count:
            return None
        learning_map = {
            0: self.ignore_index, 1: self.ignore_index, 2: 6, 3: 6, 4: 6, 5: self.ignore_index, 6: 6, 7: self.ignore_index,
            8: self.ignore_index, 9: 0, 10: self.ignore_index, 11: self.ignore_index, 12: 7, 13: self.ignore_index,
            14: 1, 15: 2, 16: 2, 17: 3, 18: 4, 19: self.ignore_index, 20: self.ignore_index, 21: 5, 22: 8, 23: 9,
            24: 10, 25: 11, 26: 12, 27: 13, 28: 14, 29: self.ignore_index, 30: 15, 31: self.ignore_index,
        }
        return np.vectorize(lambda x: learning_map.get(x, self.ignore_index))(segment_data).astype(np.int64)

    def _load_occ_segment(self, data, points):
        self._build_scene_lookup()
        token = data["token"]
        scene_name = self._scene_of_token.get(token)
        if scene_name is None:
            return None
        occ_path = os.path.join(self.data_root, "gts", scene_name, token, "labels.npz")
        if not os.path.exists(occ_path):
            return None

        occ = np.load(occ_path)
        semantics = occ["semantics"]
        mask_lidar = occ["mask_lidar"].astype(bool)

        lidar2ego = np.array(data["lidar_points"]["lidar2ego"], dtype=np.float32)
        pts_h = np.concatenate([points[:, :3], np.ones((points.shape[0], 1), dtype=np.float32)], axis=1)
        ego_xyz = (pts_h @ lidar2ego.T)[:, :3]
        ijk = np.floor((ego_xyz - self.occ_origin) / self.occ_voxel_size).astype(np.int32)
        in_bounds = (
            (ijk[:, 0] >= 0) & (ijk[:, 0] < self.occ_shape[0]) &
            (ijk[:, 1] >= 0) & (ijk[:, 1] < self.occ_shape[1]) &
            (ijk[:, 2] >= 0) & (ijk[:, 2] < self.occ_shape[2])
        )

        segment = np.full((points.shape[0],), self.ignore_index, dtype=np.int64)
        if not np.any(in_bounds):
            return segment

        idx = ijk[in_bounds]
        occ_sem = semantics[idx[:, 0], idx[:, 1], idx[:, 2]]
        occ_vis = mask_lidar[idx[:, 0], idx[:, 1], idx[:, 2]]
        mapped = np.vectorize(lambda x: self.OCCUPANCY_LABEL_MAP.get(int(x), self.ignore_index))(occ_sem).astype(np.int64)
        valid = occ_vis & (mapped != self.ignore_index)
        segment[np.where(in_bounds)[0][valid]] = mapped[valid]
        return segment

    def get_data(self, idx):
        data = self.data_list[idx % len(self.data_list)]
        token = data["token"]
        points, coord, strength = self._load_points(data)

        segment = self._load_lidarseg(token, points.shape[0])
        if segment is None:
            segment = self._load_occ_segment(data, points)
        if segment is None:
            segment = np.full((points.shape[0],), self.ignore_index, dtype=np.int64)

        data_dict = dict(
            coord=coord,
            strength=strength,
            segment=segment,
            name=self.get_data_name(idx),
        )
        return data_dict










@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(object):
    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)
        self.img_ratio = img_num / (6 * sweeps)
        super().__init__(ignore_index=ignore_index, **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

    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]
        )

        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"]
                    # sensor2lidar = cam_lidar_tm @ sensor2lidar
                    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 += 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_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]

        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((points.shape[0],), dtype=np.int64) * self.ignore_index

        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

    @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