supervisely 6.73.273__py3-none-any.whl → 6.73.275__py3-none-any.whl

supervisely/convert/pointcloud_episodes/nuscenes_conv/nuscenes_converter.py

@@ -0,0 +1,305 @@
+from os import path as osp
+from pathlib import Path
+from typing import Dict, Optional
+
+import supervisely.convert.pointcloud_episodes.nuscenes_conv.nuscenes_helper as helpers
+import supervisely.io.fs as fs
+from supervisely._utils import is_development
+from supervisely.annotation.obj_class import ObjClass
+from supervisely.annotation.tag_meta import TagMeta, TagValueType
+from supervisely.api.api import Api, ApiField
+from supervisely.convert.base_converter import AvailablePointcloudConverters
+from supervisely.convert.pointcloud_episodes.pointcloud_episodes_converter import (
+    PointcloudEpisodeConverter,
+)
+from supervisely.geometry.cuboid_3d import Cuboid3d
+from supervisely.pointcloud_annotation.pointcloud_episode_annotation import (
+    PointcloudEpisodeAnnotation,
+)
+from supervisely.pointcloud_annotation.pointcloud_episode_frame import (
+    PointcloudEpisodeFrame,
+)
+from supervisely.pointcloud_annotation.pointcloud_episode_frame_collection import (
+    PointcloudEpisodeFrameCollection,
+)
+from supervisely.pointcloud_annotation.pointcloud_episode_object import (
+    PointcloudEpisodeObject,
+)
+from supervisely.pointcloud_annotation.pointcloud_episode_object_collection import (
+    PointcloudEpisodeObjectCollection,
+)
+from supervisely.pointcloud_annotation.pointcloud_episode_tag_collection import (
+    PointcloudEpisodeTagCollection,
+)
+from supervisely.pointcloud_annotation.pointcloud_figure import PointcloudFigure
+from supervisely.project.project_meta import ProjectMeta
+from supervisely.sly_logger import logger
+from supervisely.tiny_timer import TinyTimer
+
+
+class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
+    """Converter for NuScenes pointcloud episodes format."""
+
+    def __init__(
+        self,
+        input_data: str,
+        labeling_interface: str,
+        upload_as_links: bool,
+        remote_files_map: Optional[Dict[str, str]] = None,
+    ):
+        super().__init__(input_data, labeling_interface, upload_as_links, remote_files_map)
+        self._nuscenes = None
+
+    def __str__(self) -> str:
+        return AvailablePointcloudConverters.NUSCENES
+
+    def validate_format(self) -> bool:
+        try:
+            from nuscenes import NuScenes
+        except ImportError:
+            logger.warning("Please, run 'pip install nuscenes-devkit' to import NuScenes data.")
+            return False
+
+        def filter_fn(path):
+            return all(
+                [
+                    (Path(path) / name).exists()
+                    for name in ["maps", "samples", "sweeps", "v1.0-mini"]
+                ]
+            )
+
+        try:
+            input_path = [d for d in fs.dirs_filter(self._input_data, filter_fn)].pop()
+        except IndexError:
+            return False
+
+        sample_dir = input_path + "/samples/"
+        if any([not fs.dir_exists(f"{sample_dir}/{d}") for d in helpers.DIR_NAMES]):
+            return False
+
+        sweeps_dir = input_path + "/sweeps/"
+        if any([not fs.dir_exists(f"{sweeps_dir}/{d}") for d in helpers.DIR_NAMES]):
+            return False
+
+        ann_dir = input_path + "/v1.0-mini/"
+        if any([not fs.file_exists(f"{ann_dir}/{d}.json") for d in helpers.TABLE_NAMES]):
+            return False
+
+        try:
+            t = TinyTimer()
+            nuscenes = NuScenes(dataroot=input_path, verbose=False)
+            self._nuscenes: NuScenes = nuscenes
+            logger.info(f"NuScenes initialization took {t.get_sec():.3f} sec")
+        except Exception as e:
+            logger.debug(f"Failed to initialize NuScenes: {e}")
+            return False
+
+        return True
+
+    def to_supervisely(
+        self,
+        scene_samples,
+        meta: ProjectMeta,
+        renamed_classes: dict = {},
+        renamed_tags: dict = {},
+    ) -> PointcloudEpisodeAnnotation:
+        token_to_obj = {}
+        frames = []
+        tags = []
+        for sample_i, sample in enumerate(scene_samples):
+            figures = []
+            for obj in sample.anns:
+                instance_token = obj.instance_token
+                class_name = obj.category
+                parent_obj_token = obj.parent_token
+                parent_object = None
+                if parent_obj_token == "":
+                    # * Create a new object
+                    obj_class_name = renamed_classes.get(class_name, class_name)
+                    obj_class = meta.get_obj_class(obj_class_name)
+                    obj_tags = None  # ! TODO: fix tags
+                    pcd_ep_obj = PointcloudEpisodeObject(obj_class, obj_tags)
+                    # * Assign the object to the starting token
+                    token_to_obj[instance_token] = pcd_ep_obj
+                    parent_object = pcd_ep_obj
+                else:
+                    # * -> Figure has a parent object, get it
+                    token_to_obj[instance_token] = token_to_obj[parent_obj_token]
+                    parent_object = token_to_obj[parent_obj_token]
+                geom = obj.to_supervisely()
+                pcd_figure = PointcloudFigure(parent_object, geom, sample_i)
+                figures.append(pcd_figure)
+            frame = PointcloudEpisodeFrame(sample_i, figures)
+            frames.append(frame)
+        tag_collection = PointcloudEpisodeTagCollection(tags) if len(tags) > 0 else None
+        return PointcloudEpisodeAnnotation(
+            len(frames),
+            PointcloudEpisodeObjectCollection(list(set(token_to_obj.values()))),
+            PointcloudEpisodeFrameCollection(frames),
+            tag_collection,
+        )
+
+    def upload_dataset(self, api: Api, dataset_id: int, batch_size: int = 1, log_progress=True):
+        nuscenes = self._nuscenes
+
+        tag_metas = [TagMeta(attr["name"], TagValueType.NONE) for attr in nuscenes.attribute]
+        obj_classes = []
+        for category in nuscenes.category:
+            color = nuscenes.colormap[category["name"]]
+            description = category["description"]
+            if len(description) > 255:
+                # * Trim description to fit into 255 characters limit
+                sentences = description.split(".")
+                trimmed_description = ""
+                for sentence in sentences:
+                    if len(trimmed_description) + len(sentence) + 1 > 255:
+                        break
+                    trimmed_description += sentence + "."
+                description = trimmed_description.strip()
+            obj_classes.append(ObjClass(category["name"], Cuboid3d, color, description=description))
+
+        self._meta = ProjectMeta(obj_classes, tag_metas)
+        meta, renamed_classes, renamed_tags = self.merge_metas_with_conflicts(api, dataset_id)
+
+        dataset_info = api.dataset.get_info_by_id(dataset_id)
+        scene_name_to_dataset = {}
+
+        scene_names = [scene["name"] for scene in nuscenes.scene]
+        scene_cnt = len(scene_names)
+        total_sample_cnt = sum([scene["nbr_samples"] for scene in nuscenes.scene])
+
+        multiple_scenes = len(scene_names) > 1
+        if multiple_scenes:
+            logger.info(f"Found {scene_cnt} scenes ({total_sample_cnt} samples) in the input data.")
+            # * Create a nested dataset for each scene
+            for name in scene_names:
+                ds = api.dataset.create(
+                    dataset_info.project_id,
+                    name,
+                    change_name_if_conflict=True,
+                    parent_id=dataset_id,
+                )
+                scene_name_to_dataset[name] = ds
+        else:
+            scene_name_to_dataset[scene_names[0]] = dataset_info
+
+        if log_progress:
+            progress, progress_cb = self.get_progress(total_sample_cnt, "Converting episode scenes...")
+        else:
+            progress_cb = None
+
+        for scene in nuscenes.scene:
+            current_dataset_id = scene_name_to_dataset[scene["name"]].id
+
+            log = nuscenes.get("log", scene["log_token"])
+            sample_token = scene["first_sample_token"]
+
+            # * Extract scene's samples
+            scene_samples = []
+            for i in range(scene["nbr_samples"]):
+                sample = nuscenes.get("sample", sample_token)
+                lidar_path, boxes, _ = nuscenes.get_sample_data(sample["data"]["LIDAR_TOP"])
+                if not osp.exists(lidar_path):
+                    logger.warning(f'Scene "{scene["name"]}" has no LIDAR data.')
+                    continue
+
+                timestamp = sample["timestamp"]
+                anns = []
+                for box, name, inst_token in helpers.Sample.generate_boxes(nuscenes, boxes):
+                    current_instance_token = inst_token["token"]
+                    parent_token = inst_token["prev"]
+
+                    # get category, attributes and visibility
+                    ann = nuscenes.get("sample_annotation", current_instance_token)
+                    category = ann["category_name"]
+                    attributes = [
+                        nuscenes.get("attribute", attr)["name"] for attr in ann["attribute_tokens"]
+                    ]
+                    visibility = nuscenes.get("visibility", ann["visibility_token"])["level"]
+
+                    anns.append(
+                        helpers.AnnotationObject(
+                            name,
+                            box,
+                            current_instance_token,
+                            parent_token,
+                            category,
+                            attributes,
+                            visibility,
+                        )
+                    )
+
+                # get camera data
+                sample_data = nuscenes.get("sample_data", sample["data"]["LIDAR_TOP"])
+                cal_sensor = nuscenes.get(
+                    "calibrated_sensor", sample_data["calibrated_sensor_token"]
+                )
+                ego_pose = nuscenes.get("ego_pose", sample_data["ego_pose_token"])
+
+                camera_data = [
+                    helpers.CamData(nuscenes, sensor, token, cal_sensor, ego_pose)
+                    for sensor, token in sample["data"].items()
+                    if sensor.startswith("CAM")
+                ]
+                scene_samples.append(helpers.Sample(timestamp, lidar_path, anns, camera_data))
+                sample_token = sample["next"]
+
+            # * Convert and upload pointclouds
+            frame_to_pointcloud_ids = {}
+            for idx, sample in enumerate(scene_samples):
+                pcd_path = sample.convert_lidar_to_supervisely()
+
+                pcd_name = fs.get_file_name(pcd_path)
+                pcd_meta = {
+                    "frame": idx,
+                    "vehicle": log["vehicle"],
+                    "date": log["date_captured"],
+                    "location": log["location"],
+                    "description": scene["description"],
+                }
+                info = api.pointcloud_episode.upload_path(
+                    current_dataset_id, pcd_name, pcd_path, pcd_meta
+                )
+                fs.silent_remove(pcd_path)
+
+                pcd_id = info.id
+                frame_to_pointcloud_ids[idx] = pcd_id
+
+                # * Upload related images
+                image_jsons = []
+                camera_names = []
+                for img_path, rimage_info in [
+                    data.get_info(sample.timestamp) for data in sample.cam_data
+                ]:
+                    img = api.pointcloud_episode.upload_related_image(img_path)
+                    image_jsons.append(
+                        {
+                            ApiField.ENTITY_ID: pcd_id,
+                            ApiField.NAME: rimage_info[ApiField.NAME],
+                            ApiField.HASH: img,
+                            ApiField.META: rimage_info[ApiField.META],
+                        }
+                    )
+                    camera_names.append(rimage_info[ApiField.META]["deviceId"])
+                if len(image_jsons) > 0:
+                    api.pointcloud_episode.add_related_images(image_jsons, camera_names)
+
+                if log_progress:
+                    progress_cb(1)
+
+            # * Convert and upload annotations
+            pcd_ann = self.to_supervisely(scene_samples, meta, renamed_classes, renamed_tags)
+            try:
+                api.pointcloud_episode.annotation.append(
+                    current_dataset_id, pcd_ann, frame_to_pointcloud_ids
+                )
+                logger.info(f"Dataset ID:{current_dataset_id} has been successfully uploaded.")
+            except Exception as e:
+                error_msg = getattr(getattr(e, "response", e), "text", str(e))
+                logger.warning(
+                    f"Failed to upload annotation for scene: {scene['name']}. Message: {error_msg}"
+                )
+
+        if log_progress:
+            if is_development():
+                progress.close()

supervisely/convert/pointcloud_episodes/nuscenes_conv/nuscenes_helper.py

@@ -0,0 +1,265 @@
+from datetime import datetime
+from os import path as osp
+from pathlib import Path
+from typing import List
+
+import numpy as np
+
+from supervisely import fs, logger
+from supervisely.geometry.cuboid_3d import Cuboid3d, Vector3d
+
+DIR_NAMES = [
+    "CAM_BACK",
+    "CAM_BACK_LEFT",
+    "CAM_BACK_RIGHT",
+    "CAM_FRONT",
+    "CAM_FRONT_LEFT",
+    "CAM_FRONT_RIGHT",
+    "LIDAR_TOP",
+    "RADAR_FRONT",
+    "RADAR_FRONT_LEFT",
+    "RADAR_FRONT_RIGHT",
+    "RADAR_BACK_LEFT",
+    "RADAR_BACK_RIGHT",
+]
+
+TABLE_NAMES = [
+    "category",
+    "attribute",
+    "visibility",
+    "instance",
+    "sensor",
+    "calibrated_sensor",
+    "ego_pose",
+    "log",
+    "scene",
+    "sample",
+    "sample_data",
+    "sample_annotation",
+    "map",
+]
+
+
+class Sample:
+    """
+    A class to represent a sample from the NuScenes dataset.
+    """
+
+    def __init__(self, timestamp, lidar_path, anns, cam_data):
+        self.timestamp = datetime.utcfromtimestamp(timestamp / 1e6).isoformat()
+        self.lidar_path = lidar_path
+        self.anns = anns
+        self.cam_data = cam_data
+
+    @staticmethod
+    def generate_boxes(nuscenes, boxes):
+        """
+        Generate ground truth boxes for a given set of boxes.
+
+        Yields:
+            tuple: A tuple containing:
+                - gt_box (np.ndarray): A numpy array representing the ground truth box with concatenated location,
+                                       dimensions, and rotation.
+                - name (str): The name of the object.
+                - instance_token (str): The instance token associated with the box.
+        """
+        locs = np.array([b.center for b in boxes]).reshape(-1, 3)
+        dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)
+        rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes]).reshape(-1, 1)
+
+        gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1)
+        names = np.array([b.name for b in boxes])
+        instance_tokens = [nuscenes.get("sample_annotation", box.token) for box in boxes]
+
+        yield from zip(gt_boxes, names, instance_tokens)
+
+    def convert_lidar_to_supervisely(self):
+        """
+        Converts a LiDAR point cloud file to the Supervisely format and saves it as a .pcd file.
+
+        Returns:
+            str: The file path of the saved .pcd file.
+        """
+        import open3d as o3d  # pylint: disable=import-error
+
+        bin_file = Path(self.lidar_path)
+        save_path = str(bin_file.with_suffix(".pcd"))
+
+        b = np.fromfile(bin_file, dtype=np.float32).reshape(-1, 5)
+        points = b[:, 0:3]
+        intensity = b[:, 3]
+        ring_index = b[:, 4]
+        intensity_fake_rgb = np.zeros((intensity.shape[0], 3))
+        intensity_fake_rgb[:, 0] = (
+            intensity  # red The intensity measures the reflectivity of the objects
+        )
+        intensity_fake_rgb[:, 1] = (
+            ring_index  # green ring index is the index of the laser ranging from 0 to 31
+        )
+        try:
+            pc = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points))
+            pc.colors = o3d.utility.Vector3dVector(intensity_fake_rgb)
+            o3d.io.write_point_cloud(save_path, pc)
+        except Exception as e:
+            logger.warning(f"Error converting lidar to supervisely format: {e}")
+        return save_path
+
+
+class AnnotationObject:
+    """
+    A class to represent an annotation object in the NuScenes dataset.
+
+    Attributes:
+    -----------
+    name : str
+        The name of the annotation object.
+    bbox : np.ndarray
+        The bounding box coordinates.
+    instance_token : str
+        The instance token associated with the annotation object.
+    parent_token : str
+        The token of instance preceding the current object instance.
+    category : str
+        The class name of the annotation object.
+    attributes : List[str]
+        The attribute names associated with the annotation object.
+    visibility : str
+        The visibility level of the annotation object.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        bbox: np.ndarray,
+        instance_token: str,
+        parent_token: str,
+        category: str,
+        attributes: List[str],
+        visibility: str,
+    ):
+        self.name = name
+        self.bbox = bbox
+        self.instance_token = instance_token
+        self.parent_token = parent_token
+
+        self.category = category
+        self.attributes = attributes
+        self.visibility = visibility
+
+    def to_supervisely(self):
+        box = self.convert_nuscenes_to_BEVBox3D()
+
+        bbox = box.to_xyzwhlr()
+        dim = bbox[[3, 5, 4]]
+        pos = bbox[:3] + [0, 0, dim[1] / 2]
+        yaw = bbox[-1]
+
+        position = Vector3d(float(pos[0]), float(pos[1]), float(pos[2]))
+        rotation = Vector3d(0, 0, float(-yaw))
+        dimension = Vector3d(float(dim[0]), float(dim[2]), float(dim[1]))
+        geometry = Cuboid3d(position, rotation, dimension)
+
+        return geometry
+
+    def convert_nuscenes_to_BEVBox3D(self):
+        import open3d as o3d  # pylint: disable=import-error
+
+        box = self.bbox
+        center = [float(box[0]), float(box[1]), float(box[2])]
+        size = [float(box[3]), float(box[5]), float(box[4])]
+        ry = float(box[6])
+        yaw = ry - np.pi
+        yaw = yaw - np.floor(yaw / (2 * np.pi) + 0.5) * 2 * np.pi
+        world_cam = None
+        return o3d.ml.datasets.utils.BEVBox3D(center, size, yaw, self.name, -1.0, world_cam)
+
+
+class CamData:
+    """
+    A class to represent camera data and perform transformations between different coordinate systems.
+
+    Attributes:
+    -----------
+    name : str
+        The name of the sensor.
+    path : str
+        The path to the image file.
+    imsize : tuple
+        The size of the image (width, height).
+    extrinsic : np.ndarray
+        The extrinsic matrix (4x4) representing the transformation from the lidar to the camera coordinate system.
+    intrinsic : np.ndarray
+        The intrinsic matrix (3x3) representing the camera's intrinsic parameters.
+    """
+
+    def __init__(self, nuscenes, sensor_name, sensor_token, cs_record, ego_record):
+        from nuscenes.utils.data_classes import (  # pylint: disable=import-error
+            transform_matrix,
+        )
+        from pyquaternion import Quaternion  # pylint: disable=import-error
+
+        img_path, boxes, cam_intrinsic = nuscenes.get_sample_data(sensor_token)
+        if not osp.exists(img_path):
+            return None
+
+        sd_record_cam = nuscenes.get("sample_data", sensor_token)
+        cs_record_cam = nuscenes.get("calibrated_sensor", sd_record_cam["calibrated_sensor_token"])
+        ego_record_cam = nuscenes.get("ego_pose", sd_record_cam["ego_pose_token"])
+        lid_to_ego = transform_matrix(
+            cs_record["translation"],
+            Quaternion(cs_record["rotation"]),
+            inverse=False,
+        )
+        lid_ego_to_world = transform_matrix(
+            ego_record["translation"],
+            Quaternion(ego_record["rotation"]),
+            inverse=False,
+        )
+        world_to_cam_ego = transform_matrix(
+            ego_record_cam["translation"],
+            Quaternion(ego_record_cam["rotation"]),
+            inverse=True,
+        )
+        ego_to_cam = transform_matrix(
+            cs_record_cam["translation"],
+            Quaternion(cs_record_cam["rotation"]),
+            inverse=True,
+        )
+        velo_to_cam = np.dot(
+            ego_to_cam, np.dot(world_to_cam_ego, np.dot(lid_ego_to_world, lid_to_ego))
+        )
+        velo_to_cam_rot = velo_to_cam[:3, :3]
+        velo_to_cam_trans = velo_to_cam[:3, 3]
+
+        self.name = sensor_name
+        self.path = str(img_path)
+        self.imsize = (sd_record_cam["width"], sd_record_cam["height"])
+        self.extrinsic = np.hstack((velo_to_cam_rot, velo_to_cam_trans.reshape(3, 1)))
+        self.intrinsic = np.asarray(cs_record_cam["camera_intrinsic"])
+
+    def get_info(self, timestamp):
+        """
+        Retrieves information about the image and its metadata.
+
+        Args:
+            timestamp (int): The timestamp associated with the image.
+
+        Returns:
+            tuple: A tuple containing the image path and a dictionary with image metadata.
+        """
+        sensors_to_skip = ["_intrinsic", "_extrinsic", "_imsize"]
+        if not any([self.name.endswith(s) for s in sensors_to_skip]):
+            image_name = fs.get_file_name_with_ext(self.path)
+            sly_path_img = osp.join(osp.dirname(self.path), image_name)
+            img_info = {
+                "name": image_name,
+                "meta": {
+                    "deviceId": self.name,
+                    "timestamp": timestamp,
+                    "sensorsData": {
+                        "extrinsicMatrix": list(self.extrinsic.flatten().astype(float)),
+                        "intrinsicMatrix": list(self.intrinsic.flatten().astype(float)),
+                    },
+                },
+            }
+            return (sly_path_img, img_info)