Skip to content

Commit 0f66bc7

Browse files
sasha-scalesasha-scale
committed
nits
1 parent 7e7bb42 commit 0f66bc7

File tree

6 files changed

+273
-165
lines changed

6 files changed

+273
-165
lines changed

‎nucleus/__init__.py

Lines changed: 54 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,11 @@
6969

7070
from .dataset import Dataset
7171
from .dataset_item import DatasetItem
72-
from .annotation import BoxAnnotation, PolygonAnnotation, SegmentationAnnotation
72+
from .annotation import (
73+
BoxAnnotation,
74+
PolygonAnnotation,
75+
SegmentationAnnotation,
76+
)
7377
from .prediction import BoxPrediction, PolygonPrediction
7478
from .model_run import ModelRun
7579
from .slice import Slice
@@ -79,6 +83,7 @@
7983
construct_annotation_payload,
8084
construct_model_creation_payload,
8185
construct_box_predictions_payload,
86+
construct_segmentation_payload,
8287
)
8388
from .constants import (
8489
NUCLEUS_ENDPOINT,
@@ -464,7 +469,9 @@ def exception_handler(request, exception):
464469
def annotate_dataset(
465470
self,
466471
dataset_id: str,
467-
annotations: List[Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]],
472+
annotations: List[
473+
Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
474+
],
468475
update: bool,
469476
batch_size: int = 100,
470477
):
@@ -477,8 +484,16 @@ def annotate_dataset(
477484
"""
478485

479486
# Split payload into segmentations and Box/Polygon
480-
segmentations = [ann for ann in annotations if isinstance(ann, SegmentationAnnotation)]
481-
other_annotations = [ann for ann in annotations if not isinstance(ann, SegmentationAnnotation)]
487+
segmentations = [
488+
ann
489+
for ann in annotations
490+
if isinstance(ann, SegmentationAnnotation)
491+
]
492+
other_annotations = [
493+
ann
494+
for ann in annotations
495+
if not isinstance(ann, SegmentationAnnotation)
496+
]
482497

483498
batches = [
484499
other_annotations[i : i + batch_size]
@@ -496,39 +511,42 @@ def annotate_dataset(
496511
ANNOTATIONS_IGNORED_KEY: 0,
497512
}
498513

499-
tqdm_batches = self.tqdm_bar(batches)
514+
total_batches = len(batches) + len(semseg_batches)
500515

501-
for batch in tqdm_batches:
502-
payload = construct_annotation_payload(batch, update)
503-
response = self._make_request(
504-
payload, f"dataset/{dataset_id}/annotate"
505-
)
506-
if STATUS_CODE_KEY in response:
507-
agg_response[ERRORS_KEY] = response
508-
else:
509-
agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
510-
ANNOTATIONS_PROCESSED_KEY
511-
]
512-
agg_response[ANNOTATIONS_IGNORED_KEY] += response[
513-
ANNOTATIONS_IGNORED_KEY
514-
]
516+
tqdm_batches = self.tqdm_bar(batches)
515517

516-
for s_batch in semseg_batches:
517-
payload = {"segmentations": [seg.to_payload() for seg in s_batch]}
518-
if update:
519-
payload["force"] = update
520-
response = self._make_request(
521-
payload, f"dataset/{dataset_id}/annotate_segmentation"
522-
)
523-
if STATUS_CODE_KEY in response:
524-
agg_response[ERRORS_KEY] = response
525-
else:
526-
agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
527-
ANNOTATIONS_PROCESSED_KEY
528-
]
529-
agg_response[ANNOTATIONS_IGNORED_KEY] += response[
530-
ANNOTATIONS_IGNORED_KEY
531-
]
518+
with self.tqdm_bar(total=total_batches) as pbar:
519+
for batch in tqdm_batches:
520+
payload = construct_annotation_payload(batch, update)
521+
response = self._make_request(
522+
payload, f"dataset/{dataset_id}/annotate"
523+
)
524+
pbar.update(1)
525+
if STATUS_CODE_KEY in response:
526+
agg_response[ERRORS_KEY] = response
527+
else:
528+
agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
529+
ANNOTATIONS_PROCESSED_KEY
530+
]
531+
agg_response[ANNOTATIONS_IGNORED_KEY] += response[
532+
ANNOTATIONS_IGNORED_KEY
533+
]
534+
535+
for s_batch in semseg_batches:
536+
payload = construct_segmentation_payload(s_batch, update)
537+
response = self._make_request(
538+
payload, f"dataset/{dataset_id}/annotate_segmentation"
539+
)
540+
pbar.update(1)
541+
if STATUS_CODE_KEY in response:
542+
agg_response[ERRORS_KEY] = response
543+
else:
544+
agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
545+
ANNOTATIONS_PROCESSED_KEY
546+
]
547+
agg_response[ANNOTATIONS_IGNORED_KEY] += response[
548+
ANNOTATIONS_IGNORED_KEY
549+
]
532550

533551
return agg_response
534552

@@ -635,7 +653,7 @@ def predict(
635653

636654
for batch in tqdm_batches:
637655
batch_payload = construct_box_predictions_payload(
638-
annotations,
656+
batch,
639657
update,
640658
)
641659
response = self._make_request(

‎nucleus/annotation.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ def __str__(self):
3636
@classmethod
3737
def from_json(cls, payload: dict):
3838
return cls(
39-
label=payload.get(LABEL_KEY),
40-
index=payload.get(INDEX_KEY),
39+
label=payload.get(LABEL_KEY, ""),
40+
index=payload.get(INDEX_KEY, None),
4141
metadata=payload.get(METADATA_KEY, None),
4242
)
4343

@@ -75,7 +75,10 @@ def __str__(self):
7575
def from_json(cls, payload: dict):
7676
return cls(
7777
mask_url=payload[MASK_URL_KEY],
78-
annotations=[Segment.from_json(ann) for ann in payload.get(ANNOTATIONS_KEY, [])],
78+
annotations=[
79+
Segment.from_json(ann)
80+
for ann in payload.get(ANNOTATIONS_KEY, [])
81+
],
7982
reference_id=payload.get(REFERENCE_ID_KEY, None),
8083
item_id=payload.get(ITEM_ID_KEY, None),
8184
)
@@ -205,4 +208,4 @@ def to_payload(self) -> dict:
205208
}
206209

207210
def __str__(self):
208-
return str(self.to_payload())
211+
return str(self.to_payload())

‎nucleus/constants.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,4 @@
5151
AUTOTAGS_KEY = "autotags"
5252
MASK_URL_KEY = "mask_url"
5353
INDEX_KEY = "index"
54+
SEGMENTATIONS_KEY = "SEGMENTATIONS"

‎nucleus/payload_constructor.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
from typing import List, Optional, Dict, Union
22
from .dataset_item import DatasetItem
3-
from .annotation import BoxAnnotation, PolygonAnnotation
3+
from .annotation import (
4+
BoxAnnotation,
5+
PolygonAnnotation,
6+
SegmentationAnnotation,
7+
)
48
from .prediction import BoxPrediction, PolygonPrediction
59
from .constants import (
610
ANNOTATION_UPDATE_KEY,
@@ -11,6 +15,7 @@
1115
ITEMS_KEY,
1216
FORCE_KEY,
1317
MODEL_ID_KEY,
18+
SEGMENTATIONS_KEY,
1419
)
1520

1621

@@ -39,6 +44,17 @@ def construct_annotation_payload(
3944
return {ANNOTATIONS_KEY: annotations, ANNOTATION_UPDATE_KEY: update}
4045

4146

47+
def construct_segmentation_payload(
48+
annotation_items: List[SegmentationAnnotation],
49+
update: bool,
50+
) -> dict:
51+
annotations = []
52+
for annotation_item in annotation_items:
53+
annotations.append(annotation_item.to_payload())
54+
55+
return {SEGMENTATIONS_KEY: annotations, "force": update}
56+
57+
4258
def construct_box_predictions_payload(
4359
box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
4460
update: bool,

‎tests/helpers.py

Lines changed: 68 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -2,31 +2,34 @@
22
from urllib.parse import urlparse
33
import boto3
44

5-
PRESIGN_EXPIRY_SECONDS = 60*60*24*2 #2 days
5+
PRESIGN_EXPIRY_SECONDS = 60 * 60 * 24 * 2 # 2 days
66

7-
TEST_MODEL_NAME = '[PyTest] Test Model'
8-
TEST_MODEL_REFERENCE = '[PyTest] Test Model Reference'
9-
TEST_MODEL_RUN = '[PyTest] Test Model Run'
10-
TEST_DATASET_NAME = '[PyTest] Test Dataset'
11-
TEST_SLICE_NAME = '[PyTest] Test Slice'
7+
TEST_MODEL_NAME = "[PyTest] Test Model"
8+
TEST_MODEL_REFERENCE = "[PyTest] Test Model Reference"
9+
TEST_MODEL_RUN = "[PyTest] Test Model Run"
10+
TEST_DATASET_NAME = "[PyTest] Test Dataset"
11+
TEST_SLICE_NAME = "[PyTest] Test Slice"
1212
TEST_IMG_URLS = [
13-
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg',
14-
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg',
15-
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg',
16-
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg',
17-
's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg',
13+
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg",
14+
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg",
15+
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg",
16+
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg",
17+
"s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg",
1818
]
1919

20+
2021
def get_signed_url(url):
2122
bucket, key = get_s3_details(url)
2223
return s3_sign(bucket, key)
2324

25+
2426
def get_s3_details(url):
2527
# Expects S3 URL format to be https://<BUCKET>.s3.amazonaws.com/<KEY>
2628
parsed = urlparse(url)
27-
bucket = parsed.netloc[:parsed.netloc.find(".")]
29+
bucket = parsed.netloc[: parsed.netloc.find(".")]
2830
return bucket, parsed.path[1:]
2931

32+
3033
def s3_sign(bucket, key):
3134
s3 = boto3.client("s3")
3235
return s3.generate_presigned_url(
@@ -37,35 +40,37 @@ def s3_sign(bucket, key):
3740
},
3841
ExpiresIn=PRESIGN_EXPIRY_SECONDS,
3942
)
40-
43+
44+
4145
def reference_id_from_url(url):
4246
return Path(url).name
4347

48+
4449
TEST_BOX_ANNOTATIONS = [
4550
{
46-
'label': f'[Pytest] Box Annotation ${i}',
47-
'x': 50 + i * 10,
48-
'y': 60 + i * 10,
49-
'width': 70 + i * 10,
50-
'height': 80 + i * 10,
51-
'reference_id': reference_id_from_url(TEST_IMG_URLS[i]),
52-
'annotation_id': f'[Pytest] Box Annotation Annotation Id{i}',
51+
"label": f"[Pytest] Box Annotation ${i}",
52+
"x": 50 + i * 10,
53+
"y": 60 + i * 10,
54+
"width": 70 + i * 10,
55+
"height": 80 + i * 10,
56+
"reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
57+
"annotation_id": f"[Pytest] Box Annotation Annotation Id{i}",
5358
}
5459
for i in range(len(TEST_IMG_URLS))
5560
]
5661

5762
TEST_POLYGON_ANNOTATIONS = [
5863
{
59-
'label': f'[Pytest] Polygon Annotation ${i}',
60-
'vertices': [
64+
"label": f"[Pytest] Polygon Annotation ${i}",
65+
"vertices": [
6166
{
62-
'x': 50 + i * 10 + j,
63-
'y': 60 + i * 10 + j,
67+
"x": 50 + i * 10 + j,
68+
"y": 60 + i * 10 + j,
6469
}
6570
for j in range(3)
6671
],
67-
'reference_id': reference_id_from_url(TEST_IMG_URLS[i]),
68-
'annotation_id': f'[Pytest] Polygon Annotation Annotation Id{i}',
72+
"reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
73+
"annotation_id": f"[Pytest] Polygon Annotation Annotation Id{i}",
6974
}
7075
for i in range(len(TEST_IMG_URLS))
7176
]
@@ -76,52 +81,62 @@ def reference_id_from_url(url):
7681
"reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
7782
"mask_url": get_signed_url(TEST_MASK_URL),
7883
"annotations": [
79-
{"label": "bear", "index": 2}, {"label": "grass-merged", "index": 1}
80-
]
84+
{"label": "bear", "index": 2},
85+
{"label": "grass-merged", "index": 1},
86+
],
8187
}
8288
for i in range(len(TEST_IMG_URLS))
8389
]
8490

8591
TEST_BOX_PREDICTIONS = [
86-
{
87-
**TEST_BOX_ANNOTATIONS[i],
88-
'confidence': 0.10 * i
89-
}
92+
{**TEST_BOX_ANNOTATIONS[i], "confidence": 0.10 * i}
9093
for i in range(len(TEST_BOX_ANNOTATIONS))
9194
]
9295

9396
TEST_POLYGON_PREDICTIONS = [
94-
{
95-
**TEST_POLYGON_ANNOTATIONS[i],
96-
'confidence': 0.10 * i
97-
}
97+
{**TEST_POLYGON_ANNOTATIONS[i], "confidence": 0.10 * i}
9898
for i in range(len(TEST_POLYGON_ANNOTATIONS))
9999
]
100100

101101

102102
# Asserts that a box annotation instance matches a dict representing its properties.
103103
# Useful to check annotation uploads/updates match.
104104
def assert_box_annotation_matches_dict(annotation_instance, annotation_dict):
105-
assert annotation_instance.label == annotation_dict['label']
106-
assert annotation_instance.x == annotation_dict['x']
107-
assert annotation_instance.y == annotation_dict['y']
108-
assert annotation_instance.height == annotation_dict['height']
109-
assert annotation_instance.width == annotation_dict['width']
110-
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
111-
112-
def assert_polygon_annotation_matches_dict(annotation_instance, annotation_dict):
113-
assert annotation_instance.label == annotation_dict['label']
114-
assert annotation_instance.annotation_id == annotation_dict['annotation_id']
115-
for instance_pt, dict_pt in zip(annotation_instance.vertices, annotation_dict['vertices']):
116-
assert instance_pt['x'] == dict_pt['x']
117-
assert instance_pt['y'] == dict_pt['y']
105+
assert annotation_instance.label == annotation_dict["label"]
106+
assert annotation_instance.x == annotation_dict["x"]
107+
assert annotation_instance.y == annotation_dict["y"]
108+
assert annotation_instance.height == annotation_dict["height"]
109+
assert annotation_instance.width == annotation_dict["width"]
110+
assert (
111+
annotation_instance.annotation_id == annotation_dict["annotation_id"]
112+
)
113+
114+
115+
def assert_polygon_annotation_matches_dict(
116+
annotation_instance, annotation_dict
117+
):
118+
assert annotation_instance.label == annotation_dict["label"]
119+
assert (
120+
annotation_instance.annotation_id == annotation_dict["annotation_id"]
121+
)
122+
for instance_pt, dict_pt in zip(
123+
annotation_instance.vertices, annotation_dict["vertices"]
124+
):
125+
assert instance_pt["x"] == dict_pt["x"]
126+
assert instance_pt["y"] == dict_pt["y"]
127+
118128

119129
# Asserts that a box prediction instance matches a dict representing its properties.
120130
# Useful to check prediction uploads/updates match.
121131
def assert_box_prediction_matches_dict(prediction_instance, prediction_dict):
122132
assert_box_annotation_matches_dict(prediction_instance, prediction_dict)
123-
assert prediction_instance.confidence == prediction_dict['confidence']
133+
assert prediction_instance.confidence == prediction_dict["confidence"]
124134

125-
def assert_polygon_prediction_matches_dict(prediction_instance, prediction_dict):
126-
assert_polygon_annotation_matches_dict(prediction_instance, prediction_dict)
127-
assert prediction_instance.confidence == prediction_dict['confidence']
135+
136+
def assert_polygon_prediction_matches_dict(
137+
prediction_instance, prediction_dict
138+
):
139+
assert_polygon_annotation_matches_dict(
140+
prediction_instance, prediction_dict
141+
)
142+
assert prediction_instance.confidence == prediction_dict["confidence"]

0 commit comments

Comments
 (0)