Skip to content

Commit 3619a54

Browse files
author
Diego Ardila
committed
Checking in load test script
1 parent 7bde819 commit 3619a54

File tree

9 files changed

+222
-117
lines changed

9 files changed

+222
-117
lines changed

‎nucleus/__init__.py

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,6 @@
8282
ANNOTATIONS_PROCESSED_KEY,
8383
AUTOTAGS_KEY,
8484
DATASET_ID_KEY,
85-
DATASET_ITEM_IDS_KEY,
8685
DEFAULT_NETWORK_TIMEOUT_SEC,
8786
EMBEDDING_DIMENSION_KEY,
8887
EMBEDDINGS_URL_KEY,
@@ -1111,8 +1110,7 @@ def delete_annotations(
11111110
def append_to_slice(
11121111
self,
11131112
slice_id: str,
1114-
dataset_item_ids: List[str] = None,
1115-
reference_ids: List[str] = None,
1113+
reference_ids: List[str],
11161114
) -> dict:
11171115
"""
11181116
Appends to a slice from items already present in a dataset.
@@ -1128,18 +1126,10 @@ def append_to_slice(
11281126
"slice_id": str,
11291127
}
11301128
"""
1131-
if dataset_item_ids and reference_ids:
1132-
raise Exception(
1133-
"You cannot specify both dataset_item_ids and reference_ids"
1134-
)
1135-
1136-
ids_to_append: Dict[str, Any] = {}
1137-
if dataset_item_ids:
1138-
ids_to_append[DATASET_ITEM_IDS_KEY] = dataset_item_ids
1139-
if reference_ids:
1140-
ids_to_append[REFERENCE_IDS_KEY] = reference_ids
11411129

1142-
response = self.make_request(ids_to_append, f"slice/{slice_id}/append")
1130+
response = self.make_request(
1131+
{REFERENCE_IDS_KEY: reference_ids}, f"slice/{slice_id}/append"
1132+
)
11431133
return response
11441134

11451135
def list_autotags(self, dataset_id: str) -> List[str]:

‎nucleus/annotation.py

Lines changed: 15 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,10 @@
99
ANNOTATIONS_KEY,
1010
BOX_TYPE,
1111
CUBOID_TYPE,
12-
DATASET_ITEM_ID_KEY,
1312
DIMENSIONS_KEY,
1413
GEOMETRY_KEY,
1514
HEIGHT_KEY,
1615
INDEX_KEY,
17-
ITEM_ID_KEY,
1816
LABEL_KEY,
1917
MASK_TYPE,
2018
MASK_URL_KEY,
@@ -33,16 +31,7 @@
3331

3432

3533
class Annotation:
36-
reference_id: Optional[str] = None
37-
item_id: Optional[str] = None
38-
39-
def _check_ids(self):
40-
if self.reference_id and self.item_id:
41-
self.item_id = None # Prefer reference id to item id.
42-
if not (self.reference_id or self.item_id):
43-
raise Exception(
44-
"You must specify either a reference_id or an item_id for an annotation."
45-
)
34+
reference_id: str
4635

4736
@classmethod
4837
def from_json(cls, payload: dict):
@@ -93,14 +82,12 @@ def to_payload(self) -> dict:
9382
class SegmentationAnnotation(Annotation):
9483
mask_url: str
9584
annotations: List[Segment]
85+
reference_id: str
9686
annotation_id: Optional[str] = None
97-
reference_id: Optional[str] = None
98-
item_id: Optional[str] = None
9987

10088
def __post_init__(self):
10189
if not self.mask_url:
10290
raise Exception("You must specify a mask_url.")
103-
self._check_ids()
10491

10592
@classmethod
10693
def from_json(cls, payload: dict):
@@ -112,8 +99,7 @@ def from_json(cls, payload: dict):
11299
Segment.from_json(ann)
113100
for ann in payload.get(ANNOTATIONS_KEY, [])
114101
],
115-
reference_id=payload.get(REFERENCE_ID_KEY, None),
116-
item_id=payload.get(ITEM_ID_KEY, None),
102+
reference_id=payload[REFERENCE_ID_KEY],
117103
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
118104
)
119105

@@ -124,10 +110,9 @@ def to_payload(self) -> dict:
124110
ANNOTATIONS_KEY: [ann.to_payload() for ann in self.annotations],
125111
ANNOTATION_ID_KEY: self.annotation_id,
126112
}
127-
if self.reference_id:
128-
payload[REFERENCE_ID_KEY] = self.reference_id
129-
else:
130-
payload[ITEM_ID_KEY] = self.item_id
113+
114+
payload[REFERENCE_ID_KEY] = self.reference_id
115+
131116
return payload
132117

133118

@@ -144,14 +129,14 @@ class BoxAnnotation(Annotation): # pylint: disable=R0902
144129
y: Union[float, int]
145130
width: Union[float, int]
146131
height: Union[float, int]
147-
reference_id: Optional[str] = None
148-
item_id: Optional[str] = None
132+
reference_id: str
149133
annotation_id: Optional[str] = None
150134
metadata: Optional[Dict] = None
151135

152136
def __post_init__(self):
153-
self._check_ids()
154137
self.metadata = self.metadata if self.metadata else {}
138+
if self.annotation_id is None:
139+
self.annotation_id = f"{self.label}-{self.x}-{self.y}-{self.width}-{self.height}-{self.reference_id}"
155140

156141
@classmethod
157142
def from_json(cls, payload: dict):
@@ -162,8 +147,7 @@ def from_json(cls, payload: dict):
162147
y=geometry.get(Y_KEY, 0),
163148
width=geometry.get(WIDTH_KEY, 0),
164149
height=geometry.get(HEIGHT_KEY, 0),
165-
reference_id=payload.get(REFERENCE_ID_KEY, None),
166-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
150+
reference_id=payload[REFERENCE_ID_KEY],
167151
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
168152
metadata=payload.get(METADATA_KEY, {}),
169153
)
@@ -215,13 +199,11 @@ def to_payload(self) -> dict:
215199
class PolygonAnnotation(Annotation):
216200
label: str
217201
vertices: List[Point]
218-
reference_id: Optional[str] = None
219-
item_id: Optional[str] = None
202+
reference_id: str
220203
annotation_id: Optional[str] = None
221204
metadata: Optional[Dict] = None
222205

223206
def __post_init__(self):
224-
self._check_ids()
225207
self.metadata = self.metadata if self.metadata else {}
226208
if len(self.vertices) > 0:
227209
if not hasattr(self.vertices[0], X_KEY) or not hasattr(
@@ -245,8 +227,7 @@ def from_json(cls, payload: dict):
245227
vertices=[
246228
Point.from_json(_) for _ in geometry.get(VERTICES_KEY, [])
247229
],
248-
reference_id=payload.get(REFERENCE_ID_KEY, None),
249-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
230+
reference_id=payload[REFERENCE_ID_KEY],
250231
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
251232
metadata=payload.get(METADATA_KEY, {}),
252233
)
@@ -271,13 +252,11 @@ class CuboidAnnotation(Annotation): # pylint: disable=R0902
271252
position: Point3D
272253
dimensions: Point3D
273254
yaw: float
274-
reference_id: Optional[str] = None
275-
item_id: Optional[str] = None
255+
reference_id: str
276256
annotation_id: Optional[str] = None
277257
metadata: Optional[Dict] = None
278258

279259
def __post_init__(self):
280-
self._check_ids()
281260
self.metadata = self.metadata if self.metadata else {}
282261

283262
@classmethod
@@ -288,8 +267,7 @@ def from_json(cls, payload: dict):
288267
position=Point3D.from_json(geometry.get(POSITION_KEY, {})),
289268
dimensions=Point3D.from_json(geometry.get(DIMENSIONS_KEY, {})),
290269
yaw=geometry.get(YAW_KEY, 0),
291-
reference_id=payload.get(REFERENCE_ID_KEY, None),
292-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
270+
reference_id=payload[REFERENCE_ID_KEY],
293271
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
294272
metadata=payload.get(METADATA_KEY, {}),
295273
)
@@ -304,8 +282,7 @@ def to_payload(self) -> dict:
304282
YAW_KEY: self.yaw,
305283
},
306284
}
307-
if self.reference_id:
308-
payload[REFERENCE_ID_KEY] = self.reference_id
285+
payload[REFERENCE_ID_KEY] = self.reference_id
309286
if self.annotation_id:
310287
payload[ANNOTATION_ID_KEY] = self.annotation_id
311288
if self.metadata:

‎nucleus/constants.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818
CX_KEY = "cx"
1919
CY_KEY = "cy"
2020
DATASET_ID_KEY = "dataset_id"
21-
DATASET_ITEM_IDS_KEY = "dataset_item_ids"
22-
DATASET_ITEM_ID_KEY = "dataset_item_id"
2321
DATASET_LENGTH_KEY = "length"
2422
DATASET_MODEL_RUNS_KEY = "model_run_ids"
2523
DATASET_NAME_KEY = "name"
@@ -46,7 +44,6 @@
4644
INDEX_KEY = "index"
4745
INDEX_CONTINUOUS_ENABLE_KEY = "enable"
4846
ITEMS_KEY = "items"
49-
ITEM_ID_KEY = "item_id"
5047
ITEM_KEY = "item"
5148
ITEMS_KEY = "items"
5249
ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"

‎nucleus/dataset.py

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Dict, List, Optional, Union
1+
from typing import Any, Dict, List, Optional, Sequence, Union
22

33
import requests
44

@@ -15,7 +15,6 @@
1515
check_all_mask_paths_remote,
1616
)
1717
from .constants import (
18-
DATASET_ITEM_IDS_KEY,
1918
DATASET_LENGTH_KEY,
2019
DATASET_MODEL_RUNS_KEY,
2120
DATASET_NAME_KEY,
@@ -154,7 +153,7 @@ def create_model_run(
154153

155154
def annotate(
156155
self,
157-
annotations: List[Annotation],
156+
annotations: Sequence[Annotation],
158157
update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
159158
batch_size: int = 5000,
160159
asynchronous: bool = False,
@@ -200,7 +199,7 @@ def ingest_tasks(self, task_ids: dict):
200199

201200
def append(
202201
self,
203-
items: Union[List[DatasetItem], List[LidarScene]],
202+
items: Union[Sequence[DatasetItem], Sequence[LidarScene]],
204203
update: Optional[bool] = False,
205204
batch_size: Optional[int] = 20,
206205
asynchronous=False,
@@ -353,8 +352,7 @@ def loc(self, dataset_item_id: str) -> dict:
353352
def create_slice(
354353
self,
355354
name: str,
356-
dataset_item_ids: List[str] = None,
357-
reference_ids: List[str] = None,
355+
reference_ids: List[str],
358356
):
359357
"""
360358
Creates a slice from items already present in a dataset.
@@ -367,16 +365,9 @@ def create_slice(
367365
368366
:return: new Slice object
369367
"""
370-
if bool(dataset_item_ids) == bool(reference_ids):
371-
raise Exception(
372-
"You must specify exactly one of dataset_item_ids or reference_ids."
373-
)
374-
payload: Dict[str, Any] = {NAME_KEY: name}
375-
if dataset_item_ids:
376-
payload[DATASET_ITEM_IDS_KEY] = dataset_item_ids
377-
if reference_ids:
378-
payload[REFERENCE_IDS_KEY] = reference_ids
379-
return self._client.create_slice(self.id, payload)
368+
return self._client.create_slice(
369+
self.id, {NAME_KEY: name, REFERENCE_IDS_KEY: reference_ids}
370+
)
380371

381372
def delete_item(self, item_id: str = None, reference_id: str = None):
382373
if bool(item_id) == bool(reference_id):

‎nucleus/prediction.py

Lines changed: 7 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
)
1111
from .constants import (
1212
ANNOTATION_ID_KEY,
13-
DATASET_ITEM_ID_KEY,
1413
REFERENCE_ID_KEY,
1514
METADATA_KEY,
1615
GEOMETRY_KEY,
@@ -23,7 +22,6 @@
2322
CONFIDENCE_KEY,
2423
VERTICES_KEY,
2524
ANNOTATIONS_KEY,
26-
ITEM_ID_KEY,
2725
MASK_URL_KEY,
2826
POSITION_KEY,
2927
DIMENSIONS_KEY,
@@ -42,8 +40,7 @@ def from_json(cls, payload: dict):
4240
Segment.from_json(ann)
4341
for ann in payload.get(ANNOTATIONS_KEY, [])
4442
],
45-
reference_id=payload.get(REFERENCE_ID_KEY, None),
46-
item_id=payload.get(ITEM_ID_KEY, None),
43+
reference_id=payload[REFERENCE_ID_KEY],
4744
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
4845
)
4946

@@ -56,8 +53,7 @@ def __init__(
5653
y: int,
5754
width: int,
5855
height: int,
59-
reference_id: Optional[str] = None,
60-
item_id: Optional[str] = None,
56+
reference_id: str,
6157
confidence: Optional[float] = None,
6258
annotation_id: Optional[str] = None,
6359
metadata: Optional[Dict] = None,
@@ -70,7 +66,6 @@ def __init__(
7066
width=width,
7167
height=height,
7268
reference_id=reference_id,
73-
item_id=item_id,
7469
annotation_id=annotation_id,
7570
metadata=metadata,
7671
)
@@ -95,8 +90,7 @@ def from_json(cls, payload: dict):
9590
y=geometry.get(Y_KEY, 0),
9691
width=geometry.get(WIDTH_KEY, 0),
9792
height=geometry.get(HEIGHT_KEY, 0),
98-
reference_id=payload.get(REFERENCE_ID_KEY, None),
99-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
93+
reference_id=payload[REFERENCE_ID_KEY],
10094
confidence=payload.get(CONFIDENCE_KEY, None),
10195
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
10296
metadata=payload.get(METADATA_KEY, {}),
@@ -109,8 +103,7 @@ def __init__(
109103
self,
110104
label: str,
111105
vertices: List[Point],
112-
reference_id: Optional[str] = None,
113-
item_id: Optional[str] = None,
106+
reference_id: str,
114107
confidence: Optional[float] = None,
115108
annotation_id: Optional[str] = None,
116109
metadata: Optional[Dict] = None,
@@ -120,7 +113,6 @@ def __init__(
120113
label=label,
121114
vertices=vertices,
122115
reference_id=reference_id,
123-
item_id=item_id,
124116
annotation_id=annotation_id,
125117
metadata=metadata,
126118
)
@@ -144,8 +136,7 @@ def from_json(cls, payload: dict):
144136
vertices=[
145137
Point.from_json(_) for _ in geometry.get(VERTICES_KEY, [])
146138
],
147-
reference_id=payload.get(REFERENCE_ID_KEY, None),
148-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
139+
reference_id=payload[REFERENCE_ID_KEY],
149140
confidence=payload.get(CONFIDENCE_KEY, None),
150141
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
151142
metadata=payload.get(METADATA_KEY, {}),
@@ -160,8 +151,7 @@ def __init__(
160151
position: Point3D,
161152
dimensions: Point3D,
162153
yaw: float,
163-
reference_id: Optional[str] = None,
164-
item_id: Optional[str] = None,
154+
reference_id: str,
165155
confidence: Optional[float] = None,
166156
annotation_id: Optional[str] = None,
167157
metadata: Optional[Dict] = None,
@@ -173,7 +163,6 @@ def __init__(
173163
dimensions=dimensions,
174164
yaw=yaw,
175165
reference_id=reference_id,
176-
item_id=item_id,
177166
annotation_id=annotation_id,
178167
metadata=metadata,
179168
)
@@ -197,8 +186,7 @@ def from_json(cls, payload: dict):
197186
position=Point3D.from_json(geometry.get(POSITION_KEY, {})),
198187
dimensions=Point3D.from_json(geometry.get(DIMENSIONS_KEY, {})),
199188
yaw=geometry.get(YAW_KEY, 0),
200-
reference_id=payload.get(REFERENCE_ID_KEY, None),
201-
item_id=payload.get(DATASET_ITEM_ID_KEY, None),
189+
reference_id=payload[REFERENCE_ID_KEY],
202190
confidence=payload.get(CONFIDENCE_KEY, None),
203191
annotation_id=payload.get(ANNOTATION_ID_KEY, None),
204192
metadata=payload.get(METADATA_KEY, {}),

0 commit comments

Comments
 (0)