diff --git a/CHANGELOG.md b/CHANGELOG.md index f7c76a77..8a441697 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,11 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.6.6](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.6.6) - 2021-02-18 + +### Added +- Video upload support + ## [0.6.5](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.6.5) - 2021-02-16 ### Fixed diff --git a/nucleus/__init__.py b/nucleus/__init__.py index ec0f31ed..0273a7c6 100644 --- a/nucleus/__init__.py +++ b/nucleus/__init__.py @@ -14,9 +14,8 @@ "DatasetItem", "DatasetItemRetrievalError", "Frame", - "Frame", - "LidarScene", "LidarScene", + "VideoScene", "Model", "ModelCreationError", # "MultiCategoryAnnotation", # coming soon! @@ -124,7 +123,7 @@ SegmentationPrediction, ) from .retry_strategy import RetryStrategy -from .scene import Frame, LidarScene +from .scene import Frame, LidarScene, VideoScene from .slice import Slice from .upload_response import UploadResponse from .validate import Validate diff --git a/nucleus/constants.py b/nucleus/constants.py index de29b482..45926be4 100644 --- a/nucleus/constants.py +++ b/nucleus/constants.py @@ -44,6 +44,7 @@ ERROR_CODES = "error_codes" ERROR_ITEMS = "upload_errors" ERROR_PAYLOAD = "error_payload" +FRAME_RATE_KEY = "frame_rate" FRAMES_KEY = "frames" FX_KEY = "fx" FY_KEY = "fy" @@ -101,6 +102,10 @@ UPLOAD_TO_SCALE_KEY = "upload_to_scale" URL_KEY = "url" VERTICES_KEY = "vertices" +VIDEO_FRAME_LOCATION_KEY = "video_frame_location" +VIDEO_FRAME_URL_KEY = "video_frame_url" +VIDEO_KEY = "video" +VIDEO_UPLOAD_TYPE_KEY = "video_upload_type" WIDTH_KEY = "width" YAW_KEY = "yaw" W_KEY = "w" diff --git a/nucleus/dataset.py b/nucleus/dataset.py index 6fbc7941..16d16e67 100644 --- a/nucleus/dataset.py +++ b/nucleus/dataset.py @@ -47,6 +47,7 @@ REQUEST_ID_KEY, SLICE_ID_KEY, UPDATE_KEY, + VIDEO_UPLOAD_TYPE_KEY, ) from .data_transfer_object.dataset_info import DatasetInfo from .data_transfer_object.dataset_size import DatasetSize @@ -65,7 +66,7 @@ construct_model_run_creation_payload, construct_taxonomy_payload, ) -from .scene import LidarScene, Scene, check_all_scene_paths_remote +from .scene import LidarScene, Scene, VideoScene, check_all_scene_paths_remote from .slice import Slice from .upload_response import UploadResponse @@ -405,7 +406,9 @@ def ingest_tasks(self, task_ids: List[str]) -> dict: def append( self, - items: Union[Sequence[DatasetItem], Sequence[LidarScene]], + items: Union[ + Sequence[DatasetItem], Sequence[LidarScene], Sequence[VideoScene] + ], update: bool = False, batch_size: int = 20, asynchronous: bool = False, @@ -413,8 +416,7 @@ def append( """Appends items or scenes to a dataset. .. note:: - Datasets can only accept one of :class:`DatasetItems ` - or :class:`Scenes `, never both. + Datasets can only accept one of DatasetItems or Scenes, never both. This behavior is set during Dataset :meth:`creation ` with the ``is_scene`` flag. @@ -478,13 +480,14 @@ def append( Union[ \ Sequence[:class:`DatasetItem`], \ Sequence[:class:`LidarScene`] \ + Sequence[:class:`VideoScene`] ]): List of items or scenes to upload. batch_size: Size of the batch for larger uploads. Default is 20. update: Whether or not to overwrite metadata on reference ID collision. Default is False. asynchronous: Whether or not to process the upload asynchronously (and - return an :class:`AsyncJob` object). This is highly encouraged for - 3D data to drastically increase throughput. Default is False. + return an :class:`AsyncJob` object). This is required when uploading + scenes. Default is False. Returns: For scenes @@ -508,17 +511,26 @@ def append( dataset_items = [ item for item in items if isinstance(item, DatasetItem) ] - scenes = [item for item in items if isinstance(item, LidarScene)] - if dataset_items and scenes: + lidar_scenes = [item for item in items if isinstance(item, LidarScene)] + video_scenes = [item for item in items if isinstance(item, VideoScene)] + if dataset_items and (lidar_scenes or video_scenes): raise Exception( "You must append either DatasetItems or Scenes to the dataset." ) - if scenes: + if lidar_scenes: assert ( asynchronous - ), "In order to avoid timeouts, you must set asynchronous=True when uploading scenes." + ), "In order to avoid timeouts, you must set asynchronous=True when uploading 3D scenes." - return self._append_scenes(scenes, update, asynchronous) + return self._append_scenes(lidar_scenes, update, asynchronous) + if video_scenes: + assert ( + asynchronous + ), "In order to avoid timeouts, you must set asynchronous=True when uploading videos." + + return self._append_video_scenes( + video_scenes, update, asynchronous + ) check_for_duplicate_reference_ids(dataset_items) @@ -601,6 +613,51 @@ def _append_scenes( ) return response + def _append_video_scenes( + self, + scenes: List[VideoScene], + update: Optional[bool] = False, + asynchronous: Optional[bool] = False, + ) -> Union[dict, AsyncJob]: + # TODO: make private in favor of Dataset.append invocation + if not self.is_scene: + raise Exception( + "Your dataset is not a scene dataset but only supports single dataset items. " + "In order to be able to add scenes, please create another dataset with " + "client.create_dataset(, is_scene=True) or add the scenes to " + "an existing scene dataset." + ) + + for scene in scenes: + scene.validate() + + if not asynchronous: + print( + "WARNING: Processing videos usually takes several seconds. As a result, synchronous video scene upload" + "requests are likely to timeout. For large uploads, we recommend using the flag asynchronous=True " + "to avoid HTTP timeouts. Please see" + "https://dashboard.scale.com/nucleus/docs/api?language=python#guide-for-large-ingestions" + " for details." + ) + + if asynchronous: + # TODO check_all_scene_paths_remote(scenes) + request_id = serialize_and_write_to_presigned_url( + scenes, self.id, self._client + ) + response = self._client.make_request( + payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update}, + route=f"{self.id}/upload_video_scenes?async=1", + ) + return AsyncJob.from_json(response, self._client) + + payload = construct_append_scenes_payload(scenes, update) + response = self._client.make_request( + payload=payload, + route=f"{self.id}/upload_video_scenes", + ) + return response + def iloc(self, i: int) -> dict: """Retrieves dataset item by absolute numerical index. @@ -1082,13 +1139,14 @@ def get_scene(self, reference_id: str) -> Scene: :class:`Scene`: A scene object containing frames, which in turn contain pointcloud or image items. """ - return LidarScene.from_json( - self._client.make_request( - payload=None, - route=f"dataset/{self.id}/scene/{reference_id}", - requests_command=requests.get, - ) + response = self._client.make_request( + payload=None, + route=f"dataset/{self.id}/scene/{reference_id}", + requests_command=requests.get, ) + if VIDEO_UPLOAD_TYPE_KEY in response: + return VideoScene.from_json(response) + return LidarScene.from_json(response) def export_predictions(self, model): """Fetches all predictions of a model that were uploaded to the dataset. diff --git a/nucleus/dataset_item.py b/nucleus/dataset_item.py index 7f6b912c..b4974719 100644 --- a/nucleus/dataset_item.py +++ b/nucleus/dataset_item.py @@ -22,6 +22,7 @@ TYPE_KEY, UPLOAD_TO_SCALE_KEY, URL_KEY, + VIDEO_FRAME_URL_KEY, W_KEY, X_KEY, Y_KEY, @@ -120,34 +121,42 @@ def to_payload(self) -> dict: class DatasetItemType(Enum): IMAGE = "image" POINTCLOUD = "pointcloud" + VIDEO = "video" @dataclass # pylint: disable=R0902 class DatasetItem: # pylint: disable=R0902 - """A dataset item is an image or pointcloud that has associated metadata. + """A dataset item is an image, pointcloud or video frame that has associated metadata. Note: for 3D data, please include a :class:`CameraParams` object under a key named "camera_params" within the metadata dictionary. This will allow for projecting 3D annotations to any image within a scene. Args: - image_location (Optional[str]): Required if pointcloud_location not present: The - location containing the image for the given row of data. This can be a - local path, or a remote URL. Remote formats supported include any URL - (``http://`` or ``https://``) or URIs for AWS S3, Azure, or GCS - (i.e. ``s3://``, ``gcs://``). - - pointcloud_location (Optional[str]): Required if image_location not - present: The remote URL containing the pointcloud JSON. Remote - formats supported include any URL (``http://`` or ``https://``) or - URIs for AWS S3, Azure, or GCS (i.e. ``s3://``, ``gcs://``). + image_location (Optional[str]): Required if pointcloud_location and + video_frame_location are not present: The location containing the image for + the given row of data. This can be a local path, or a remote URL. Remote + formats supported include any URL (``http://`` or ``https://``) or URIs for + AWS S3, Azure, or GCS (i.e. ``s3://``, ``gcs://``). + + pointcloud_location (Optional[str]): Required if image_location and + video_frame_location are not present: The remote URL containing the + pointcloud JSON. Remote formats supported include any URL (``http://`` + or ``https://``) or URIs for AWS S3, Azure, or GCS (i.e. ``s3://``, + ``gcs://``). + + video_frame_location (Optional[str]): Required if image_location and + pointcloud_location are not present: The remote URL containing the + video frame image. Remote formats supported include any URL (``http://`` + or ``https://``) or URIs for AWS S3, Azure, or GCS (i.e. ``s3://``, + ``gcs://``). reference_id (Optional[str]): A user-specified identifier to reference the item. metadata (Optional[dict]): Extra information about the particular dataset item. ints, floats, string values will be made searchable in - the query bar by the key in this dict For example, ``{"animal": + the query bar by the key in this dict. For example, ``{"animal": "dog"}`` will become searchable via ``metadata.animal = "dog"``. Categorical data can be passed as a string and will be treated @@ -190,9 +199,10 @@ class DatasetItem: # pylint: disable=R0902 upload_to_scale (Optional[bool]): Set this to false in order to use `privacy mode `_. - Setting this to false means the actual data within the item (i.e. the - image or pointcloud) will not be uploaded to scale meaning that you can - send in links that are only accessible to certain users, and not to Scale. + Setting this to false means the actual data within the item will not be + uploaded to scale meaning that you can send in links that are only accessible + to certain users, and not to Scale. Skipping upload to Scale is currently only + implemented for images. """ image_location: Optional[str] = None @@ -202,15 +212,21 @@ class DatasetItem: # pylint: disable=R0902 metadata: Optional[dict] = None pointcloud_location: Optional[str] = None upload_to_scale: Optional[bool] = True + video_frame_location: Optional[str] = None def __post_init__(self): assert self.reference_id != "DUMMY_VALUE", "reference_id is required." - assert bool(self.image_location) != bool( - self.pointcloud_location - ), "Must specify exactly one of the image_location, pointcloud_location parameters" - if self.pointcloud_location and not self.upload_to_scale: + assert ( + bool(self.image_location) + + bool(self.pointcloud_location) + + bool(self.video_frame_location) + == 1 + ), "Must specify exactly one of the image_location, pointcloud_location, video_frame_location parameters" + if ( + self.pointcloud_location or self.video_frame_location + ) and not self.upload_to_scale: raise NotImplementedError( - "Skipping upload to Scale is not currently implemented for pointclouds." + "Skipping upload to Scale is not currently implemented for pointclouds and videos." ) self.local = ( is_local_path(self.image_location) if self.image_location else None @@ -218,7 +234,11 @@ def __post_init__(self): self.type = ( DatasetItemType.IMAGE if self.image_location - else DatasetItemType.POINTCLOUD + else ( + DatasetItemType.POINTCLOUD + if self.pointcloud_location + else DatasetItemType.VIDEO + ) ) camera_params = ( self.metadata.get(CAMERA_PARAMS_KEY, None) @@ -238,6 +258,7 @@ def from_json(cls, payload: dict): return cls( image_location=image_url, pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None), + video_frame_location=payload.get(VIDEO_FRAME_URL_KEY, None), reference_id=payload.get(REFERENCE_ID_KEY, None), metadata=payload.get(METADATA_KEY, {}), upload_to_scale=payload.get(UPLOAD_TO_SCALE_KEY, True), @@ -260,13 +281,15 @@ def to_payload(self, is_scene=False) -> dict: payload[URL_KEY] = self.image_location elif self.pointcloud_location: payload[URL_KEY] = self.pointcloud_location + elif self.video_frame_location: + payload[URL_KEY] = self.video_frame_location payload[TYPE_KEY] = self.type.value if self.camera_params: payload[CAMERA_PARAMS_KEY] = self.camera_params.to_payload() else: assert ( self.image_location - ), "Must specify image_location for DatasetItems not in a LidarScene" + ), "Must specify image_location for DatasetItems not in a LidarScene or VideoScene" payload[IMAGE_URL_KEY] = self.image_location payload[UPLOAD_TO_SCALE_KEY] = self.upload_to_scale diff --git a/nucleus/payload_constructor.py b/nucleus/payload_constructor.py index 6bafd70b..38609368 100644 --- a/nucleus/payload_constructor.py +++ b/nucleus/payload_constructor.py @@ -32,7 +32,7 @@ PolygonPrediction, SegmentationPrediction, ) -from .scene import LidarScene +from .scene import LidarScene, VideoScene def construct_append_payload( @@ -50,7 +50,8 @@ def construct_append_payload( def construct_append_scenes_payload( - scene_list: List[LidarScene], update: Optional[bool] = False + scene_list: Union[List[LidarScene], List[VideoScene]], + update: Optional[bool] = False, ) -> dict: scenes = [] for scene in scene_list: diff --git a/nucleus/scene.py b/nucleus/scene.py index 339b8f4e..b587581a 100644 --- a/nucleus/scene.py +++ b/nucleus/scene.py @@ -1,9 +1,11 @@ import json from abc import ABC from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from enum import Enum +from typing import Any, Dict, List, Optional, Union from nucleus.constants import ( + FRAME_RATE_KEY, FRAMES_KEY, IMAGE_LOCATION_KEY, LENGTH_KEY, @@ -11,6 +13,8 @@ NUM_SENSORS_KEY, POINTCLOUD_LOCATION_KEY, REFERENCE_ID_KEY, + VIDEO_FRAME_LOCATION_KEY, + VIDEO_UPLOAD_TYPE_KEY, ) from .annotation import is_local_path @@ -409,7 +413,179 @@ def flatten(t): return [item for sublist in t for item in sublist] -def check_all_scene_paths_remote(scenes: List[LidarScene]): +class _VideoUploadType(Enum): + IMAGE = "image" + VIDEO = "video" + + +@dataclass +class VideoScene(ABC): + """ + Nucleus video datasets are comprised of VideoScenes, which are in turn + comprised of a sequence of :class:`DatasetItems ` which are + equivalent to frames. + + VideoScenes are uploaded to a :class:`Dataset` with any accompanying + metadata. Each of :class:`DatasetItems ` representing a frame + also accepts metadata. + + Note: Uploads with different items will error out (only on scenes that + now differ). Existing video are expected to retain the same frames, and only + metadata can be updated. If a video definition is changed (for example, + additional frames added) the update operation will be ignored. If you would + like to alter the structure of a video scene, please delete the scene and + re-upload. + + Parameters: + reference_id (str): User-specified identifier to reference the scene. + frame_rate (int): Frame rate of the video. + attachment_type (str): The type of attachments being uploaded as a string literal. + Currently, videos can only be uploaded as an array of frames, so the only + accepted attachment_type is "image". + items (Optional[List[:class:`DatasetItem`]]): List of items to be a part of + the scene. A scene can be created before items have been added to it, + but must be non-empty when uploading to a :class:`Dataset`. + metadata (Optional[Dict]): Optional metadata to include with the scene. + + Refer to our `guide to uploading video data + `_ for more info! + """ + + reference_id: str + frame_rate: int + attachment_type: _VideoUploadType + items: List[DatasetItem] = field(default_factory=list) + metadata: Optional[dict] = field(default_factory=dict) + + def __post_init__(self): + assert ( + self.attachment_type != _VideoUploadType.IMAGE + ), "Videos can currently only be uploaded from frames" + if self.metadata is None: + self.metadata = {} + + def __eq__(self, other): + return all( + [ + self.reference_id == other.reference_id, + self.items == other.items, + self.metadata == other.metadata, + ] + ) + + @property + def length(self) -> int: + """Number of items in the scene.""" + return len(self.items) + + def validate(self): + # TODO: make private + assert self.frame_rate > 0, "Frame rate must be at least 1" + assert self.length > 0, "Must have at least 1 item in a scene" + for item in self.items: + assert isinstance( + item, DatasetItem + ), "Each item in a scene must be a DatasetItem object" + assert ( + item.video_frame_location is not None + ), "Each item in a scene must have a video_frame_location" + + def add_item( + self, item: DatasetItem, index: int = None, update: bool = False + ) -> None: + """Adds DatasetItem to the specified index. + + Parameters: + item (:class:`DatasetItem`): Video item to add. + index: Serial index at which to add the item. + update: Whether to overwrite the item at the specified index, if it + exists. Default is False. + """ + if index is None: + index = len(self.items) + assert ( + 0 <= index <= len(self.items) + ), f"Video scenes must be contiguous so index must be at least 0 and at most {len(self.items)}." + if index < len(self.items) and update: + self.items[index] = item + else: + self.items.append(item) + + def get_item(self, index: int) -> DatasetItem: + """Fetches the DatasetItem at the specified index. + + Parameters: + index: Serial index for which to retrieve the DatasetItem. + + Return: + :class:`DatasetItem`: DatasetItem at the specified index.""" + if index < 0 or index > len(self.items): + raise ValueError( + f"This scene does not have an item at index {index}" + ) + return self.items[index] + + def get_items(self) -> List[DatasetItem]: + """Fetches a sorted list of DatasetItems of the scene. + + Returns: + List[:class:`DatasetItem`]: List of DatasetItems, sorted by index ascending. + """ + return self.items + + def info(self): + """Fetches information about the scene. + + Returns: + Payload containing:: + + { + "reference_id": str, + "length": int, + "num_sensors": int + } + """ + return { + REFERENCE_ID_KEY: self.reference_id, + FRAME_RATE_KEY: self.frame_rate, + LENGTH_KEY: self.length, + } + + @classmethod + def from_json(cls, payload: dict): + """Instantiates scene object from schematized JSON dict payload.""" + items_payload = payload.get(FRAMES_KEY, []) + items = [DatasetItem.from_json(item) for item in items_payload] + return cls( + reference_id=payload[REFERENCE_ID_KEY], + frame_rate=payload[FRAME_RATE_KEY], + attachment_type=payload[VIDEO_UPLOAD_TYPE_KEY], + items=items, + metadata=payload.get(METADATA_KEY, {}), + ) + + def to_payload(self) -> dict: + """Serializes scene object to schematized JSON dict.""" + self.validate() + items_payload = [item.to_payload(is_scene=True) for item in self.items] + payload: Dict[str, Any] = { + REFERENCE_ID_KEY: self.reference_id, + VIDEO_UPLOAD_TYPE_KEY: self.attachment_type, + FRAME_RATE_KEY: self.frame_rate, + FRAMES_KEY: items_payload, + } + if self.metadata: + payload[METADATA_KEY] = self.metadata + return payload + + def to_json(self) -> str: + """Serializes scene object to schematized JSON string.""" + return json.dumps(self.to_payload(), allow_nan=False) + + +def check_all_scene_paths_remote( + scenes: Union[List[LidarScene], List[VideoScene]] +): for scene in scenes: for item in scene.get_items(): pointcloud_location = getattr(item, POINTCLOUD_LOCATION_KEY) @@ -424,3 +600,9 @@ def check_all_scene_paths_remote(scenes: List[LidarScene]): f"All paths for DatasetItems in a Scene must be remote, but {item.image_location} is either " "local, or a remote URL type that is not supported." ) + video_frame_location = getattr(item, VIDEO_FRAME_LOCATION_KEY) + if video_frame_location and is_local_path(video_frame_location): + raise ValueError( + f"All paths for DatasetItems in a Scene must be remote, but {item.video_frame_location} is either " + "local, or a remote URL type that is not supported." + ) diff --git a/nucleus/utils.py b/nucleus/utils.py index 98071acf..3be51e77 100644 --- a/nucleus/utils.py +++ b/nucleus/utils.py @@ -39,7 +39,7 @@ PolygonPrediction, SegmentationPrediction, ) -from .scene import LidarScene +from .scene import LidarScene, VideoScene STRING_REPLACEMENTS = { "\\\\n": "\n", @@ -215,7 +215,9 @@ def convert_export_payload(api_payload): def serialize_and_write( - upload_units: Sequence[Union[DatasetItem, Annotation, LidarScene]], + upload_units: Sequence[ + Union[DatasetItem, Annotation, LidarScene, VideoScene] + ], file_pointer, ): if len(upload_units) == 0: @@ -224,7 +226,9 @@ def serialize_and_write( ) for unit in upload_units: try: - if isinstance(unit, (DatasetItem, Annotation, LidarScene)): + if isinstance( + unit, (DatasetItem, Annotation, LidarScene, VideoScene) + ): file_pointer.write(unit.to_json() + "\n") else: file_pointer.write(json.dumps(unit) + "\n") @@ -254,7 +258,9 @@ def upload_to_presigned_url(presigned_url: str, file_pointer: IO): def serialize_and_write_to_presigned_url( - upload_units: Sequence[Union[DatasetItem, Annotation, LidarScene]], + upload_units: Sequence[ + Union[DatasetItem, Annotation, LidarScene, VideoScene] + ], dataset_id: str, client, ): diff --git a/pyproject.toml b/pyproject.toml index 915cf435..b710abbb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ exclude = ''' [tool.poetry] name = "scale-nucleus" -version = "0.6.5" +version = "0.6.6" description = "The official Python client library for Nucleus, the Data Platform for AI" license = "MIT" authors = ["Scale AI Nucleus Team "] diff --git a/tests/helpers.py b/tests/helpers.py index 4781c33e..beac972c 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -84,6 +84,32 @@ "update": False, } +TEST_VIDEO_SCENES = { + "scenes": [ + { + "reference_id": "scene_1", + "video_upload_type": "image", + "frame_rate": 15, + "frames": [ + { + "video_frame_url": TEST_IMG_URLS[0], + "type": "video", + "reference_id": "video_frame_0", + "metadata": {"time": 123, "foo": "bar"}, + }, + { + "video_frame_url": TEST_IMG_URLS[1], + "type": "video", + "reference_id": "video_frame_1", + "metadata": {"time": 124, "foo": "bar_2"}, + }, + ], + "metadata": {"timestamp": "1234", "weather": "rainy"}, + } + ], + "update": False, +} + def reference_id_from_url(url): return Path(url).name @@ -96,6 +122,33 @@ def reference_id_from_url(url): DatasetItem(TEST_IMG_URLS[3], reference_id_from_url(TEST_IMG_URLS[3])), ] +TEST_VIDEO_ITEMS = [ + DatasetItem( + None, + reference_id_from_url(TEST_IMG_URLS[0]), + None, + None, + True, + TEST_IMG_URLS[0], + ), + DatasetItem( + None, + reference_id_from_url(TEST_IMG_URLS[1]), + None, + None, + True, + TEST_IMG_URLS[1], + ), + DatasetItem( + None, + reference_id_from_url(TEST_IMG_URLS[2]), + None, + None, + True, + TEST_IMG_URLS[2], + ), +] + TEST_LIDAR_ITEMS = [ DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[0], reference_id="1"), DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[1], reference_id="2"), diff --git a/tests/test_scene.py b/tests/test_scene.py index 47dce8c4..3069ef75 100644 --- a/tests/test_scene.py +++ b/tests/test_scene.py @@ -3,9 +3,16 @@ import pytest -from nucleus import CuboidAnnotation, DatasetItem, Frame, LidarScene +from nucleus import ( + CuboidAnnotation, + DatasetItem, + Frame, + LidarScene, + VideoScene, +) from nucleus.constants import ( ANNOTATIONS_KEY, + FRAME_RATE_KEY, FRAMES_KEY, IMAGE_KEY, IMAGE_URL_KEY, @@ -19,6 +26,8 @@ TYPE_KEY, UPDATE_KEY, URL_KEY, + VIDEO_KEY, + VIDEO_UPLOAD_TYPE_KEY, ) from nucleus.scene import flatten @@ -28,6 +37,8 @@ TEST_DATASET_ITEMS, TEST_LIDAR_ITEMS, TEST_LIDAR_SCENES, + TEST_VIDEO_ITEMS, + TEST_VIDEO_SCENES, assert_cuboid_annotation_matches_dict, ) @@ -258,6 +269,59 @@ def test_scene_add_frame(): } +def test_video_scene_property_methods(): + payload = TEST_VIDEO_SCENES + expected_frame_rate = TEST_VIDEO_SCENES["scenes"][0]["frame_rate"] + scene_json = payload[SCENES_KEY][0] + scene = VideoScene.from_json(scene_json) + + expected_length = len(scene_json[FRAMES_KEY]) + assert scene.length == expected_length + assert scene.info() == { + REFERENCE_ID_KEY: scene_json[REFERENCE_ID_KEY], + LENGTH_KEY: expected_length, + FRAME_RATE_KEY: expected_frame_rate, + } + + +def test_video_scene_add_item(): + scene_ref_id = "scene_1" + frame_rate = 20 + video_upload_type = "image" + scene = VideoScene(scene_ref_id, frame_rate, video_upload_type) + scene.add_item(TEST_VIDEO_ITEMS[0]) + scene.add_item(TEST_VIDEO_ITEMS[1], index=1) + scene.add_item(TEST_VIDEO_ITEMS[2], index=0, update=True) + + assert scene.get_item(0) == TEST_VIDEO_ITEMS[2] + assert scene.get_item(1) == TEST_VIDEO_ITEMS[1] + + assert scene.get_items() == [ + TEST_VIDEO_ITEMS[2], + TEST_VIDEO_ITEMS[1], + ] + + assert scene.to_payload() == { + REFERENCE_ID_KEY: scene_ref_id, + VIDEO_UPLOAD_TYPE_KEY: video_upload_type, + FRAME_RATE_KEY: frame_rate, + FRAMES_KEY: [ + { + URL_KEY: TEST_VIDEO_ITEMS[2].video_frame_location, + REFERENCE_ID_KEY: TEST_VIDEO_ITEMS[2].reference_id, + TYPE_KEY: VIDEO_KEY, + METADATA_KEY: TEST_VIDEO_ITEMS[2].metadata or {}, + }, + { + URL_KEY: TEST_VIDEO_ITEMS[1].video_frame_location, + REFERENCE_ID_KEY: TEST_VIDEO_ITEMS[1].reference_id, + TYPE_KEY: VIDEO_KEY, + METADATA_KEY: TEST_VIDEO_ITEMS[1].metadata or {}, + }, + ], + } + + @pytest.mark.skip("Deactivated sync upload for scenes") def test_scene_upload_sync(dataset_scene): payload = TEST_LIDAR_SCENES @@ -499,3 +563,155 @@ def test_scene_metadata_update(dataset_scene): updated_scene = dataset_scene.get_scene(scene_ref_id) actual_metadata = updated_scene.metadata assert expected_new_metadata == actual_metadata + + +@pytest.mark.integration +def test_video_scene_upload_async(dataset_scene): + payload = TEST_VIDEO_SCENES + scenes = [ + VideoScene.from_json(scene_json) for scene_json in payload[SCENES_KEY] + ] + update = payload[UPDATE_KEY] + + job = dataset_scene.append(scenes, update=update, asynchronous=True) + job.sleep_until_complete() + status = job.status() + + assert status == { + "job_id": job.job_id, + "status": "Completed", + "message": { + "scene_upload_progress": { + "errors": [], + "dataset_id": dataset_scene.id, + "new_scenes": len(scenes), + "ignored_scenes": 0, + "scenes_errored": 0, + "updated_scenes": 0, + } + }, + "job_progress": "1.00", + "completed_steps": 1, + "total_steps": 1, + } + + uploaded_scenes = dataset_scene.scenes + assert len(uploaded_scenes) == len(scenes) + assert all( + u["reference_id"] == o.reference_id + for u, o in zip(uploaded_scenes, scenes) + ) + assert all( + u["metadata"] == o.metadata or (not u["metadata"] and not o.metadata) + for u, o in zip(uploaded_scenes, scenes) + ) + + +@pytest.mark.integration +def test_video_scene_upload_and_update(dataset_scene): + payload = TEST_VIDEO_SCENES + scenes = [ + VideoScene.from_json(scene_json) for scene_json in payload[SCENES_KEY] + ] + update = payload[UPDATE_KEY] + + job = dataset_scene.append(scenes, update=update, asynchronous=True) + job.sleep_until_complete() + status = job.status() + + assert status == { + "job_id": job.job_id, + "status": "Completed", + "message": { + "scene_upload_progress": { + "errors": [], + "dataset_id": dataset_scene.id, + "new_scenes": len(scenes), + "ignored_scenes": 0, + "scenes_errored": 0, + "updated_scenes": 0, + } + }, + "job_progress": "1.00", + "completed_steps": 1, + "total_steps": 1, + } + + uploaded_scenes = dataset_scene.scenes + assert len(uploaded_scenes) == len(scenes) + assert all( + u["reference_id"] == o.reference_id + for u, o in zip(uploaded_scenes, scenes) + ) + assert all( + u["metadata"] == o.metadata or (not u["metadata"] and not o.metadata) + for u, o in zip(uploaded_scenes, scenes) + ) + + job2 = dataset_scene.append(scenes, update=True, asynchronous=True) + job2.sleep_until_complete() + status2 = job2.status() + + assert status2 == { + "job_id": job2.job_id, + "status": "Completed", + "message": { + "scene_upload_progress": { + "errors": [], + "dataset_id": dataset_scene.id, + "new_scenes": 0, + "ignored_scenes": 0, + "scenes_errored": 0, + "updated_scenes": len(scenes), + } + }, + "job_progress": "1.00", + "completed_steps": 1, + "total_steps": 1, + } + + +@pytest.mark.integration +def test_video_scene_deletion(dataset_scene): + payload = TEST_VIDEO_SCENES + scenes = [ + VideoScene.from_json(scene_json) for scene_json in payload[SCENES_KEY] + ] + update = payload[UPDATE_KEY] + + job = dataset_scene.append(scenes, update=update, asynchronous=True) + job.sleep_until_complete() + + uploaded_scenes = dataset_scene.scenes + assert len(uploaded_scenes) == len(scenes) + assert all( + u["reference_id"] == o.reference_id + for u, o in zip(uploaded_scenes, scenes) + ) + + for scene in uploaded_scenes: + dataset_scene.delete_scene(scene.reference_id) + time.sleep(1) + scenes = dataset_scene.scenes + assert len(scenes) == 0, f"Expected to delete all scenes, got: {scenes}" + + +@pytest.mark.integration +def test_video_scene_metadata_update(dataset_scene): + payload = TEST_VIDEO_SCENES + scenes = [ + VideoScene.from_json(scene_json) for scene_json in payload[SCENES_KEY] + ] + update = payload[UPDATE_KEY] + + job = dataset_scene.append(scenes, update=update, asynchronous=True) + job.sleep_until_complete() + + scene_ref_id = scenes[0].reference_id + additional_metadata = {"some_new_key": 123} + dataset_scene.update_scene_metadata({scene_ref_id: additional_metadata}) + expected_new_metadata = {**scenes[0].metadata, **additional_metadata} + + updated_scene = dataset_scene.get_scene(scene_ref_id) + actual_metadata = updated_scene.metadata + assert expected_new_metadata == actual_metadata