Skip to content

Commit 31a802d

Browse files
Drew KaulDrew Kaul
Drew Kaul
authored and
Drew Kaul
committed
reformat
1 parent a1f87ef commit 31a802d

File tree

6 files changed

+71
-30
lines changed

6 files changed

+71
-30
lines changed

‎nucleus/__init__.py

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -559,21 +559,23 @@ def annotate_dataset(
559559
self,
560560
dataset_id: str,
561561
annotations: List[
562-
Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
562+
Union[
563+
BoxAnnotation,
564+
PolygonAnnotation,
565+
CuboidAnnotation,
566+
SegmentationAnnotation,
567+
]
563568
],
564569
update: bool,
565570
batch_size: int = 5000,
566571
):
567572
"""
568573
Uploads ground truth annotations for a given dataset.
569574
:param dataset_id: id of the dataset
570-
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
575+
:param annotations: List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]]
571576
:param update: whether to update or ignore conflicting annotations
572577
:return: {"dataset_id: str, "annotations_processed": int}
573578
"""
574-
if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
575-
raise NotImplementedError("Cuboid annotations not yet supported")
576-
577579
# Split payload into segmentations and Box/Polygon
578580
segmentations = [
579581
ann
@@ -714,14 +716,19 @@ def predict(
714716
self,
715717
model_run_id: str,
716718
annotations: List[
717-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
719+
Union[
720+
BoxPrediction,
721+
PolygonPrediction,
722+
CuboidPrediction,
723+
SegmentationPrediction,
724+
]
718725
],
719726
update: bool,
720727
batch_size: int = 5000,
721728
):
722729
"""
723730
Uploads model outputs as predictions for a model_run. Returns info about the upload.
724-
:param annotations: List[Union[BoxPrediction, PolygonPrediction]],
731+
:param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
725732
:param update: bool
726733
:return:
727734
{
@@ -731,9 +738,6 @@ def predict(
731738
"predictions_ignored": int,
732739
}
733740
"""
734-
if any((isinstance(ann, CuboidPrediction) for ann in annotations)):
735-
raise NotImplementedError("Cuboid predictions not yet supported")
736-
737741
segmentations = [
738742
ann
739743
for ann in annotations
@@ -885,7 +889,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
885889
:param reference_id: reference_id of a dataset item.
886890
:return:
887891
{
888-
"annotations": List[BoxPrediction],
892+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
889893
}
890894
"""
891895
return self.make_request(
@@ -910,7 +914,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
910914
:param i: absolute number of Dataset Item for a dataset corresponding to the model run.
911915
:return:
912916
{
913-
"annotations": List[BoxPrediction],
917+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
914918
}
915919
"""
916920
return self.make_request(
@@ -939,7 +943,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
939943
:param dataset_item_id: dataset_item_id of a dataset item.
940944
:return:
941945
{
942-
"annotations": List[BoxPrediction],
946+
"annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
943947
}
944948
"""
945949
return self.make_request(

‎nucleus/dataset.py

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

1313
from .annotation import (
1414
Annotation,
15-
CuboidAnnotation,
1615
check_all_annotation_paths_remote,
1716
)
1817
from .constants import (
@@ -168,9 +167,6 @@ def annotate(
168167
"ignored_items": int,
169168
}
170169
"""
171-
if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
172-
raise NotImplementedError("Cuboid annotations not yet supported")
173-
174170
if asynchronous:
175171
check_all_annotation_paths_remote(annotations)
176172

@@ -257,7 +253,7 @@ def iloc(self, i: int) -> dict:
257253
:return:
258254
{
259255
"item": DatasetItem,
260-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
256+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
261257
}
262258
"""
263259
response = self._client.dataitem_iloc(self.id, i)
@@ -270,7 +266,7 @@ def refloc(self, reference_id: str) -> dict:
270266
:return:
271267
{
272268
"item": DatasetItem,
273-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
269+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
274270
}
275271
"""
276272
response = self._client.dataitem_ref_id(self.id, reference_id)
@@ -283,7 +279,7 @@ def loc(self, dataset_item_id: str) -> dict:
283279
:return:
284280
{
285281
"item": DatasetItem,
286-
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
282+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
287283
}
288284
"""
289285
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
)
@@ -42,7 +43,12 @@ def create_run(
4243
name: str,
4344
dataset: Dataset,
4445
predictions: List[
45-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
46+
Union[
47+
BoxPrediction,
48+
PolygonPrediction,
49+
CuboidPrediction,
50+
SegmentationPrediction,
51+
]
4652
],
4753
metadata: Optional[Dict] = None,
4854
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
JOB_ID_KEY,
1213
POLYGON_TYPE,
@@ -16,6 +17,7 @@
1617
)
1718
from .prediction import (
1819
BoxPrediction,
20+
CuboidPrediction,
1921
PolygonPrediction,
2022
SegmentationPrediction,
2123
)
@@ -90,14 +92,19 @@ def commit(self, payload: Optional[dict] = None) -> dict:
9092
def predict(
9193
self,
9294
annotations: List[
93-
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
95+
Union[
96+
BoxPrediction,
97+
PolygonPrediction,
98+
CuboidPrediction,
99+
SegmentationPrediction,
100+
]
94101
],
95102
update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
96103
asynchronous: bool = False,
97104
) -> Union[dict, AsyncJob]:
98105
"""
99106
Uploads model outputs as predictions for a model_run. Returns info about the upload.
100-
:param annotations: List[Union[BoxPrediction, PolygonPrediction]],
107+
:param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
101108
:return:
102109
{
103110
"model_run_id": str,
@@ -124,7 +131,7 @@ def iloc(self, i: int):
124131
"""
125132
Returns Model Run Info For Dataset Item by its number.
126133
:param i: absolute number of Dataset Item for a dataset corresponding to the model run.
127-
:return: List[Union[BoxPrediction, PolygonPrediction]],
134+
:return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
128135
}
129136
"""
130137
response = self._client.predictions_iloc(self.model_run_id, i)
@@ -134,7 +141,7 @@ def refloc(self, reference_id: str):
134141
"""
135142
Returns Model Run Info For Dataset Item by its reference_id.
136143
:param reference_id: reference_id of a dataset item.
137-
:return: List[Union[BoxPrediction, PolygonPrediction]],
144+
:return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
138145
"""
139146
response = self._client.predictions_ref_id(
140147
self.model_run_id, reference_id
@@ -159,7 +166,14 @@ def _format_prediction_response(
159166
self, response: dict
160167
) -> Union[
161168
dict,
162-
List[Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]],
169+
List[
170+
Union[
171+
BoxPrediction,
172+
PolygonPrediction,
173+
CuboidPrediction,
174+
SegmentationPrediction,
175+
]
176+
],
163177
]:
164178
annotation_payload = response.get(ANNOTATIONS_KEY, None)
165179
if not annotation_payload:
@@ -171,11 +185,13 @@ def _format_prediction_response(
171185
Union[
172186
Type[BoxPrediction],
173187
Type[PolygonPrediction],
188+
Type[CuboidPrediction],
174189
Type[SegmentationPrediction],
175190
],
176191
] = {
177192
BOX_TYPE: BoxPrediction,
178193
POLYGON_TYPE: PolygonPrediction,
194+
CUBOID_TYPE: CuboidPrediction,
179195
SEGMENTATION_TYPE: SegmentationPrediction,
180196
}
181197
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 = []

‎nucleus/utils.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from nucleus.annotation import (
1212
Annotation,
1313
BoxAnnotation,
14+
CuboidAnnotation,
1415
PolygonAnnotation,
1516
SegmentationAnnotation,
1617
)
@@ -19,13 +20,14 @@
1920
ANNOTATION_TYPES,
2021
ANNOTATIONS_KEY,
2122
BOX_TYPE,
23+
CUBOID_TYPE,
2224
ITEM_KEY,
2325
POLYGON_TYPE,
2426
REFERENCE_ID_KEY,
2527
SEGMENTATION_TYPE,
2628
)
2729
from .dataset_item import DatasetItem
28-
from .prediction import BoxPrediction, PolygonPrediction
30+
from .prediction import BoxPrediction, CuboidPrediction, PolygonPrediction
2931

3032

3133
def _get_all_field_values(metadata_list: List[dict], key: str):
@@ -34,7 +36,10 @@ def _get_all_field_values(metadata_list: List[dict], key: str):
3436

3537
def suggest_metadata_schema(
3638
data: Union[
37-
List[DatasetItem], List[BoxPrediction], List[PolygonPrediction]
39+
List[DatasetItem],
40+
List[BoxPrediction],
41+
List[PolygonPrediction],
42+
List[CuboidPrediction],
3843
]
3944
):
4045
metadata_list: List[dict] = [
@@ -106,6 +111,9 @@ def convert_export_payload(api_payload):
106111
for box in row[BOX_TYPE]:
107112
box[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
108113
annotations[BOX_TYPE].append(BoxAnnotation.from_json(box))
114+
for cuboid in row[CUBOID_TYPE]:
115+
cuboid[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
116+
annotations[CUBOID_TYPE].append(CuboidAnnotation.from_json(cuboid))
109117
return_payload_row[ANNOTATIONS_KEY] = annotations
110118
return_payload.append(return_payload_row)
111119
return return_payload

0 commit comments

Comments
 (0)