Skip to content

Commit 527ab6a

Browse files
sasha-scalesasha-scale
committed
add pytest unit tests
1 parent e688fdb commit 527ab6a

File tree

5 files changed

+122
-43
lines changed

5 files changed

+122
-43
lines changed

‎nucleus/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
BoxAnnotation,
7474
PolygonAnnotation,
7575
SegmentationAnnotation,
76+
Segment,
7677
)
7778
from .prediction import (
7879
BoxPrediction,

‎nucleus/constants.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
1+
# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
2+
NUCLEUS_ENDPOINT = "http://localhost:3000/v1/nucleus"
23
ITEMS_KEY = "items"
34
ITEM_KEY = "item"
45
REFERENCE_ID_KEY = "reference_id"

‎nucleus/model_run.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
from typing import Optional, List, Dict, Any, Union
22
from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
3-
from .prediction import BoxPrediction, PolygonPrediction
3+
from .prediction import (
4+
BoxPrediction,
5+
PolygonPrediction,
6+
SegmentationPrediction,
7+
)
48
from .payload_constructor import construct_box_predictions_payload
59

610

@@ -62,7 +66,9 @@ def commit(self, payload: Optional[dict] = None) -> dict:
6266

6367
def predict(
6468
self,
65-
annotations: List[Union[BoxPrediction, PolygonPrediction]],
69+
annotations: List[
70+
Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
71+
],
6672
update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
6773
) -> dict:
6874
"""

‎tests/helpers.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,8 @@ def reference_id_from_url(url):
8888
for i in range(len(TEST_IMG_URLS))
8989
]
9090

91+
TEST_SEGMENTATION_PREDICTIONS = TEST_SEGMENTATION_ANNOTATIONS
92+
9193
TEST_BOX_PREDICTIONS = [
9294
{**TEST_BOX_ANNOTATIONS[i], "confidence": 0.10 * i}
9395
for i in range(len(TEST_BOX_ANNOTATIONS))

‎tests/test_prediction.py

Lines changed: 109 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -8,36 +8,41 @@
88
TEST_IMG_URLS,
99
TEST_BOX_PREDICTIONS,
1010
TEST_POLYGON_PREDICTIONS,
11+
TEST_SEGMENTATION_PREDICTIONS,
1112
reference_id_from_url,
1213
assert_box_prediction_matches_dict,
1314
assert_polygon_prediction_matches_dict,
1415
)
1516

16-
from nucleus import BoxPrediction, PolygonPrediction, DatasetItem
17+
from nucleus import (
18+
BoxPrediction,
19+
PolygonPrediction,
20+
SegmentationPrediction,
21+
DatasetItem,
22+
)
1723
from nucleus.constants import ERROR_PAYLOAD
1824

25+
1926
@pytest.fixture()
2027
def model_run(CLIENT):
2128
ds = CLIENT.create_dataset(TEST_DATASET_NAME)
2229
ds_items = []
2330
for url in TEST_IMG_URLS:
24-
ds_items.append(DatasetItem(
25-
image_location=url,
26-
reference_id=reference_id_from_url(url),
27-
))
31+
ds_items.append(
32+
DatasetItem(
33+
image_location=url,
34+
reference_id=reference_id_from_url(url),
35+
)
36+
)
2837

2938
response = ds.append(ds_items)
3039
assert ERROR_PAYLOAD not in response.json()
3140

3241
model = CLIENT.add_model(
33-
name=TEST_MODEL_NAME,
34-
reference_id=TEST_MODEL_REFERENCE
42+
name=TEST_MODEL_NAME, reference_id=TEST_MODEL_REFERENCE
3543
)
3644

37-
run = model.create_run(
38-
name=TEST_MODEL_RUN,
39-
dataset=ds,
40-
predictions=[])
45+
run = model.create_run(name=TEST_MODEL_RUN, dataset=ds, predictions=[])
4146

4247
yield run
4348

@@ -46,13 +51,14 @@ def model_run(CLIENT):
4651
response = CLIENT.delete_model(model.id)
4752
assert response == {}
4853

54+
4955
def test_box_pred_upload(model_run):
5056
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
5157
response = model_run.predict(annotations=[prediction])
5258

53-
assert response['model_run_id'] == model_run.model_run_id
54-
assert response['predictions_processed'] == 1
55-
assert response['predictions_ignored'] == 0
59+
assert response["model_run_id"] == model_run.model_run_id
60+
assert response["predictions_processed"] == 1
61+
assert response["predictions_ignored"] == 0
5662

5763
response = model_run.refloc(prediction.reference_id)
5864
assert len(response) == 1
@@ -63,31 +69,63 @@ def test_polygon_pred_upload(model_run):
6369
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
6470
response = model_run.predict(annotations=[prediction])
6571

66-
assert response['model_run_id'] == model_run.model_run_id
67-
assert response['predictions_ignored'] == 0
68-
assert response['predictions_ignored'] == 0
72+
assert response["model_run_id"] == model_run.model_run_id
73+
assert response["predictions_ignored"] == 0
74+
assert response["predictions_ignored"] == 0
6975

7076
response = model_run.refloc(prediction.reference_id)
7177
assert len(response) == 1
72-
assert_polygon_prediction_matches_dict(response[0], TEST_POLYGON_PREDICTIONS[0])
78+
assert_polygon_prediction_matches_dict(
79+
response[0], TEST_POLYGON_PREDICTIONS[0]
80+
)
81+
82+
83+
def test_segmentation_pred_upload(model_run):
84+
prediction = SegmentationPrediction.from_json(
85+
TEST_SEGMENTATION_PREDICTIONS[0]
86+
)
87+
response = model_run.predict(annotations=[prediction])
88+
89+
assert response["model_run_id"] == model_run.model_run_id
90+
assert response["predictions_processed"] == 1
91+
assert response["predictions_ignored"] == 0
92+
93+
94+
def test_segmentation_pred_upload_ignore(model_run):
95+
prediction = SegmentationPrediction.from_json(
96+
TEST_SEGMENTATION_PREDICTIONS[0]
97+
)
98+
response1 = model_run.predict(annotations=[prediction])
99+
100+
assert response1["predictions_processed"] == 1
101+
102+
# Upload Duplicate annotation
103+
response = model_run.predict(annotations=[prediction])
104+
assert response["model_run_id"] == model_run.model_run_id
105+
assert response["predictions_processed"] == 0
106+
assert response["predictions_ignored"] == 1
73107

74108

75109
def test_box_pred_upload_update(model_run):
76110
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
77111
response = model_run.predict(annotations=[prediction])
78112

79-
assert response['predictions_processed'] == 1
113+
assert response["predictions_processed"] == 1
80114

81115
# Copy so we don't modify the original.
82116
prediction_update_params = dict(TEST_BOX_PREDICTIONS[1])
83-
prediction_update_params['annotation_id'] = TEST_BOX_PREDICTIONS[0]['annotation_id']
84-
prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
117+
prediction_update_params["annotation_id"] = TEST_BOX_PREDICTIONS[0][
118+
"annotation_id"
119+
]
120+
prediction_update_params["reference_id"] = TEST_BOX_PREDICTIONS[0][
121+
"reference_id"
122+
]
85123

86124
prediction_update = BoxPrediction(**prediction_update_params)
87125
response = model_run.predict(annotations=[prediction_update], update=True)
88126

89-
assert response['predictions_processed'] == 1
90-
assert response['predictions_ignored'] == 0
127+
assert response["predictions_processed"] == 1
128+
assert response["predictions_ignored"] == 0
91129

92130
response = model_run.refloc(prediction.reference_id)
93131
assert len(response) == 1
@@ -98,18 +136,22 @@ def test_box_pred_upload_ignore(model_run):
98136
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
99137
response = model_run.predict(annotations=[prediction])
100138

101-
assert response['predictions_processed'] == 1
139+
assert response["predictions_processed"] == 1
102140

103141
# Copy so we don't modify the original.
104142
prediction_update_params = dict(TEST_BOX_PREDICTIONS[1])
105-
prediction_update_params['annotation_id'] = TEST_BOX_PREDICTIONS[0]['annotation_id']
106-
prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
143+
prediction_update_params["annotation_id"] = TEST_BOX_PREDICTIONS[0][
144+
"annotation_id"
145+
]
146+
prediction_update_params["reference_id"] = TEST_BOX_PREDICTIONS[0][
147+
"reference_id"
148+
]
107149
prediction_update = BoxPrediction(**prediction_update_params)
108150
# Default behavior is ignore.
109151
response = model_run.predict(annotations=[prediction_update])
110152

111-
assert response['predictions_processed'] == 1
112-
assert response['predictions_ignored'] == 1
153+
assert response["predictions_processed"] == 1
154+
assert response["predictions_ignored"] == 1
113155

114156
response = model_run.refloc(prediction.reference_id)
115157
assert len(response) == 1
@@ -120,42 +162,69 @@ def test_polygon_pred_upload_update(model_run):
120162
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
121163
response = model_run.predict(annotations=[prediction])
122164

123-
assert response['predictions_processed'] == 1
165+
assert response["predictions_processed"] == 1
124166

125167
# Copy so we don't modify the original.
126168
prediction_update_params = dict(TEST_POLYGON_PREDICTIONS[1])
127-
prediction_update_params['annotation_id'] = TEST_POLYGON_PREDICTIONS[0]['annotation_id']
128-
prediction_update_params['reference_id'] = TEST_POLYGON_PREDICTIONS[0]['reference_id']
169+
prediction_update_params["annotation_id"] = TEST_POLYGON_PREDICTIONS[0][
170+
"annotation_id"
171+
]
172+
prediction_update_params["reference_id"] = TEST_POLYGON_PREDICTIONS[0][
173+
"reference_id"
174+
]
129175

130176
prediction_update = PolygonPrediction(**prediction_update_params)
131177
response = model_run.predict(annotations=[prediction_update], update=True)
132178

133-
assert response['predictions_processed'] == 1
134-
assert response['predictions_ignored'] == 0
179+
assert response["predictions_processed"] == 1
180+
assert response["predictions_ignored"] == 0
135181

136182
response = model_run.refloc(prediction.reference_id)
137183
assert len(response) == 1
138-
assert_polygon_prediction_matches_dict(response[0], prediction_update_params)
184+
assert_polygon_prediction_matches_dict(
185+
response[0], prediction_update_params
186+
)
139187

140188

141189
def test_polygon_pred_upload_ignore(model_run):
142190
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
143191
response = model_run.predict(annotations=[prediction])
144192

145-
assert response['predictions_processed'] == 1
193+
assert response["predictions_processed"] == 1
146194

147195
# Copy so we don't modify the original.
148196
prediction_update_params = dict(TEST_POLYGON_PREDICTIONS[1])
149-
prediction_update_params['annotation_id'] = TEST_POLYGON_PREDICTIONS[0]['annotation_id']
150-
prediction_update_params['reference_id'] = TEST_POLYGON_PREDICTIONS[0]['reference_id']
197+
prediction_update_params["annotation_id"] = TEST_POLYGON_PREDICTIONS[0][
198+
"annotation_id"
199+
]
200+
prediction_update_params["reference_id"] = TEST_POLYGON_PREDICTIONS[0][
201+
"reference_id"
202+
]
151203

152204
prediction_update = PolygonPrediction(**prediction_update_params)
153205
# Default behavior is ignore.
154206
response = model_run.predict(annotations=[prediction_update])
155207

156-
assert response['predictions_processed'] == 1
157-
assert response['predictions_ignored'] == 1
208+
assert response["predictions_processed"] == 1
209+
assert response["predictions_ignored"] == 1
158210

159211
response = model_run.refloc(prediction.reference_id)
160212
assert len(response) == 1
161-
assert_polygon_prediction_matches_dict(response[0], TEST_POLYGON_PREDICTIONS[0])
213+
assert_polygon_prediction_matches_dict(
214+
response[0], TEST_POLYGON_PREDICTIONS[0]
215+
)
216+
217+
218+
def test_mixed_pred_upload(model_run):
219+
prediction_semseg = SegmentationPrediction.from_json(
220+
TEST_SEGMENTATION_PREDICTIONS[0]
221+
)
222+
prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
223+
prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
224+
response = model_run.predict(
225+
annotations=[prediction_semseg, prediction_polygon, prediction_bbox]
226+
)
227+
228+
assert response["model_run_id"] == model_run.model_run_id
229+
assert response["predictions_processed"] == 3
230+
assert response["predictions_ignored"] == 0

0 commit comments

Comments
 (0)