supervisely 6.73.410__py3-none-any.whl → 6.73.470__py3-none-any.whl

supervisely/convert/pointcloud_episodes/nuscenes_conv/nuscenes_converter.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import os
+import uuid
 from os import path as osp
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import supervisely.convert.pointcloud_episodes.nuscenes_conv.nuscenes_helper as helpers
 import supervisely.io.fs as fs
@@ -9,6 +12,7 @@ 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.api.dataset_api import DatasetInfo
 from supervisely.convert.base_converter import AvailablePointcloudConverters
 from supervisely.convert.pointcloud_episodes.pointcloud_episodes_converter import (
     PointcloudEpisodeConverter,
@@ -36,20 +40,14 @@ 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
+from supervisely.video_annotation.key_id_map import KeyIdMap
 
 
 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
+    _nuscenes: "NuScenes" = None  # type: ignore
+    _custom_data: Dict = {}
 
     def __str__(self) -> str:
         return AvailablePointcloudConverters.NUSCENES
@@ -61,8 +59,10 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
             logger.warning("Please, run 'pip install nuscenes-devkit' to import NuScenes data.")
             return False
 
-        def _contains_tables(p):
-            return all(fs.file_exists(Path(p) / f"{name}.json") for name in helpers.TABLE_NAMES)
+        table_json_filenames = [f"{name}.json" for name in helpers.TABLE_NAMES]
+
+        def _contains_tables(dir_path: str) -> bool:
+            return all(fs.file_exists(osp.join(dir_path, table)) for table in table_json_filenames)
 
         def _filter_fn(path):
             has_tables = False
@@ -83,18 +83,19 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
         version = osp.basename(ann_dir)
         try:
             t = TinyTimer()
-            nuscenes = NuScenes(version=version, dataroot=input_path, verbose=False)
-            self._nuscenes: NuScenes = nuscenes
+            self._nuscenes: NuScenes = NuScenes(version=version, dataroot=input_path, verbose=False)
             logger.debug(f"NuScenes initialization took {t.get_sec():.3f} sec")
         except Exception as e:
             logger.debug(f"Failed to initialize NuScenes: {e}")
             return False
 
+        self._custom_data["nuscenes_version"] = version
+        self._custom_data["dataroot"] = input_path
         return True
 
     def to_supervisely(
         self,
-        scene_samples,
+        scene_samples: Dict[str, helpers.Sample],
         meta: ProjectMeta,
         renamed_classes: dict = {},
         renamed_tags: dict = {},
@@ -102,9 +103,13 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
         token_to_obj = {}
         frames = []
         tags = []
-        for sample_i, sample in enumerate(scene_samples):
+        frame_idx_to_scene_sample_token = {}
+        if "frame_token_map" not in self._custom_data:
+            self._custom_data["frame_token_map"] = {}
+        for sample_i, (token, sample) in enumerate(scene_samples.items()):
             figures = []
             for obj in sample.anns:
+                ann_token = uuid.UUID(obj.token)
                 instance_token = obj.instance_token
                 class_name = obj.category
                 parent_obj_token = obj.parent_token
@@ -114,7 +119,9 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
                     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)
+                    pcd_ep_obj = PointcloudEpisodeObject(
+                        obj_class, obj_tags, uuid.UUID(instance_token)
+                    )
                     # * Assign the object to the starting token
                     token_to_obj[instance_token] = pcd_ep_obj
                     parent_object = pcd_ep_obj
@@ -123,42 +130,45 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
                     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)
+                pcd_figure = PointcloudFigure(parent_object, geom, sample_i, ann_token)
                 figures.append(pcd_figure)
+                frame_idx_to_scene_sample_token[sample_i] = token
             frame = PointcloudEpisodeFrame(sample_i, figures)
             frames.append(frame)
         tag_collection = PointcloudEpisodeTagCollection(tags) if len(tags) > 0 else None
+        self._custom_data["frame_token_map"][self._current_ds_id] = frame_idx_to_scene_sample_token
+        key_uuid = uuid.UUID(token)
         return PointcloudEpisodeAnnotation(
             len(frames),
             PointcloudEpisodeObjectCollection(list(set(token_to_obj.values()))),
             PointcloudEpisodeFrameCollection(frames),
             tag_collection,
+            key=key_uuid,
         )
 
     def upload_dataset(self, api: Api, dataset_id: int, batch_size: int = 1, log_progress=True):
-        nuscenes = self._nuscenes
+        from nuscenes import NuScenes  # pylint: disable=import-error
+
+        nuscenes: NuScenes = self._nuscenes
+        key_id_map = KeyIdMap()
 
         tag_metas = [TagMeta(attr["name"], TagValueType.NONE) for attr in nuscenes.attribute]
-        obj_classes = []
+        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)
+            description = helpers.trim_description(category.get("description", ""))
+            token = category["token"]
+            obj_classes[token] = ObjClass(
+                category["name"], Cuboid3d, color, description=description
+            )
+
+        self._custom_data["classes_token_map"] = {k: v.name for k, v in obj_classes.items()}
+
+        self._meta = ProjectMeta(list(obj_classes.values()), 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_name_to_dataset: Dict[str, DatasetInfo] = {}
 
         scene_names = [scene["name"] for scene in nuscenes.scene]
         scene_cnt = len(scene_names)
@@ -188,12 +198,13 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
 
         for scene in nuscenes.scene:
             current_dataset_id = scene_name_to_dataset[scene["name"]].id
+            self._current_ds_id = current_dataset_id
 
             log = nuscenes.get("log", scene["log_token"])
             sample_token = scene["first_sample_token"]
 
             # * Extract scene's samples
-            scene_samples = []
+            scene_samples: Dict[str, helpers.Sample] = {}
             for i in range(scene["nbr_samples"]):
                 sample = nuscenes.get("sample", sample_token)
                 lidar_path, boxes, _ = nuscenes.get_sample_data(sample["data"]["LIDAR_TOP"])
@@ -207,25 +218,25 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
                     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,
-                        )
+                    ann_token = ann["token"]
+
+                    ann = helpers.AnnotationObject(
+                        name=name,
+                        bbox=box,
+                        token=ann_token,
+                        instance_token=current_instance_token,
+                        parent_token=parent_token,
+                        category=category,
+                        attributes=attributes,
+                        visibility=visibility,
                     )
+                    anns.append(ann)
 
                 # get camera data
                 sample_data = nuscenes.get("sample_data", sample["data"]["LIDAR_TOP"])
@@ -239,15 +250,18 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
                     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["token"]
+                scene_samples[sample_token] = 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):
+            for idx, sample in enumerate(scene_samples.values()):
                 pcd_path = sample.convert_lidar_to_supervisely()
 
-                pcd_name = fs.get_file_name(pcd_path)
+                pcd_name = fs.get_file_name_with_ext(pcd_path)
                 pcd_meta = {
                     "frame": idx,
                     "vehicle": log["vehicle"],
@@ -287,9 +301,10 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
 
             # * 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
+                    current_dataset_id, pcd_ann, frame_to_pointcloud_ids, key_id_map=key_id_map
                 )
                 logger.info(f"Dataset ID:{current_dataset_id} has been successfully uploaded.")
             except Exception as e:
@@ -297,6 +312,15 @@ class NuscenesEpisodesConverter(PointcloudEpisodeConverter):
                 logger.warning(
                     f"Failed to upload annotation for scene: {scene['name']}. Message: {error_msg}"
                 )
+        key_id_map = key_id_map.to_dict()
+        key_id_map.pop("tags")
+        key_id_map.pop("videos")
+        self._custom_data["key_id_map"] = key_id_map
+
+        project_id = dataset_info.project_id
+        current_custom_data = api.project.get_custom_data(project_id)
+        current_custom_data.update(self._custom_data)
+        api.project.update_custom_data(project_id, current_custom_data)
 
         if log_progress:
             if is_development():

supervisely/convert/pointcloud_episodes/nuscenes_conv/nuscenes_helper.py

@@ -1,7 +1,8 @@
+from dataclasses import dataclass, field
 from datetime import datetime
 from os import path as osp
 from pathlib import Path
-from typing import List
+from typing import Dict, Generator, List, Tuple
 
 import numpy as np
 
@@ -40,113 +41,50 @@ TABLE_NAMES = [
 ]
 
 
-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
+def trim_description(description: str, max_length: int = 255) -> str:
+    if len(description) > max_length:
+        sentences = description.split(".")
+        trimmed_description = ""
+        for sentence in sentences:
+            if len(trimmed_description) + len(sentence) + 1 > max_length:
+                break
+            trimmed_description += sentence + "."
+        description = trimmed_description.strip()
+    return description
 
 
+@dataclass
 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.
+    :param name: The name of the annotation object
+    :type name: str
+    :param bbox: The bounding box coordinates in NuScenes format
+    :type bbox: np.ndarray
+    :param token: The unique token identifying the annotation object
+    :type token: str
+    :param instance_token: The instance token associated with the annotation object
+    :type instance_token: str
+    :param parent_token: The token of instance preceding the current object instance
+    :type parent_token: str
+    :param category: The class name of the annotation object
+    :type category: str
+    :param attributes: The attribute names associated with the annotation object
+    :type attributes: List[str]
+    :param visibility: The visibility level of the annotation object
+    :type visibility: str
     """
-
-    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):
+    name: str
+    bbox: np.ndarray
+    token: str
+    instance_token: str
+    parent_token: str
+    category: str
+    attributes: List[str]
+    visibility: str
+
+    def to_supervisely(self) -> Cuboid3d:
         box = self.convert_nuscenes_to_BEVBox3D()
 
         bbox = box.to_xyzwhlr()
@@ -176,29 +114,33 @@ class AnnotationObject:
 
 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.
+    This class handles camera sensor data from the nuScenes dataset, including coordinate system
+    transformations from lidar to camera space and extraction of camera calibration parameters.
+
+    :param nuscenes: The nuScenes dataset instance
+    :type nuscenes: NuScenes
+    :param sensor_name: The name of the camera sensor
+    :type sensor_name: str
+    :param sensor_token: The token identifying the specific sensor sample
+    :type sensor_token: str
+    :param cs_record: The calibrated sensor record for the lidar
+    :type cs_record: dict
+    :param ego_record: The ego pose record for the lidar
+    :type ego_record: dict
     """
 
-    def __init__(self, nuscenes, sensor_name, sensor_token, cs_record, ego_record):
+    def __init__(
+        self, nuscenes, sensor_name: str, sensor_token: str, cs_record: dict, ego_record: dict
+    ):
+        from nuscenes import NuScenes  # pylint: disable=import-error
         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)
+        nuscenes: NuScenes = nuscenes
+
+        img_path, _, _ = nuscenes.get_sample_data(sensor_token)
         if not osp.exists(img_path):
             return None
 
@@ -237,15 +179,14 @@ class CamData:
         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):
+    def get_info(self, timestamp: str) -> Tuple[str, Dict]:
         """
-        Retrieves information about the image and its metadata.
-
-        Args:
-            timestamp (int): The timestamp associated with the image.
+        Generates image info based on the camera data.
 
-        Returns:
-            tuple: A tuple containing the image path and a dictionary with image metadata.
+        :param timestamp: The timestamp associated with the image
+        :type timestamp: str
+        :return: A tuple containing the image path and a dictionary with image metadata.
+        :rtype: tuple
         """
         sensors_to_skip = ["_intrinsic", "_extrinsic", "_imsize"]
         if not any([self.name.endswith(s) for s in sensors_to_skip]):
@@ -263,3 +204,75 @@ class CamData:
                 },
             }
             return (sly_path_img, img_info)
+
+
+@dataclass
+class Sample:
+    """
+    A class to represent a sample from the NuScenes dataset.
+    """
+    timestamp_us: float
+    lidar_path: str
+    anns: List[AnnotationObject]
+    cam_data: List[CamData]
+    timestamp: str = field(init=False)
+
+    def __post_init__(self):
+        self.timestamp = datetime.utcfromtimestamp(self.timestamp_us / 1e6).isoformat()
+
+    @staticmethod
+    def generate_boxes(nuscenes, boxes: List) -> Generator:
+        """
+        Generate ground truth boxes for a given set of boxes.
+
+        :param nuscenes: The nuScenes dataset instance
+        :type nuscenes: NuScenes
+        :param boxes: A list of boxes to generate ground truth for
+        :type boxes: List
+        :return: A generator that yields tuples containing the ground truth box, name, and instance token.
+        :rtype: generator
+        """
+        from nuscenes.utils.data_classes import Box  # pylint: disable=import-error
+
+        boxes: List[Box] = boxes
+
+        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) -> str:
+        """
+        Converts a LiDAR point cloud file to the Supervisely format and saves it as a .pcd file.
+
+        :return: The file path of the saved .pcd file.
+        :rtype: str
+        """
+        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

supervisely/convert/video/video_converter.py

@@ -266,8 +266,8 @@ class VideoConverter(BaseConverter):
             if codec_type not in ["video", "audio"]:
                 continue
             codec_name = stream["codecName"]
-            if codec_type == "video" and codec_name != "h264":
-                logger.info(f"Video codec is not h264, transcoding is required: {codec_name}")
+            if codec_type == "video" and codec_name not in ["h264", "h265", "hevc", "av1"]:
+                logger.info(f"Video codec is not h264/h265/hevc/av1, transcoding is required: {codec_name}")
                 need_video_transc = True
             elif codec_type == "audio" and codec_name != "aac":
                 logger.info(f"Audio codec is not aac, transcoding is required: {codec_name}")