Skip to content

Commit af14a2e

Browse files
sasha-scalesasha-scale
authored
Merge pull request #84 from scaleapi/upload_scene_integration_test
Upload scene integration test
2 parents 3b6aeab + f0a57cf commit af14a2e

File tree

10 files changed

+258
-35
lines changed

10 files changed

+258
-35
lines changed

‎nucleus/__init__.py

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -615,21 +615,23 @@ def annotate_dataset(
615615
self,
616616
dataset_id: str,
617617
annotations: List[
618-
Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
618+
Union[
619+
BoxAnnotation,
620+
PolygonAnnotation,
621+
CuboidAnnotation,
622+
SegmentationAnnotation,
623+
]
619624
],
620625
update: bool,
621626
batch_size: int = 5000,
622627
):
623628
"""
624629
Uploads ground truth annotations for a given dataset.
625630
:param dataset_id: id of the dataset
626-
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
631+
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]]
627632
:param update: whether to update or ignore conflicting annotations
628633
:return: {"dataset_id: str, "annotations_processed": int}
629634
"""
630-
if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
631-
raise NotImplementedError("Cuboid annotations not yet supported")
632-
633635
# Split payload into segmentations and Box/Polygon
634636
segmentations = [
635637
ann
@@ -770,14 +772,19 @@ def predict(
770772
self,
771773
model_run_id: str,
772774
annotations: List[
773-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
775+
Union[
776+
BoxPrediction,
777+
PolygonPrediction,
778+
CuboidPrediction,
779+
SegmentationPrediction,
780+
]
774781
],
775782
update: bool,
776783
batch_size: int = 5000,
777784
):
778785
"""
779786
Uploads model outputs as predictions for a model_run. Returns info about the upload.
780-
:param annotations: List[Union[BoxPrediction, PolygonPrediction]],
787+
:param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
781788
:param update: bool
782789
:return:
783790
{
@@ -787,9 +794,6 @@ def predict(
787794
"predictions_ignored": int,
788795
}
789796
"""
790-
if any((isinstance(ann, CuboidPrediction) for ann in annotations)):
791-
raise NotImplementedError("Cuboid predictions not yet supported")
792-
793797
segmentations = [
794798
ann
795799
for ann in annotations
@@ -941,7 +945,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
941945
:param reference_id: reference_id of a dataset item.
942946
:return:
943947
{
944-
"annotations": List[BoxPrediction],
948+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
945949
}
946950
"""
947951
return self.make_request(
@@ -966,7 +970,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
966970
:param i: absolute number of Dataset Item for a dataset corresponding to the model run.
967971
:return:
968972
{
969-
"annotations": List[BoxPrediction],
973+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
970974
}
971975
"""
972976
return self.make_request(
@@ -995,7 +999,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
995999
:param dataset_item_id: dataset_item_id of a dataset item.
9961000
:return:
9971001
{
998-
"annotations": List[BoxPrediction],
1002+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
9991003
}
10001004
"""
10011005
return self.make_request(

‎nucleus/annotation.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -299,16 +299,26 @@ def to_payload(self) -> dict:
299299
LABEL_KEY: self.label,
300300
TYPE_KEY: CUBOID_TYPE,
301301
GEOMETRY_KEY: {
302-
POSITION_KEY: self.position,
303-
DIMENSIONS_KEY: self.dimensions,
302+
POSITION_KEY: self.position.to_payload(),
303+
DIMENSIONS_KEY: self.dimensions.to_payload(),
304304
YAW_KEY: self.yaw,
305305
},
306306
REFERENCE_ID_KEY: self.reference_id,
307+
ITEM_ID_KEY: self.item_id,
307308
ANNOTATION_ID_KEY: self.annotation_id,
308309
METADATA_KEY: self.metadata,
309310
}
310311

311312

313+
def check_all_frame_paths_remote(frames: List[str]):
314+
for frame_url in frames:
315+
if is_local_path(frame_url):
316+
raise ValueError(
317+
f"All paths must be remote, but {frame_url} is either "
318+
"local, or a remote URL type that is not supported."
319+
)
320+
321+
312322
def check_all_mask_paths_remote(
313323
annotations: Sequence[Union[Annotation]],
314324
):

‎nucleus/constants.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
ERROR_CODES = "error_codes"
3131
ERROR_ITEMS = "upload_errors"
3232
ERROR_PAYLOAD = "error_payload"
33+
FRAMES = "frames"
3334
GEOMETRY_KEY = "geometry"
3435
HEIGHT_KEY = "height"
3536
IGNORED_ITEMS = "ignored_items"
@@ -62,6 +63,7 @@
6263
REFERENCE_IDS_KEY = "reference_ids"
6364
REFERENCE_ID_KEY = "reference_id"
6465
REQUEST_ID_KEY = "requestId"
66+
SCENES = "scenes"
6567
SEGMENTATIONS_KEY = "segmentations"
6668
SLICE_ID_KEY = "slice_id"
6769
STATUS_CODE_KEY = "status_code"
@@ -70,6 +72,7 @@
7072
TYPE_KEY = "type"
7173
UPDATED_ITEMS = "updated_items"
7274
UPDATE_KEY = "update"
75+
URL = "url"
7376
VERTICES_KEY = "vertices"
7477
WIDTH_KEY = "width"
7578
YAW_KEY = "yaw"

‎nucleus/dataset.py

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212

1313
from .annotation import (
1414
Annotation,
15-
CuboidAnnotation,
1615
check_all_mask_paths_remote,
16+
check_all_frame_paths_remote,
1717
)
1818
from .constants import (
1919
DATASET_ITEM_IDS_KEY,
@@ -23,10 +23,13 @@
2323
DATASET_SLICES_KEY,
2424
DEFAULT_ANNOTATION_UPDATE_MODE,
2525
EXPORTED_ROWS,
26+
FRAMES,
2627
NAME_KEY,
2728
REFERENCE_IDS_KEY,
2829
REQUEST_ID_KEY,
30+
SCENES,
2931
UPDATE_KEY,
32+
URL,
3033
)
3134
from .dataset_item import (
3235
DatasetItem,
@@ -167,9 +170,6 @@ def annotate(
167170
"ignored_items": int,
168171
}
169172
"""
170-
if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
171-
raise NotImplementedError("Cuboid annotations not yet supported")
172-
173173
check_all_mask_paths_remote(annotations)
174174

175175
if asynchronous:
@@ -248,14 +248,53 @@ def append(
248248
batch_size=batch_size,
249249
)
250250

251+
def upload_scenes(
252+
self,
253+
payload: dict,
254+
update: Optional[bool] = False,
255+
asynchronous: bool = False,
256+
) -> Union[dict, AsyncJob]:
257+
"""
258+
Uploads scenes with given frames to the dataset
259+
260+
Parameters:
261+
:param payload: dictionary containing scenes to be uploaded
262+
:param update: if True, overwrite scene on collision
263+
:param aynchronous: if True, return a job object representing asynchronous ingestion job
264+
:return:
265+
{
266+
'dataset_id': str,
267+
'new_scenes': int,
268+
}
269+
"""
270+
if asynchronous:
271+
for scene in payload[SCENES]:
272+
for frame in scene[FRAMES]:
273+
check_all_frame_paths_remote(frame[URL])
274+
request_id = serialize_and_write_to_presigned_url(
275+
[payload], self.id, self._client
276+
)
277+
response = self._client.make_request(
278+
payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
279+
route=f"{self.id}/upload_scenes?async=1",
280+
)
281+
return AsyncJob.from_json(response, self._client)
282+
283+
# TODO: create client method for sync scene upload
284+
response = self._client.make_request(
285+
payload=payload,
286+
route=f"{self.id}/upload_scenes",
287+
)
288+
return response
289+
251290
def iloc(self, i: int) -> dict:
252291
"""
253292
Returns Dataset Item Info By Dataset Item Number.
254293
:param i: absolute number of dataset item for the given dataset.
255294
:return:
256295
{
257296
"item": DatasetItem,
258-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
297+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
259298
}
260299
"""
261300
response = self._client.dataitem_iloc(self.id, i)
@@ -268,7 +307,7 @@ def refloc(self, reference_id: str) -> dict:
268307
:return:
269308
{
270309
"item": DatasetItem,
271-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
310+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
272311
}
273312
"""
274313
response = self._client.dataitem_ref_id(self.id, reference_id)
@@ -281,7 +320,7 @@ def loc(self, dataset_item_id: str) -> dict:
281320
:return:
282321
{
283322
"item": DatasetItem,
284-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
323+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
285324
}
286325
"""
287326
response = self._client.dataitem_loc(self.id, dataset_item_id)

‎nucleus/model.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from .dataset import Dataset
33
from .prediction import (
44
BoxPrediction,
5+
CuboidPrediction,
56
PolygonPrediction,
67
SegmentationPrediction,
78
)
@@ -57,7 +58,12 @@ def create_run(
5758
name: str,
5859
dataset: Dataset,
5960
predictions: List[
60-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
61+
Union[
62+
BoxPrediction,
63+
PolygonPrediction,
64+
CuboidPrediction,
65+
SegmentationPrediction,
66+
]
6167
],
6268
metadata: Optional[Dict] = None,
6369
asynchronous: bool = False,

‎nucleus/model_run.py

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from .constants import (
88
ANNOTATIONS_KEY,
99
BOX_TYPE,
10+
CUBOID_TYPE,
1011
DEFAULT_ANNOTATION_UPDATE_MODE,
1112
POLYGON_TYPE,
1213
REQUEST_ID_KEY,
@@ -15,6 +16,7 @@
1516
)
1617
from .prediction import (
1718
BoxPrediction,
19+
CuboidPrediction,
1820
PolygonPrediction,
1921
SegmentationPrediction,
2022
)
@@ -89,14 +91,19 @@ def commit(self, payload: Optional[dict] = None) -> dict:
8991
def predict(
9092
self,
9193
annotations: List[
92-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
94+
Union[
95+
BoxPrediction,
96+
PolygonPrediction,
97+
CuboidPrediction,
98+
SegmentationPrediction,
99+
]
93100
],
94101
update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
95102
asynchronous: bool = False,
96103
) -> Union[dict, AsyncJob]:
97104
"""
98105
Uploads model outputs as predictions for a model_run. Returns info about the upload.
99-
:param annotations: List[Union[BoxPrediction, PolygonPrediction]],
106+
:param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
100107
:return:
101108
{
102109
"model_run_id": str,
@@ -122,7 +129,7 @@ def iloc(self, i: int):
122129
"""
123130
Returns Model Run Info For Dataset Item by its number.
124131
:param i: absolute number of Dataset Item for a dataset corresponding to the model run.
125-
:return: List[Union[BoxPrediction, PolygonPrediction]],
132+
:return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
126133
}
127134
"""
128135
response = self._client.predictions_iloc(self.model_run_id, i)
@@ -132,7 +139,7 @@ def refloc(self, reference_id: str):
132139
"""
133140
Returns Model Run Info For Dataset Item by its reference_id.
134141
:param reference_id: reference_id of a dataset item.
135-
:return: List[Union[BoxPrediction, PolygonPrediction]],
142+
:return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
136143
"""
137144
response = self._client.predictions_ref_id(
138145
self.model_run_id, reference_id
@@ -167,7 +174,14 @@ def _format_prediction_response(
167174
self, response: dict
168175
) -> Union[
169176
dict,
170-
List[Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]],
177+
List[
178+
Union[
179+
BoxPrediction,
180+
PolygonPrediction,
181+
CuboidPrediction,
182+
SegmentationPrediction,
183+
]
184+
],
171185
]:
172186
annotation_payload = response.get(ANNOTATIONS_KEY, None)
173187
if not annotation_payload:
@@ -179,11 +193,13 @@ def _format_prediction_response(
179193
Union[
180194
Type[BoxPrediction],
181195
Type[PolygonPrediction],
196+
Type[CuboidPrediction],
182197
Type[SegmentationPrediction],
183198
],
184199
] = {
185200
BOX_TYPE: BoxPrediction,
186201
POLYGON_TYPE: PolygonPrediction,
202+
CUBOID_TYPE: CuboidPrediction,
187203
SEGMENTATION_TYPE: SegmentationPrediction,
188204
}
189205
for type_key in annotation_payload:

‎nucleus/payload_constructor.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22
from .dataset_item import DatasetItem
33
from .annotation import (
44
BoxAnnotation,
5+
CuboidAnnotation,
56
PolygonAnnotation,
67
SegmentationAnnotation,
78
)
89
from .prediction import (
910
BoxPrediction,
11+
CuboidPrediction,
1012
PolygonPrediction,
1113
SegmentationPrediction,
1214
)
@@ -39,7 +41,14 @@ def construct_append_payload(
3941

4042

4143
def construct_annotation_payload(
42-
annotation_items: List[Union[BoxAnnotation, PolygonAnnotation]],
44+
annotation_items: List[
45+
Union[
46+
BoxAnnotation,
47+
PolygonAnnotation,
48+
CuboidAnnotation,
49+
SegmentationAnnotation,
50+
]
51+
],
4352
update: bool,
4453
) -> dict:
4554
annotations = []
@@ -63,7 +72,9 @@ def construct_segmentation_payload(
6372

6473

6574
def construct_box_predictions_payload(
66-
box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
75+
box_predictions: List[
76+
Union[BoxPrediction, PolygonPrediction, CuboidPrediction]
77+
],
6778
update: bool,
6879
) -> dict:
6980
predictions = []

0 commit comments

Comments
 (0)