update tests

Author: Matt Sokoloff

labelbox/data/generator.py

@@ -2,6 +2,7 @@
 import threading
 from queue import Queue
 from typing import Any, Iterable
+from concurrent.futures import ThreadPoolExecutor
 
 logger = logging.getLogger(__name__)
 
@@ -32,10 +33,7 @@ class PrefetchGenerator:
     Useful for modifying the generator results based on data from a network
     """
 
-    def __init__(self,
-                 data: Iterable[Any],
-                 prefetch_limit=20,
-                 max_concurrency=4):
+    def __init__(self, data: Iterable[Any], prefetch_limit=20, num_executors=4):
         if isinstance(data, (list, tuple)):
             self._data = (r for r in data)
         else:
@@ -44,14 +42,13 @@ def __init__(self,
         self.queue = Queue(prefetch_limit)
         self._data = ThreadSafeGen(self._data)
         self.completed_threads = 0
-        self.max_concurrency = max_concurrency
-        self.threads = [
-            threading.Thread(target=self.fill_queue)
-            for _ in range(max_concurrency)
-        ]
-        for thread in self.threads:
-            thread.daemon = True
-            thread.start()
+        # Can only iterate over once it the queue.get hangs forever.
+        self.done = False
+        self.num_executors = num_executors
+        with ThreadPoolExecutor(max_workers=num_executors) as executor:
+            self.futures = [
+                executor.submit(self.fill_queue) for _ in range(num_executors)
+            ]
 
     def _process(self, value) -> Any:
         raise NotImplementedError("Abstract method needs to be implemented")
@@ -73,10 +70,13 @@ def __iter__(self):
         return self
 
     def __next__(self) -> Any:
+        if self.done:
+            raise StopIteration
         value = self.queue.get()
         while value is None:
             self.completed_threads += 1
-            if self.completed_threads == self.max_concurrency:
+            if self.completed_threads == self.num_executors:
+                self.done = True
                 raise StopIteration
             value = self.queue.get()
         return value

labelbox/data/serialization/ndjson/classification.py

@@ -189,10 +189,6 @@ def from_common(
             raise TypeError(
                 f"Unable to convert object to MAL format. `{type(annotation.value)}`"
             )
-        if len(annotation.classifications):
-            raise ValueError(
-                "Nested classifications not supported by this format")
-
         return classify_obj.from_common(annotation.value, annotation.schema_id,
                                         annotation.extra, data)
 

labelbox/data/serialization/ndjson/converter.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Dict, Generator, Iterable
 
-from ...annotation_types.collection import LabelData, LabelGenerator
+from ...annotation_types.collection import LabelCollection, LabelGenerator
 from .label import NDLabel
 
 logger = logging.getLogger(__name__)
@@ -23,7 +23,8 @@ def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator:
         return data.to_common()
 
     @staticmethod
-    def serialize(labels: LabelData) -> Generator[Dict[str, Any], None, None]:
+    def serialize(
+            labels: LabelCollection) -> Generator[Dict[str, Any], None, None]:
         """
         Converts a labelbox common object to the labelbox ndjson format (prediction import format)
 
@@ -32,7 +33,7 @@ def serialize(labels: LabelData) -> Generator[Dict[str, Any], None, None]:
         We will continue to improve the error messages and add helper functions to deal with this.
 
         Args:
-            labels: Either a LabelCollection or a LabelGenerator
+            labels: Either a LabelList or a LabelGenerator
         Returns:
             A generator for accessing the ndjson representation of the data
         """

labelbox/data/serialization/ndjson/label.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoClassificationAnnotation
-from ...annotation_types.collection import LabelData, LabelGenerator
+from ...annotation_types.collection import LabelCollection, LabelGenerator
 from ...annotation_types.data import RasterData, TextData, VideoData
 from ...annotation_types.label import Label
 from ...annotation_types.ner import TextEntity
@@ -25,7 +25,8 @@ def to_common(self) -> LabelGenerator:
             data=self._generate_annotations(grouped_annotations))
 
     @classmethod
-    def from_common(cls, data: LabelData) -> Generator["NDLabel", None, None]:
+    def from_common(cls,
+                    data: LabelCollection) -> Generator["NDLabel", None, None]:
         for label in data:
             yield from cls._create_non_video_annotations(label)
             yield from cls._create_video_annotations(label)

labelbox/data/serialization/ndjson/objects.py

@@ -123,7 +123,7 @@ class NDMask(NDBaseObject):
 
     def to_common(self) -> Mask:
         return Mask(mask=RasterData(url=self.mask.instanceURI),
-                    color_rgb=self.mask.colorRGB)
+                    color=self.mask.colorRGB)
 
     @classmethod
     def from_common(cls, mask: Mask,
@@ -134,8 +134,7 @@ def from_common(cls, mask: Mask,
             raise ValueError(
                 "Mask does not have a url. Use `LabelGenerator.add_url_to_masks`, `LabelCollection.add_url_to_masks`, or `Label.add_url_to_masks`."
             )
-        return cls(mask=_Mask(instanceURI=mask.mask.url,
-                              colorRGB=mask.color_rgb),
+        return cls(mask=_Mask(instanceURI=mask.mask.url, colorRGB=mask.color),
                    dataRow=DataRow(id=data.uid),
                    schema_id=schema_id,
                    uuid=extra.get('uuid'),

tests/data/serialization/ndjson/test_classification.py

@@ -7,6 +7,6 @@ def test_classification():
     with open('tests/data/assets/ndjson/classification_import.json',
               'r') as file:
         data = json.load(file)
-    res = NDJsonConverter.deserialize(data).as_collection()
+    res = NDJsonConverter.deserialize(data).as_list()
     res = list(NDJsonConverter.serialize(res))
     assert res == data

tests/data/serialization/ndjson/test_image.py

@@ -20,7 +20,7 @@ def test_image():
     with open('tests/data/assets/ndjson/image_import.json', 'r') as file:
         data = json.load(file)
 
-    res = NDJsonConverter.deserialize(data).as_collection()
+    res = NDJsonConverter.deserialize(data).as_list()
     res = list(NDJsonConverter.serialize(res))
     for r in res:
         r.pop('classifications', None)

tests/data/serialization/ndjson/test_nested.py

@@ -6,6 +6,6 @@
 def test_nested():
     with open('tests/data/assets/ndjson/nested_import.json', 'r') as file:
         data = json.load(file)
-    res = NDJsonConverter.deserialize(data).as_collection()
+    res = NDJsonConverter.deserialize(data).as_list()
     res = list(NDJsonConverter.serialize(res))
     assert res == data

tests/data/serialization/ndjson/test_text.py

@@ -6,6 +6,6 @@
 def test_text():
     with open('tests/data/assets/ndjson/text_import.json', 'r') as file:
         data = json.load(file)
-    res = NDJsonConverter.deserialize(data).as_collection()
+    res = NDJsonConverter.deserialize(data).as_list()
     res = list(NDJsonConverter.serialize(res))
     assert res == data

tests/data/serialization/ndjson/test_video.py

@@ -7,6 +7,6 @@ def test_video():
     with open('tests/data/assets/ndjson/video_import.json', 'r') as file:
         data = json.load(file)
 
-    res = NDJsonConverter.deserialize(data).as_collection()
+    res = NDJsonConverter.deserialize(data).as_list()
     res = list(NDJsonConverter.serialize(res))
     assert res == [data[2], data[0], data[1]]