bug fixes and test updates

Author: Matt Sokoloff

labelbox/__init__.py

@@ -1,9 +1,9 @@
 name = "labelbox"
 __version__ = "2.7.0"
 
+from labelbox.schema.project import Project
 from labelbox.client import Client
 from labelbox.schema.bulk_import_request import BulkImportRequest
-from labelbox.schema.project import Project
 from labelbox.schema.dataset import Dataset
 from labelbox.schema.data_row import DataRow
 from labelbox.schema.label import Label
@@ -16,6 +16,6 @@
 from labelbox.schema.asset_attachment import AssetAttachment
 from labelbox.schema.webhook import Webhook
 from labelbox.schema.prediction import Prediction, PredictionModel
-from labelbox.schema.ontology import Ontology
+from labelbox.schema.ontology import Ontology, OntologyBuilder, Classification, Option, Tool
 from labelbox.schema.role import Role, ProjectRole
 from labelbox.schema.invite import Invite, InviteLimit

labelbox/data/annotation_types/__init__.py

@@ -0,0 +1,28 @@
+from .geometry import Line
+from .geometry import Point
+from .geometry import Mask
+from .geometry import Polygon
+from .geometry import Rectangle
+from .geometry import Geometry
+
+from .annotation import ClassificationAnnotation
+from .annotation import VideoClassificationAnnotation
+from .annotation import ObjectAnnotation
+from .annotation import VideoObjectAnnotation
+
+from .ner import TextEntity
+
+from .classification import Checklist
+from .classification import ClassificationAnswer
+from .classification import Dropdown
+from .classification import Radio
+from .classification import Text
+
+from .data import RasterData
+from .data import TextData
+from .data import VideoData
+
+from .label import Label
+
+from .collection import LabelList
+from .collection import LabelGenerator

labelbox/data/annotation_types/annotation.py

@@ -1,5 +1,7 @@
 from typing import Any, Dict, List, Union
 
+from pydantic.main import BaseModel
+
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
 from .geometry import Geometry
@@ -9,22 +11,19 @@
 class BaseAnnotation(FeatureSchema):
     """ Base annotation class. Shouldn't be directly instantiated
     """
-    classifications: List["ClassificationAnnotation"] = []
     extra: Dict[str, Any] = {}
 
 
-class ObjectAnnotation(BaseAnnotation):
-    """Class representing objects annotations (non classifications or annotations that have a location)
-    """
-    value: Union[TextEntity, Geometry]
-
-
 class ClassificationAnnotation(BaseAnnotation):
-    """Class represneting classification annotations (annotations that don't have a location) """
+    """Class representing classification annotations (annotations that don't have a location) """
     value: Union[Text, Checklist, Radio, Dropdown]
 
 
-ClassificationAnnotation.update_forward_refs()
+class ObjectAnnotation(BaseAnnotation):
+    """Class representing objects annotations (non classifications or annotations that have a location)
+    """
+    value: Union[TextEntity, Geometry]
+    classifications: List[ClassificationAnnotation] = []
 
 
 class VideoObjectAnnotation(ObjectAnnotation):
@@ -47,7 +46,3 @@ class VideoClassificationAnnotation(ClassificationAnnotation):
         frame: The frame index that this annotation corresponds to
     """
     frame: int
-
-
-VideoObjectAnnotation.update_forward_refs()
-ObjectAnnotation.update_forward_refs()

labelbox/data/annotation_types/collection.py

@@ -5,15 +5,16 @@
 
 from tqdm import tqdm
 
-from labelbox.schema.ontology import OntologyBuilder
+from labelbox.schema import ontology
 from labelbox.orm.model import Entity
+from ..ontology import get_classifications, get_tools
 from ..generator import PrefetchGenerator
 from .label import Label
 
 logger = logging.getLogger(__name__)
 
 
-class LabelCollection:
+class LabelList:
     """
     A container for interacting with a collection of labels.
     Less memory efficient than LabelGenerator but more performant and convenient to use.
@@ -25,15 +26,15 @@ def __init__(self, data: Iterable[Label]):
         self._index = 0
 
     def assign_schema_ids(
-            self, ontology_builder: OntologyBuilder) -> "LabelCollection":
+            self, ontology_builder: "ontology.OntologyBuilder") -> "LabelList":
         """
         Adds schema ids to all FeatureSchema objects in the Labels.
         This is necessary for MAL.
 
         Args:
-            ontology_builder: The ontology that matches the feature names assigned to objects in this LabelCollection
+            ontology_builder: The ontology that matches the feature names assigned to objects in this LabelList
         Returns:
-            LabelCollection. useful for chaining these modifying functions
+            LabelList. useful for chaining these modifying functions
         """
         for label in self._data:
             label.assign_schema_ids(ontology_builder)
@@ -42,7 +43,7 @@ def assign_schema_ids(
     def add_to_dataset(self,
                        dataset: "Entity.Dataset",
                        signer: Callable[[bytes], str],
-                       max_concurrency=20) -> "LabelCollection":
+                       max_concurrency=20) -> "LabelList":
         """
         Creates data rows from each labels data object and attaches the data to the given dataset.
         Updates the label's data object to have the same external_id and uid as the data row.
@@ -55,15 +56,15 @@ def add_to_dataset(self,
             dataset: labelbox dataset object to add the new data row to
             signer: A function that accepts bytes and returns a signed url.
         Returns:
-            LabelCollection with updated references to new data rows
+            LabelList with updated references to new data rows
         """
         self._ensure_unique_external_ids()
         self.add_url_to_data(signer, max_concurrency=max_concurrency)
         upload_task = dataset.create_data_rows([{
-            Entity.DataRow.row_data: label.data.url,
-            Entity.DataRow.external_id: label.data.external_id
+            'row_data': label.data.url,
+            'external_id': label.data.external_id
         } for label in self._data])
-        upload_task.wait_til_done()
+        upload_task.wait_till_done()
 
         data_row_lookup = {
             data_row.external_id: data_row.uid
@@ -73,9 +74,9 @@ def add_to_dataset(self,
             label.data.uid = data_row_lookup[label.data.external_id]
         return self
 
-    def add_url_to_masks(self, signer, max_concurrency=20) -> "LabelCollection":
+    def add_url_to_masks(self, signer, max_concurrency=20) -> "LabelList":
         """
-        Creates signed urls for all masks in the LabelCollection.
+        Creates signed urls for all masks in the LabelList.
         Multiple masks can reference the same RasterData mask so this makes sure we only upload that url once.
         Only uploads url if one doesn't already exist.
 
@@ -84,15 +85,15 @@ def add_url_to_masks(self, signer, max_concurrency=20) -> "LabelCollection":
             max_concurrency: how many threads to use for uploading.
                 Should be balanced to match the signing services capabilities.
         Returns:
-            LabelCollection with updated references to the new mask urls
+            LabelList with updated references to the new mask urls
         """
         for row in self._apply_threaded(
             [label.add_url_to_masks for label in self._data], max_concurrency,
                 signer):
             ...
         return self
 
-    def add_url_to_data(self, signer, max_concurrency=20) -> "LabelCollection":
+    def add_url_to_data(self, signer, max_concurrency=20) -> "LabelList":
         """
         Creates signed urls for the data
         Only uploads url if one doesn't already exist.
@@ -102,32 +103,46 @@ def add_url_to_data(self, signer, max_concurrency=20) -> "LabelCollection":
             max_concurrency: how many threads to use for uploading.
                 Should be balanced to match the signing services capabilities.
         Returns:
-            LabelCollection with updated references to the new data urls
+            LabelList with updated references to the new data urls
         """
         for row in self._apply_threaded(
             [label.add_url_to_data for label in self._data], max_concurrency,
                 signer):
             ...
         return self
 
+    def get_ontology(self) -> ontology.OntologyBuilder:
+        classifications = []
+        tools = []
+        for label in self._data:
+            tools = get_tools(label.object_annotations(), tools)
+            classifications = get_classifications(
+                label.classification_annotations(), classifications)
+        return ontology.OntologyBuilder(tools=tools,
+                                        classifications=classifications)
+
     def _ensure_unique_external_ids(self) -> None:
         external_ids = set()
         for label in self._data:
             if label.data.external_id is None:
-                label.data.external_id = uuid4()
+                label.data.external_id = str(uuid4())
             else:
                 if label.data.external_id in external_ids:
                     raise ValueError(
                         f"External ids must be unique for bulk uploading. Found {label.data.external_id} more than once."
                     )
             external_ids.add(label.data.external_id)
 
-    def __iter__(self) -> "LabelCollection":
+    def append(self, label: Label):
+        self._data.append(label)
+
+    def __iter__(self) -> "LabelList":
         self._index = 0
         return self
 
     def __next__(self) -> Label:
         if self._index == len(self._data):
+            self._index = 0
             raise StopIteration
 
         value = self._data[self._index]
@@ -154,18 +169,19 @@ class LabelGenerator(PrefetchGenerator):
     A container for interacting with a collection of labels.
 
     Use this class if you have larger data. It is slightly harder to work with
-    than the LabelCollection but will be much more memory efficient.
+    than the LabelList but will be much more memory efficient.
     """
 
     def __init__(self, data: Generator[Label, None, None], *args, **kwargs):
         self._fns = {}
         super().__init__(data, *args, **kwargs)
 
-    def as_collection(self) -> "LabelCollection":
-        return LabelCollection(data=list(self))
+    def as_list(self) -> "LabelList":
+        return LabelList(data=list(self))
 
     def assign_schema_ids(
-            self, ontology_builder: OntologyBuilder) -> "LabelGenerator":
+            self,
+            ontology_builder: "ontology.OntologyBuilder") -> "LabelGenerator":
 
         def _assign_ids(label: Label):
             label.assign_schema_ids(ontology_builder)
@@ -190,7 +206,7 @@ def _add_url_to_data(label: Label):
             label.add_url_to_data(signer)
             return label
 
-        self._fns['_add_url_to_data'] = _add_url_to_data
+        self._fns['add_url_to_data'] = _add_url_to_data
         return self
 
     def add_to_dataset(self, dataset: "Entity.Dataset",
@@ -199,7 +215,7 @@ def add_to_dataset(self, dataset: "Entity.Dataset",
         Creates data rows from each labels data object and attaches the data to the given dataset.
         Updates the label's data object to have the same external_id and uid as the data row.
 
-        This is a lot slower than LabelCollection.add_to_dataset but also more memory efficient.
+        This is a lot slower than LabelList.add_to_dataset but also more memory efficient.
 
         Args:
             dataset: labelbox dataset object to add the new data row to
@@ -237,6 +253,20 @@ def _add_url_to_masks(label: Label):
         self._fns['add_url_to_masks'] = _add_url_to_masks
         return self
 
+    def register_background_fn(self, fn: Callable[[Label], Label],
+                               name: str) -> "LabelGenerator":
+        """
+        Allows users to add arbitrary io functions to the generator.
+        These functions will be exectuted in parallel and added to a prefetch queue.
+
+        Args:
+            fn: Callable that modifies a label and then returns the same label
+                - For performance reasons, this function shouldn't run if the object already has the desired state.
+            name: Register the name of the function. If the name already exists, then the function will be replaced.
+        """
+        self._fns[name] = fn
+        return self
+
     def __iter__(self):
         return self
 
@@ -255,4 +285,4 @@ def __next__(self):
         return self._process(value)
 
 
-LabelData = Union[LabelCollection, LabelGenerator]
+LabelCollection = Union[LabelList, LabelGenerator]