Skip to content

Commit 7e64db6

Browse files
jihan-yinjihan-yin
committed
updated pytests, gt upload idempotency tests working
1 parent a56539d commit 7e64db6

File tree

4 files changed

+215
-13
lines changed

4 files changed

+215
-13
lines changed

‎nucleus/dataset.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ def refloc(self, reference_id: str) -> dict:
169169
:return:
170170
{
171171
"item": DatasetItem,
172-
"annotations": List[Box2DAnnotation],
172+
"annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
173173
}
174174
"""
175175
response = self._client.dataitem_ref_id(self.id, reference_id)

‎tests/helpers.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from pathlib import Path
2+
3+
TEST_DATASET_NAME = '[PyTest] Test Dataset'
4+
TEST_IMG_URLS = [
5+
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg',
6+
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg',
7+
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg',
8+
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg',
9+
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg',
10+
]
11+
12+
def reference_id_from_url(url):
13+
return Path(url).name
14+
15+
TEST_BOX_ANNOTATIONS = [
16+
{
17+
'label': f'[Pytest] Box Annotation ${i}',
18+
'x': 50 + i * 10,
19+
'y': 60 + i * 10,
20+
'width': 70 + i * 10,
21+
'height': 80 + i * 10,
22+
'reference_id': reference_id_from_url(TEST_IMG_URLS[i]),
23+
'annotation_id': f'[Pytest] Box Annotation Annotation Id{i}',
24+
}
25+
for i in range(len(TEST_IMG_URLS))
26+
]
27+
28+
TEST_POLYGON_ANNOTATIONS = [
29+
{
30+
'label': f'[Pytest] Polygon Annotation ${i}',
31+
'vertices': [
32+
{
33+
'x': 50 + i * 10 + j,
34+
'y': 60 + i * 10 + j,
35+
}
36+
for j in range(3)
37+
],
38+
'reference_id': reference_id_from_url(TEST_IMG_URLS[i]),
39+
'annotation_id': f'[Pytest] Polygon Annotation Annotation Id{i}',
40+
}
41+
for i in range(len(TEST_IMG_URLS))
42+
]

‎tests/test_annotation.py

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
import pytest
2+
3+
from helpers import (
4+
TEST_DATASET_NAME,
5+
TEST_IMG_URLS,
6+
TEST_BOX_ANNOTATIONS,
7+
TEST_POLYGON_ANNOTATIONS,
8+
reference_id_from_url,
9+
)
10+
11+
from nucleus import BoxAnnotation, PolygonAnnotation, DatasetItem
12+
from nucleus.constants import ERROR_PAYLOAD
13+
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+
32+
@pytest.fixture()
33+
def dataset(CLIENT):
34+
ds = CLIENT.create_dataset(TEST_DATASET_NAME)
35+
ds_items = []
36+
for url in TEST_IMG_URLS:
37+
ds_items.append(DatasetItem(
38+
image_location=url,
39+
reference_id=reference_id_from_url(url),
40+
))
41+
42+
response = ds.append(ds_items)
43+
assert ERROR_PAYLOAD not in response.json()
44+
yield ds
45+
46+
response = CLIENT.delete_dataset(ds.id)
47+
assert response == {}
48+
49+
50+
def test_box_gt_upload(dataset):
51+
annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
52+
response = dataset.annotate(annotations=[annotation])
53+
54+
assert response['dataset_id'] == dataset.id
55+
assert response['annotations_processed'] == 1
56+
57+
response = dataset.refloc(annotation.reference_id)['annotations']
58+
assert len(response) == 1
59+
response_annotation = response[0]
60+
assert_box_annotation_matches_dict(response_annotation, TEST_BOX_ANNOTATIONS[0])
61+
62+
63+
def test_polygon_gt_upload(dataset):
64+
annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
65+
response = dataset.annotate(annotations=[annotation])
66+
67+
assert response['dataset_id'] == dataset.id
68+
assert response['annotations_processed'] == 1
69+
70+
response = dataset.refloc(annotation.reference_id)['annotations']
71+
assert len(response) == 1
72+
response_annotation = response[0]
73+
print(response_annotation)
74+
assert_polygon_annotation_matches_dict(response_annotation, TEST_POLYGON_ANNOTATIONS[0])
75+
76+
77+
def test_box_gt_upload_update(dataset):
78+
annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
79+
response = dataset.annotate(annotations=[annotation])
80+
81+
assert response['annotations_processed'] == 1
82+
83+
# Copy so we don't modify the original.
84+
annotation_update_params = dict(TEST_BOX_ANNOTATIONS[1])
85+
annotation_update_params['annotation_id'] = TEST_BOX_ANNOTATIONS[0]['annotation_id']
86+
annotation_update_params['reference_id'] = TEST_BOX_ANNOTATIONS[0]['reference_id']
87+
88+
annotation_update = BoxAnnotation(**annotation_update_params)
89+
response = dataset.annotate(annotations=[annotation_update], update=True)
90+
91+
assert response['annotations_processed'] == 1
92+
93+
response = dataset.refloc(annotation.reference_id)['annotations']
94+
assert len(response) == 1
95+
response_annotation = response[0]
96+
assert_box_annotation_matches_dict(response_annotation, annotation_update_params)
97+
98+
99+
def test_box_gt_upload_ignore(dataset):
100+
annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
101+
response = dataset.annotate(annotations=[annotation])
102+
103+
assert response['annotations_processed'] == 1
104+
105+
# Copy so we don't modify the original.
106+
annotation_update_params = dict(TEST_BOX_ANNOTATIONS[1])
107+
annotation_update_params['annotation_id'] = TEST_BOX_ANNOTATIONS[0]['annotation_id']
108+
annotation_update_params['reference_id'] = TEST_BOX_ANNOTATIONS[0]['reference_id']
109+
annotation_update = BoxAnnotation(**annotation_update_params)
110+
# Default behavior is ignore.
111+
response = dataset.annotate(annotations=[annotation_update])
112+
113+
assert response['annotations_processed'] == 1
114+
115+
response = dataset.refloc(annotation.reference_id)['annotations']
116+
assert len(response) == 1
117+
response_annotation = response[0]
118+
assert_box_annotation_matches_dict(response_annotation, TEST_BOX_ANNOTATIONS[0])
119+
120+
121+
def test_polygon_gt_upload_update(dataset):
122+
annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
123+
response = dataset.annotate(annotations=[annotation])
124+
125+
assert response['annotations_processed'] == 1
126+
127+
# Copy so we don't modify the original.
128+
annotation_update_params = dict(TEST_POLYGON_ANNOTATIONS[1])
129+
annotation_update_params['annotation_id'] = TEST_POLYGON_ANNOTATIONS[0]['annotation_id']
130+
annotation_update_params['reference_id'] = TEST_POLYGON_ANNOTATIONS[0]['reference_id']
131+
132+
annotation_update = PolygonAnnotation(**annotation_update_params)
133+
response = dataset.annotate(annotations=[annotation_update], update=True)
134+
135+
assert response['annotations_processed'] == 1
136+
137+
response = dataset.refloc(annotation.reference_id)['annotations']
138+
assert len(response) == 1
139+
response_annotation = response[0]
140+
assert_polygon_annotation_matches_dict(response_annotation, annotation_update_params)
141+
142+
143+
def test_polygon_gt_upload_ignore(dataset):
144+
annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
145+
response = dataset.annotate(annotations=[annotation])
146+
147+
assert response['annotations_processed'] == 1
148+
149+
# Copy so we don't modify the original.
150+
annotation_update_params = dict(TEST_POLYGON_ANNOTATIONS[1])
151+
annotation_update_params['annotation_id'] = TEST_POLYGON_ANNOTATIONS[0]['annotation_id']
152+
annotation_update_params['reference_id'] = TEST_POLYGON_ANNOTATIONS[0]['reference_id']
153+
154+
annotation_update = PolygonAnnotation(**annotation_update_params)
155+
# Default behavior is ignore.
156+
response = dataset.annotate(annotations=[annotation_update])
157+
158+
assert response['annotations_processed'] == 1
159+
160+
response = dataset.refloc(annotation.reference_id)['annotations']
161+
assert len(response) == 1
162+
response_annotation = response[0]
163+
assert_polygon_annotation_matches_dict(response_annotation, TEST_POLYGON_ANNOTATIONS[0])

‎tests/test_dataset.py

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
1-
from pathlib import Path
21
import pytest
2+
3+
from helpers import (
4+
TEST_DATASET_NAME,
5+
TEST_IMG_URLS,
6+
reference_id_from_url,
7+
)
8+
39
from nucleus import Dataset, DatasetItem, UploadResponse
410
from nucleus.constants import (
511
NEW_ITEMS,
@@ -10,22 +16,13 @@
1016
DATASET_ID_KEY,
1117
)
1218

13-
TEST_DATASET_NAME = '[PyTest] Test Dataset'
14-
TEST_SLICE_NAME = '[PyTest] Test Slice'
15-
TEST_IMG_URLS = [
16-
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg",
17-
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg",
18-
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg",
19-
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg",
20-
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg",
21-
]
22-
2319
@pytest.fixture()
2420
def dataset(CLIENT):
2521
ds = CLIENT.create_dataset(TEST_DATASET_NAME)
2622
yield ds
2723

28-
CLIENT.delete_dataset(ds.id)
24+
response = CLIENT.delete_dataset(ds.id)
25+
assert response == {}
2926

3027
def test_dataset_create_and_delete(CLIENT):
3128
# Creation

0 commit comments

Comments
 (0)