Merge pull request #202 from Labelbox/ms/lbv1-annotation-types

lbv1-converter

Author: msokoloff1

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

labelbox/data/serialization/__init__.py

@@ -0,0 +1 @@
+from .labelbox_v1 import LBV1Converter

labelbox/data/serialization/labelbox_v1/__init__.py

@@ -0,0 +1 @@
+from .converter import LBV1Converter

labelbox/data/serialization/labelbox_v1/classification.py

@@ -0,0 +1,118 @@
+from typing import List, Union
+
+from pydantic.main import BaseModel
+
+from ...annotation_types.annotation import ClassificationAnnotation
+from ...annotation_types.classification import Checklist, ClassificationAnswer, Radio, Text, Dropdown
+from ...annotation_types.types import Cuid
+from .feature import LBV1Feature
+
+
+class LBV1ClassificationAnswer(LBV1Feature):
+    ...
+
+
+class LBV1Radio(LBV1Feature):
+    answer: LBV1ClassificationAnswer
+
+    def to_common(self):
+        return Radio(answer=ClassificationAnswer(
+            schema_id=self.answer.schema_id,
+            name=self.answer.title,
+            extra={
+                'feature_id': self.answer.feature_id,
+                'value': self.answer.value
+            }))
+
+    @classmethod
+    def from_common(cls, radio: Radio, schema_id: Cuid, **extra) -> "LBV1Radio":
+        return cls(schema_id=schema_id,
+                   answer=LBV1ClassificationAnswer(
+                       schema_id=radio.answer.schema_id,
+                       title=radio.answer.name,
+                       value=radio.answer.extra.get('value'),
+                       feature_id=radio.answer.extra.get('feature_id')),
+                   **extra)
+
+
+class LBV1Checklist(LBV1Feature):
+    answers: List[LBV1ClassificationAnswer]
+
+    def to_common(self):
+        return Checklist(answer=[
+            ClassificationAnswer(schema_id=answer.schema_id,
+                                 name=answer.title,
+                                 extra={
+                                     'feature_id': answer.feature_id,
+                                     'value': answer.value
+                                 }) for answer in self.answers
+        ])
+
+    @classmethod
+    def from_common(cls, checklist: Checklist, schema_id: Cuid,
+                    **extra) -> "LBV1Checklist":
+        return cls(schema_id=schema_id,
+                   answers=[
+                       LBV1ClassificationAnswer(
+                           schema_id=answer.schema_id,
+                           title=answer.name,
+                           value=answer.extra.get('value'),
+                           feature_id=answer.extra.get('feature_id'))
+                       for answer in checklist.answer
+                   ],
+                   **extra)
+
+
+class LBV1Text(LBV1Feature):
+    answer: str
+
+    def to_common(self):
+        return Text(answer=self.answer)
+
+    @classmethod
+    def from_common(cls, text: Text, schema_id: Cuid, **extra) -> "LBV1Text":
+        return cls(schema_id=schema_id, answer=text.answer, **extra)
+
+
+class LBV1Classifications(BaseModel):
+    classifications: List[Union[LBV1Radio, LBV1Checklist, LBV1Text]] = []
+
+    def to_common(self) -> List[ClassificationAnnotation]:
+        classifications = [
+            ClassificationAnnotation(value=classification.to_common(),
+                                     classifications=[],
+                                     name=classification.title,
+                                     extra={
+                                         'value': classification.value,
+                                         'feature_id': classification.feature_id
+                                     })
+            for classification in self.classifications
+        ]
+        return classifications
+
+    @classmethod
+    def from_common(
+            cls, annotations: List[ClassificationAnnotation]
+    ) -> "LBV1Classifications":
+        classifications = []
+        for annotation in annotations:
+            classification = cls.lookup_classification(annotation)
+            if classification is not None:
+                classifications.append(
+                    classification.from_common(annotation.value,
+                                               annotation.schema_id,
+                                               **annotation.extra))
+            else:
+                raise TypeError(f"Unexpected type {type(annotation.value)}")
+        return cls(classifications=classifications)
+
+    @staticmethod
+    def lookup_classification(
+        annotation: ClassificationAnnotation
+    ) -> Union[LBV1Text, LBV1Checklist, LBV1Radio]:
+        return {
+            Text: LBV1Text,
+            Dropdown: LBV1Checklist,
+            Checklist: LBV1Checklist,
+            Radio: LBV1Radio
+        }.get(type(annotation.value))

labelbox/data/serialization/labelbox_v1/converter.py

@@ -0,0 +1,97 @@
+from typing import Any, Dict, Generator, Iterable
+import logging
+
+import ndjson
+import requests
+from copy import deepcopy
+from requests.exceptions import HTTPError
+from google.api_core import retry
+
+import labelbox
+from .label import LBV1Label
+from ...annotation_types.collection import (LabelCollection, LabelGenerator,
+                                            PrefetchGenerator)
+
+logger = logging.getLogger(__name__)
+
+
+class LBV1Converter:
+
+    @staticmethod
+    def deserialize_video(json_data: Iterable[Dict[str, Any]],
+                          client: labelbox.Client):
+        """
+        Converts a labelbox video export into the common labelbox format.
+
+        Args:
+            json_data: An iterable representing the labelbox video export.
+            client: The labelbox client for downloading video annotations
+        Returns:
+            LabelGenerator containing the video data.
+        """
+        label_generator = (LBV1Label(**example).to_common()
+                           for example in LBV1VideoIterator(json_data, client))
+        return LabelGenerator(data=label_generator)
+
+    @staticmethod
+    def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator:
+        """
+        Converts a labelbox export (non-video) into the common labelbox format.
+
+        Args:
+            json_data: An iterable representing the labelbox export.
+        Returns:
+            LabelGenerator containing the export data.
+        """
+
+        def label_generator():
+            for example in json_data:
+                if 'frames' in example['Label']:
+                    raise ValueError(
+                        "Use `LBV1Converter.deserialize_video` to process video"
+                    )
+                yield LBV1Label(**example).to_common()
+
+        return LabelGenerator(data=label_generator())
+
+    @staticmethod
+    def serialize(
+            labels: LabelCollection) -> Generator[Dict[str, Any], None, None]:
+        """
+        Converts a labelbox common object to the labelbox json export format
+
+        Args:
+            labels: Either a LabelList or a LabelGenerator (LabelCollection)
+        Returns:
+            A generator for accessing the labelbox json export representation of the data
+        """
+        for label in labels:
+            res = LBV1Label.from_common(label)
+            yield res.dict(by_alias=True)
+
+
+class LBV1VideoIterator(PrefetchGenerator):
+    """
+    Generator that fetches video annotations in the background to be faster.
+    """
+
+    def __init__(self, examples, client):
+        self.client = client
+        super().__init__(examples)
+
+    def _process(self, value):
+        value = deepcopy(value)
+        if 'frames' in value['Label']:
+            req = self._request(value)
+            value['Label'] = ndjson.loads(req)
+        return value
+
+    @retry.Retry(predicate=retry.if_exception_type(HTTPError))
+    def _request(self, value):
+        req = requests.get(
+            value['Label']['frames'],
+            headers={"Authorization": f"Bearer {self.client.api_key}"})
+        if req.status_code == 401:
+            raise labelbox.exceptions.AuthenticationError("Invalid API key")
+        req.raise_for_status()
+        return req.text

labelbox/data/serialization/labelbox_v1/feature.py

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from pydantic import BaseModel, root_validator
+
+from labelbox.utils import camel_case
+from ...annotation_types.types import Cuid
+
+
+class LBV1Feature(BaseModel):
+    keyframe: Optional[bool] = None
+    title: str = None
+    value: Optional[str] = None
+    schema_id: Optional[Cuid] = None
+    feature_id: Optional[Cuid] = None
+
+    @root_validator
+    def check_ids(cls, values):
+        if values.get('value') is None:
+            values['value'] = values['title']
+        return values
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+        # This means these are no video frames ..
+        if self.keyframe is None:
+            res.pop('keyframe')
+        return res
+
+    class Config:
+        allow_population_by_field_name = True
+        alias_generator = camel_case