Skip to content

Commit 117962f

Browse files
drew-kauldrew-kaul
authored
Merge branch 'master' into implement_scene_class
2 parents e9a681a + af14a2e commit 117962f

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
@@ -616,21 +616,23 @@ def annotate_dataset(
616616
self,
617617
dataset_id: str,
618618
annotations: List[
619-
Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
619+
Union[
620+
BoxAnnotation,
621+
PolygonAnnotation,
622+
CuboidAnnotation,
623+
SegmentationAnnotation,
624+
]
620625
],
621626
update: bool,
622627
batch_size: int = 5000,
623628
):
624629
"""
625630
Uploads ground truth annotations for a given dataset.
626631
:param dataset_id: id of the dataset
627-
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
632+
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]]
628633
:param update: whether to update or ignore conflicting annotations
629634
:return: {"dataset_id: str, "annotations_processed": int}
630635
"""
631-
if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
632-
raise NotImplementedError("Cuboid annotations not yet supported")
633-
634636
# Split payload into segmentations and Box/Polygon
635637
segmentations = [
636638
ann
@@ -771,14 +773,19 @@ def predict(
771773
self,
772774
model_run_id: str,
773775
annotations: List[
774-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
776+
Union[
777+
BoxPrediction,
778+
PolygonPrediction,
779+
CuboidPrediction,
780+
SegmentationPrediction,
781+
]
775782
],
776783
update: bool,
777784
batch_size: int = 5000,
778785
):
779786
"""
780787
Uploads model outputs as predictions for a model_run. Returns info about the upload.
781-
:param annotations: List[Union[BoxPrediction, PolygonPrediction]],
788+
:param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
782789
:param update: bool
783790
:return:
784791
{
@@ -788,9 +795,6 @@ def predict(
788795
"predictions_ignored": int,
789796
}
790797
"""
791-
if any((isinstance(ann, CuboidPrediction) for ann in annotations)):
792-
raise NotImplementedError("Cuboid predictions not yet supported")
793-
794798
segmentations = [
795799
ann
796800
for ann in annotations
@@ -942,7 +946,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
942946
:param reference_id: reference_id of a dataset item.
943947
:return:
944948
{
945-
"annotations": List[BoxPrediction],
949+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
946950
}
947951
"""
948952
return self.make_request(
@@ -967,7 +971,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
967971
:param i: absolute number of Dataset Item for a dataset corresponding to the model run.
968972
:return:
969973
{
970-
"annotations": List[BoxPrediction],
974+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
971975
}
972976
"""
973977
return self.make_request(
@@ -996,7 +1000,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
9961000
:param dataset_item_id: dataset_item_id of a dataset item.
9971001
:return:
9981002
{
999-
"annotations": List[BoxPrediction],
1003+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
10001004
}
10011005
"""
10021006
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
@@ -33,6 +33,7 @@
3333
ERROR_CODES = "error_codes"
3434
ERROR_ITEMS = "upload_errors"
3535
ERROR_PAYLOAD = "error_payload"
36+
FRAMES = "frames"
3637
FRAMES_KEY = "frames"
3738
FX_KEY = "fx"
3839
FY_KEY = "fy"
@@ -70,6 +71,7 @@
7071
REFERENCE_IDS_KEY = "reference_ids"
7172
REFERENCE_ID_KEY = "reference_id"
7273
REQUEST_ID_KEY = "requestId"
74+
SCENES = "scenes"
7375
SEGMENTATIONS_KEY = "segmentations"
7476
SLICE_ID_KEY = "slice_id"
7577
STATUS_CODE_KEY = "status_code"
@@ -78,6 +80,7 @@
7880
TYPE_KEY = "type"
7981
UPDATED_ITEMS = "updated_items"
8082
UPDATE_KEY = "update"
83+
URL = "url"
8184
URL_KEY = "url"
8285
VERTICES_KEY = "vertices"
8386
WIDTH_KEY = "width"

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