|
1 |
| -from typing import List, Iterable, Set, Tuple, Optional |
| 1 | +from typing import Dict, List, Iterable, Set, Tuple, Optional, Union |
2 | 2 | from nucleus.dataset_item import DatasetItem
|
3 | 3 | from nucleus.annotation import Annotation
|
4 | 4 | from nucleus.utils import format_dataset_item_response
|
5 | 5 |
|
6 |
| -from .constants import DEFAULT_ANNOTATION_UPDATE_MODE, ITEM_KEY |
| 6 | +from .constants import DEFAULT_ANNOTATION_UPDATE_MODE |
7 | 7 |
|
8 | 8 |
|
9 | 9 | class Slice:
|
@@ -73,20 +73,42 @@ def append(
|
73 | 73 | )
|
74 | 74 | return response
|
75 | 75 |
|
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 | + """ |
78 | 89 | info = self.info()
|
79 | 90 | for item_metadata in info["dataset_items"]:
|
80 | 91 | yield format_dataset_item_response(
|
81 | 92 | self._client.dataitem_loc(
|
82 | 93 | dataset_id=info["dataset_id"],
|
83 | 94 | dataset_item_id=item_metadata["id"],
|
84 | 95 | )
|
85 |
| - )[ITEM_KEY] |
| 96 | + ) |
86 | 97 |
|
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()) |
90 | 112 |
|
91 | 113 | def annotate(
|
92 | 114 | self,
|
@@ -131,6 +153,13 @@ def check_annotations_are_in_slice(
|
131 | 153 |
|
132 | 154 | annotations: Annnotations with ids referring to targets.
|
133 | 155 | 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. |
134 | 163 | """
|
135 | 164 | info = slice_to_check.info()
|
136 | 165 |
|
|
0 commit comments