Skip to content

Commit e8a5098

Browse files
committed
review + tests
1 parent f85239a commit e8a5098

File tree

3 files changed

+69
-11
lines changed

3 files changed

+69
-11
lines changed

nucleus/constants.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,7 @@
6363
MODEL_RUN_ID_KEY = "model_run_id"
6464
NAME_KEY = "name"
6565
NEW_ITEMS = "new_items"
66-
NUCLEUS_ENDPOINT = "https://jihan-yin-api.devbox.internal.scale.com/v1/nucleus"
67-
# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
66+
NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
6867
NUM_SENSORS_KEY = "num_sensors"
6968
ORIGINAL_IMAGE_URL_KEY = "original_image_url"
7069
POINTCLOUD_KEY = "pointcloud"

nucleus/dataset.py

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -92,8 +92,18 @@ def autotag_items(self, autotag_name, for_scores_greater_than=0):
9292
9393
:return: dictionary of the form
9494
{
95-
'autotagItems': AutotagItem[],
96-
'autotag': Autotag
95+
'autotagItems': {
96+
ref_id: str,
97+
score: float,
98+
model_prediction_id: str | None
99+
ground_truth_annotation_id: str | None,
100+
}[],
101+
'autotag': {
102+
id: str,
103+
name: str,
104+
status: 'started' | 'completed',
105+
autotag_level: 'Image' | 'Object'
106+
}
97107
}
98108
See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-items for more details on the return types.
99109
"""
@@ -109,8 +119,17 @@ def autotag_training_items(self, autotag_name):
109119
110120
:return: dictionary of the form
111121
{
112-
'autotagTrainingItems': AutotagTrainingItem[],
113-
'autotag': Autotag
122+
'autotagPositiveTrainingItems': {
123+
ref_id: str,
124+
model_prediction_id: str | None,
125+
ground_truth_annotation_id: str | None,
126+
}[],
127+
'autotag': {
128+
id: str,
129+
name: str,
130+
status: 'started' | 'completed',
131+
autotag_level: 'Image' | 'Object'
132+
}
114133
}
115134
See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-training-items for more details on the return types.
116135
"""

tests/test_dataset.py

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,7 @@ def test_raises_error_for_duplicate():
329329
)
330330

331331

332-
def test_dataset_export_autotag_scores(CLIENT):
332+
def test_dataset_export_autotag_tagged_items(CLIENT):
333333
# This test can only run for the test user who has an indexed dataset.
334334
# TODO: if/when we can create autotags via api, create one instead.
335335
if NUCLEUS_PYTEST_USER_ID in CLIENT.api_key:
@@ -342,11 +342,51 @@ def test_dataset_export_autotag_scores(CLIENT):
342342
in str(api_error.value)
343343
)
344344

345-
scores = dataset.autotag_scores(autotag_name="PytestTestTag")
345+
items = dataset.autotag_items(autotag_name="PytestTestTag")
346346

347-
for column in ["dataset_item_ids", "ref_ids", "scores"]:
348-
assert column in scores
349-
assert len(scores[column]) > 0
347+
assert "autotagItems" in items
348+
assert "autotag" in items
349+
350+
autotagItems = items["autotagItems"]
351+
autotag = items["autotag"]
352+
353+
assert len(autotagItems) > 0
354+
for item in autotagItems:
355+
for column in ["ref_id", "score"]:
356+
assert column in item
357+
358+
for column in ["id", "name", "status", "autotag_level"]:
359+
assert column in autotag
360+
361+
362+
def test_dataset_export_autotag_training_items(CLIENT):
363+
# This test can only run for the test user who has an indexed dataset.
364+
# TODO: if/when we can create autotags via api, create one instead.
365+
if NUCLEUS_PYTEST_USER_ID in CLIENT.api_key:
366+
dataset = CLIENT.get_dataset(DATASET_WITH_AUTOTAG)
367+
368+
with pytest.raises(NucleusAPIError) as api_error:
369+
dataset.autotag_scores(autotag_name="NONSENSE_GARBAGE")
370+
assert (
371+
f"The autotag NONSENSE_GARBAGE was not found in dataset {DATASET_WITH_AUTOTAG}"
372+
in str(api_error.value)
373+
)
374+
375+
items = dataset.autotag_training_items(autotag_name="PytestTestTag")
376+
377+
assert "autotagItems" in items
378+
assert "autotag" in items
379+
380+
autotagTrainingItems = items["autotagPositiveTrainingItems"]
381+
autotag = items["autotag"]
382+
383+
assert len(autotagTrainingItems) > 0
384+
for item in autotagTrainingItems:
385+
for column in ["ref_id"]:
386+
assert column in item
387+
388+
for column in ["id", "name", "status", "autotag_level"]:
389+
assert column in autotag
350390

351391

352392
@pytest.mark.integration

0 commit comments

Comments
 (0)