Skip to content

Commit 7e13c1e

Browse files
jihan-yinjihan-yin
committed
stash
1 parent 7e64db6 commit 7e13c1e

File tree

5 files changed

+210
-19
lines changed

5 files changed

+210
-19
lines changed

‎nucleus/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -579,7 +579,7 @@ def predict(
579579
Uploads model outputs as predictions for a model_run. Returns info about the upload.
580580
:param payload:
581581
{
582-
"annotations": List[Box2DPrediction],
582+
"annotations": List[Union[Box2DPrediction, Polygon2DPrediction]],
583583
}
584584
:return:
585585
{

‎setup.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
import setuptools
2+
import sys
3+
4+
if sys.version_info < (3,5):
5+
sys.exit('Sorry, Python < 3.5 is not supported')
26

37
with open("README.md", "r") as fh:
48
long_description = fh.read()

‎tests/helpers.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
from pathlib import Path
22

3+
TEST_MODEL_NAME = '[PyTest] Test Model'
4+
TEST_MODEL_REFERENCE = '[PyTest] Test Model Reference'
5+
TEST_MODEL_RUN = '[PyTest] Test Model Run'
36
TEST_DATASET_NAME = '[PyTest] Test Dataset'
47
TEST_IMG_URLS = [
58
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg',
@@ -40,3 +43,47 @@ def reference_id_from_url(url):
4043
}
4144
for i in range(len(TEST_IMG_URLS))
4245
]
46+
47+
TEST_BOX_PREDICTIONS = [
48+
{
49+
**TEST_BOX_ANNOTATIONS[i],
50+
'confidence': 0.10 * i
51+
}
52+
for i in range(len(TEST_BOX_ANNOTATIONS))
53+
]
54+
55+
TEST_POLYGON_PREDICTIONS = [
56+
{
57+
**TEST_POLYGON_ANNOTATIONS[i],
58+
'confidence': 0.10 * i
59+
}
60+
for i in range(len(TEST_POLYGON_ANNOTATIONS))
61+
]
62+
63+
64+
# Asserts that a box annotation instance matches a dict representing its properties.
65+
# Useful to check annotation uploads/updates match.
66+
def assert_box_annotation_matches_dict(annotation_instance, annotation_dict):
67+
assert annotation_instance.label == annotation_dict['label']
68+
assert annotation_instance.x == annotation_dict['x']
69+
assert annotation_instance.y == annotation_dict['y']
70+
assert annotation_instance.height == annotation_dict['height']
71+
assert annotation_instance.width == annotation_dict['width']
72+
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
73+
74+
def assert_polygon_annotation_matches_dict(annotation_instance, annotation_dict):
75+
assert annotation_instance.label == annotation_dict['label']
76+
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
77+
for instance_pt, dict_pt in zip(annotation_instance.vertices, annotation_dict['vertices']):
78+
assert instance_pt['x'] == dict_pt['x']
79+
assert instance_pt['y'] == dict_pt['y']
80+
81+
# Asserts that a box prediction instance matches a dict representing its properties.
82+
# Useful to check prediction uploads/updates match.
83+
def assert_box_prediction_matches_dict(prediction_instance, prediction_dict):
84+
assert_box_annotation_matches_dict(prediction_instance, prediction_dict)
85+
assert prediction_instance.confidence == prediction_dict['confidence']
86+
87+
def assert_polygon_prediction_matches_dict(prediction_instance, prediction_dict):
88+
assert_polygon_annotation_matches_dict(prediction_instance, prediction_dict)
89+
assert prediction_instance.confidence == prediction_dict['confidence']

‎tests/test_annotation.py

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,29 +6,13 @@
66
TEST_BOX_ANNOTATIONS,
77
TEST_POLYGON_ANNOTATIONS,
88
reference_id_from_url,
9+
assert_box_annotation_matches_dict,
10+
assert_polygon_annotation_matches_dict,
911
)
1012

1113
from nucleus import BoxAnnotation, PolygonAnnotation, DatasetItem
1214
from nucleus.constants import ERROR_PAYLOAD
1315

14-
15-
# Asserts that a box annotation instance matches a dict representing its properties.
16-
# Useful to check annotation uploads/updates match.
17-
def assert_box_annotation_matches_dict(annotation_instance, annotation_dict):
18-
assert annotation_instance.label == annotation_dict['label']
19-
assert annotation_instance.x == annotation_dict['x']
20-
assert annotation_instance.y == annotation_dict['y']
21-
assert annotation_instance.height == annotation_dict['height']
22-
assert annotation_instance.width == annotation_dict['width']
23-
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
24-
25-
def assert_polygon_annotation_matches_dict(annotation_instance, annotation_dict):
26-
assert annotation_instance.label == annotation_dict['label']
27-
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
28-
for instance_pt, dict_pt in zip(annotation_instance.vertices, annotation_dict['vertices']):
29-
assert instance_pt['x'] == dict_pt['x']
30-
assert instance_pt['y'] == dict_pt['y']
31-
3216
@pytest.fixture()
3317
def dataset(CLIENT):
3418
ds = CLIENT.create_dataset(TEST_DATASET_NAME)

‎tests/test_prediction.py

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
import pytest
2+
3+
from helpers import (
4+
TEST_DATASET_NAME,
5+
TEST_MODEL_NAME,
6+
TEST_MODEL_REFERENCE,
7+
TEST_IMG_URLS,
8+
TEST_BOX_PREDICTIONS,
9+
TEST_POLYGON_PREDICTIONS,
10+
reference_id_from_url,
11+
assert_box_prediction_matches_dict,
12+
assert_polygon_prediction_matches_dict,
13+
)
14+
15+
from nucleus import BoxPrediction, PolygonPrediction, DatasetItem
16+
from nucleus.constants import ERROR_PAYLOAD
17+
18+
@pytest.fixture()
19+
def dataset(CLIENT):
20+
ds = CLIENT.create_dataset(TEST_DATASET_NAME)
21+
ds_items = []
22+
for url in TEST_IMG_URLS:
23+
ds_items.append(DatasetItem(
24+
image_location=url,
25+
reference_id=reference_id_from_url(url),
26+
))
27+
28+
response = ds.append(ds_items)
29+
assert ERROR_PAYLOAD not in response.json()
30+
31+
model = CLIENT.add_model(
32+
name=TEST_MODEL_NAME,
33+
reference_id=TEST_MODEL_NAME
34+
)
35+
36+
37+
yield ds
38+
39+
response = CLIENT.delete_dataset(ds.id)
40+
assert response == {}
41+
42+
43+
def test_box_pred_upload(dataset):
44+
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
45+
response = dataset.annotate(predictions=[prediction])
46+
47+
assert response['dataset_id'] == dataset.id
48+
assert response['predictions_processed'] == 1
49+
50+
response = dataset.refloc(prediction.reference_id)['predictions']
51+
assert len(response) == 1
52+
response_prediction = response[0]
53+
assert_box_prediction_matches_dict(response_prediction, TEST_BOX_PREDICTIONS[0])
54+
55+
56+
def test_polygon_pred_upload(dataset):
57+
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
58+
response = dataset.annotate(predictions=[prediction])
59+
60+
assert response['dataset_id'] == dataset.id
61+
assert response['predictions_processed'] == 1
62+
63+
response = dataset.refloc(prediction.reference_id)['predictions']
64+
assert len(response) == 1
65+
response_prediction = response[0]
66+
print(response_prediction)
67+
assert_polygon_prediction_matches_dict(response_prediction, TEST_POLYGON_PREDICTIONS[0])
68+
69+
70+
def test_box_pred_upload_update(dataset):
71+
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
72+
response = dataset.annotate(predictions=[prediction])
73+
74+
assert response['predictions_processed'] == 1
75+
76+
# Copy so we don't modify the original.
77+
prediction_update_params = dict(TEST_BOX_PREDICTIONS[1])
78+
prediction_update_params['annotation_id'] = TEST_BOX_PREDICTIONS[0]['annotation_id']
79+
prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
80+
81+
prediction_update = BoxPrediction(**prediction_update_params)
82+
response = dataset.annotate(predictions=[prediction_update], update=True)
83+
84+
assert response['predictions_processed'] == 1
85+
86+
response = dataset.refloc(prediction.reference_id)['predictions']
87+
assert len(response) == 1
88+
response_prediction = response[0]
89+
assert_box_prediction_matches_dict(response_prediction, prediction_update_params)
90+
91+
92+
def test_box_pred_upload_ignore(dataset):
93+
prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
94+
response = dataset.annotate(predictions=[prediction])
95+
96+
assert response['predictions_processed'] == 1
97+
98+
# Copy so we don't modify the original.
99+
prediction_update_params = dict(TEST_BOX_PREDICTIONS[1])
100+
prediction_update_params['annotation_id'] = TEST_BOX_PREDICTIONS[0]['annotation_id']
101+
prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
102+
prediction_update = BoxPrediction(**prediction_update_params)
103+
# Default behavior is ignore.
104+
response = dataset.annotate(predictions=[prediction_update])
105+
106+
assert response['predictions_processed'] == 1
107+
108+
response = dataset.refloc(prediction.reference_id)['predictions']
109+
assert len(response) == 1
110+
response_prediction = response[0]
111+
assert_box_prediction_matches_dict(response_prediction, TEST_BOX_PREDICTIONS[0])
112+
113+
114+
def test_polygon_pred_upload_update(dataset):
115+
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
116+
response = dataset.annotate(predictions=[prediction])
117+
118+
assert response['predictions_processed'] == 1
119+
120+
# Copy so we don't modify the original.
121+
prediction_update_params = dict(TEST_POLYGON_PREDICTIONS[1])
122+
prediction_update_params['annotation_id'] = TEST_POLYGON_PREDICTIONS[0]['annotation_id']
123+
prediction_update_params['reference_id'] = TEST_POLYGON_PREDICTIONS[0]['reference_id']
124+
125+
prediction_update = PolygonPrediction(**prediction_update_params)
126+
response = dataset.annotate(predictions=[prediction_update], update=True)
127+
128+
assert response['predictions_processed'] == 1
129+
130+
response = dataset.refloc(prediction.reference_id)['predictions']
131+
assert len(response) == 1
132+
response_prediction = response[0]
133+
assert_polygon_prediction_matches_dict(response_prediction, prediction_update_params)
134+
135+
136+
def test_polygon_pred_upload_ignore(dataset):
137+
prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
138+
response = dataset.annotate(predictions=[prediction])
139+
140+
assert response['predictions_processed'] == 1
141+
142+
# Copy so we don't modify the original.
143+
prediction_update_params = dict(TEST_POLYGON_PREDICTIONS[1])
144+
prediction_update_params['annotation_id'] = TEST_POLYGON_PREDICTIONS[0]['annotation_id']
145+
prediction_update_params['reference_id'] = TEST_POLYGON_PREDICTIONS[0]['reference_id']
146+
147+
prediction_update = PolygonPrediction(**prediction_update_params)
148+
# Default behavior is ignore.
149+
response = dataset.annotate(predictions=[prediction_update])
150+
151+
assert response['predictions_processed'] == 1
152+
153+
response = dataset.refloc(prediction.reference_id)['predictions']
154+
assert len(response) == 1
155+
response_prediction = response[0]
156+
assert_polygon_prediction_matches_dict(response_prediction, TEST_POLYGON_PREDICTIONS[0])

0 commit comments

Comments
 (0)