Skip to content

Commit 0dbf4f7

Browse files
committed
added annotations to return value
1 parent 77023a1 commit 0dbf4f7

File tree

2 files changed

+39
-10
lines changed

2 files changed

+39
-10
lines changed

nucleus/slice.py

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
from typing import List, Iterable, Set, Tuple, Optional
1+
from typing import Dict, List, Iterable, Set, Tuple, Optional, Union
22
from nucleus.dataset_item import DatasetItem
33
from nucleus.annotation import Annotation
44
from nucleus.utils import format_dataset_item_response
55

6-
from .constants import DEFAULT_ANNOTATION_UPDATE_MODE, ITEM_KEY
6+
from .constants import DEFAULT_ANNOTATION_UPDATE_MODE
77

88

99
class Slice:
@@ -73,20 +73,42 @@ def append(
7373
)
7474
return response
7575

76-
def items_generator(self) -> Iterable[DatasetItem]:
77-
"""Returns an iterable of DatasetItem/Annotation dicts."""
76+
def items_and_annotation_generator(
77+
self,
78+
) -> Iterable[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:
79+
"""Returns an iterable of all DatasetItems and Annotations in this slice.
80+
81+
Returns:
82+
An iterable, where each item is a dict with two keys representing a row
83+
in the dataset.
84+
* One value in the dict is the DatasetItem, containing a reference to the
85+
item that was annotated, for example an image_url.
86+
* The other value is a dictionary containing all the annotations for this
87+
dataset item, sorted by annotation type.
88+
"""
7889
info = self.info()
7990
for item_metadata in info["dataset_items"]:
8091
yield format_dataset_item_response(
8192
self._client.dataitem_loc(
8293
dataset_id=info["dataset_id"],
8394
dataset_item_id=item_metadata["id"],
8495
)
85-
)[ITEM_KEY]
96+
)
8697

87-
def items(self) -> List[DatasetItem]:
88-
"""Returns a list of all DatasetItems in this slice."""
89-
return list(self.items_generator())
98+
def items_and_annotations(
99+
self,
100+
) -> List[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:
101+
"""Returns a list of all DatasetItems and Annotations in this slice.
102+
103+
Returns:
104+
A list, where each item is a dict with two keys representing a row
105+
in the dataset.
106+
* One value in the dict is the DatasetItem, containing a reference to the
107+
item that was annotated.
108+
* The other value is a dictionary containing all the annotations for this
109+
dataset item, sorted by annotation type.
110+
"""
111+
return list(self.items_and_annotation_generator())
90112

91113
def annotate(
92114
self,
@@ -131,6 +153,13 @@ def check_annotations_are_in_slice(
131153
132154
annotations: Annnotations with ids referring to targets.
133155
slice: The slice to check against.
156+
157+
158+
Returns:
159+
A tuple, where the first element is true/false whether the annotations are all
160+
in the slice.
161+
The second element is the list of item_ids not in the slice.
162+
The third element is the list of ref_ids not in the slice.
134163
"""
135164
info = slice_to_check.info()
136165

tests/test_slice.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import pytest
22
from nucleus import Slice, NucleusClient, DatasetItem, BoxAnnotation
3-
from nucleus.constants import ERROR_PAYLOAD
3+
from nucleus.constants import ERROR_PAYLOAD, ITEM_KEY
44
from helpers import (
55
TEST_DATASET_NAME,
66
TEST_IMG_URLS,
@@ -127,7 +127,7 @@ def test_slice_append(dataset):
127127
or item.reference_id == response["dataset_items"][2]["ref_id"]
128128
)
129129

130-
all_stored_items = slc.items()
130+
all_stored_items = [_[ITEM_KEY] for _ in slc.items_and_annotations()]
131131

132132
def sort_by_reference_id(items):
133133
# Remove the generated item_ids and standardize

0 commit comments

Comments
 (0)