Merge pull request #203 from Labelbox/ms/ndjson-converter

ndjson converter

Author: msokoloff1

labelbox/data/serialization/__init__.py

@@ -1 +1,2 @@
 from .labelbox_v1 import LBV1Converter
+from .ndjson import NDJsonConverter

labelbox/data/serialization/ndjson/__init__.py

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

labelbox/data/serialization/ndjson/base.py

@@ -0,0 +1,43 @@
+from uuid import uuid4
+from pydantic import BaseModel, validator, Field
+
+from labelbox.utils import camel_case
+from ...annotation_types.types import Cuid
+
+
+class DataRow(BaseModel):
+    id: str = None
+
+    @validator('id', pre=True, always=True)
+    def validate_id(cls, v):
+        if v is None:
+            raise ValueError(
+                "Data row ids are not set. Use `LabelGenerator.add_to_dataset`, `LabelList.add_to_dataset`, or `Label.create_data_row`. "
+                "You can also manually assign the id for each `BaseData` object"
+            )
+        return v
+
+
+class NDJsonBase(BaseModel):
+    uuid: str = None
+    data_row: DataRow
+
+    @validator('uuid', pre=True, always=True)
+    def set_id(cls, v):
+        return v or str(uuid4())
+
+    class Config:
+        allow_population_by_field_name = True
+        alias_generator = camel_case
+
+
+class NDAnnotation(NDJsonBase):
+    schema_id: Cuid
+
+    @validator('schema_id', pre=True, always=True)
+    def validate_id(cls, v):
+        if v is None:
+            raise ValueError(
+                "Schema ids are not set. Use `LabelGenerator.assign_schema_ids`, `LabelList.assign_schema_ids`, or `Label.assign_schema_ids`."
+            )
+        return v

labelbox/data/serialization/ndjson/classification.py

@@ -0,0 +1,212 @@
+from typing import Any, Dict, List, Union, Optional
+
+from pydantic import BaseModel, validator
+
+from labelbox.utils import camel_case
+from ...annotation_types.annotation import ClassificationAnnotation, VideoClassificationAnnotation
+from ...annotation_types.classification.classification import ClassificationAnswer, Dropdown, Text, Checklist, Radio
+from ...annotation_types.types import Cuid
+from ...annotation_types.data import TextData, VideoData, RasterData
+from .base import NDAnnotation
+
+
+class NDFeature(BaseModel):
+    schema_id: Cuid
+
+    @validator('schema_id', pre=True, always=True)
+    def validate_id(cls, v):
+        if v is None:
+            raise ValueError(
+                "Schema ids are not set. Use `LabelGenerator.assign_schema_ids`, `LabelList.assign_schema_ids`, or `Label.assign_schema_ids`."
+            )
+        return v
+
+    class Config:
+        allow_population_by_field_name = True
+        alias_generator = camel_case
+
+
+class FrameLocation(BaseModel):
+    end: int
+    start: int
+
+
+class VideoSupported(BaseModel):
+    #Note that frames are only allowed as top level inferences for video
+    frames: Optional[List[FrameLocation]] = None
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+        # This means these are no video frames ..
+        if self.frames is None:
+            res.pop('frames')
+        return res
+
+
+class NDTextSubclass(NDFeature):
+    answer: str
+
+    def to_common(self) -> Text:
+        return Text(answer=self.answer)
+
+    @classmethod
+    def from_common(cls, text: Text, schema_id: Cuid) -> "NDTextSubclass":
+        return cls(answer=text.answer, schema_id=schema_id)
+
+
+class NDChecklistSubclass(NDFeature):
+    answer: List[NDFeature]
+
+    def to_common(self) -> Checklist:
+        return Checklist(answer=[
+            ClassificationAnswer(schema_id=answer.schema_id)
+            for answer in self.answer
+        ])
+
+    @classmethod
+    def from_common(cls, checklist: Checklist,
+                    schema_id: Cuid) -> "NDChecklistSubclass":
+        return cls(answer=[
+            NDFeature(schema_id=answer.schema_id) for answer in checklist.answer
+        ],
+                   schema_id=schema_id)
+
+
+class NDRadioSubclass(NDFeature):
+    answer: NDFeature
+
+    def to_common(self) -> Radio:
+        return Radio(answer=ClassificationAnswer(
+            schema_id=self.answer.schema_id))
+
+    @classmethod
+    def from_common(cls, radio: Radio, schema_id: Cuid) -> "NDRadioSubclass":
+        return cls(answer=NDFeature(schema_id=radio.answer.schema_id),
+                   schema_id=schema_id)
+
+
+### ====== End of subclasses
+
+
+class NDText(NDAnnotation, NDTextSubclass):
+
+    @classmethod
+    def from_common(cls, text: Text, schema_id: Cuid, extra: Dict[str, Any],
+                    data: Union[TextData, RasterData]) -> "NDText":
+        return cls(
+            answer=text.answer,
+            dataRow={'id': data.uid},
+            schema_id=schema_id,
+            uuid=extra.get('uuid'),
+        )
+
+
+class NDChecklist(NDAnnotation, NDChecklistSubclass, VideoSupported):
+
+    @classmethod
+    def from_common(
+            cls, checklist: Checklist, schema_id: Cuid, extra: Dict[str, Any],
+            data: Union[VideoData, TextData, RasterData]) -> "NDChecklist":
+        return cls(answer=[
+            NDFeature(schema_id=answer.schema_id) for answer in checklist.answer
+        ],
+                   dataRow={'id': data.uid},
+                   schema_id=schema_id,
+                   uuid=extra.get('uuid'),
+                   frames=extra.get('frames'))
+
+
+class NDRadio(NDAnnotation, NDRadioSubclass, VideoSupported):
+
+    @classmethod
+    def from_common(cls, radio: Radio, schema_id: Cuid, extra: Dict[str, Any],
+                    data: Union[VideoData, TextData, RasterData]) -> "NDRadio":
+        return cls(answer=NDFeature(schema_id=radio.answer.schema_id),
+                   dataRow={'id': data.uid},
+                   schema_id=schema_id,
+                   uuid=extra.get('uuid'),
+                   frames=extra.get('frames'))
+
+
+class NDSubclassification:
+
+    @classmethod
+    def from_common(
+        cls, annotation: ClassificationAnnotation
+    ) -> Union[NDTextSubclass, NDChecklistSubclass, NDRadioSubclass]:
+        classify_obj = cls.lookup_subclassification(annotation)
+        if classify_obj is None:
+            raise TypeError(
+                f"Unable to convert object to MAL format. `{type(annotation.value)}`"
+            )
+        return classify_obj.from_common(annotation.value, annotation.schema_id)
+
+    @staticmethod
+    def to_common(
+            annotation: "NDClassificationType") -> ClassificationAnnotation:
+        return ClassificationAnnotation(value=annotation.to_common(),
+                                        schema_id=annotation.schema_id)
+
+    @staticmethod
+    def lookup_subclassification(
+        annotation: ClassificationAnnotation
+    ) -> Union[NDTextSubclass, NDChecklistSubclass, NDRadioSubclass]:
+        if isinstance(annotation, Dropdown):
+            raise TypeError("Dropdowns are not supported for MAL")
+        return {
+            Text: NDTextSubclass,
+            Checklist: NDChecklistSubclass,
+            Radio: NDRadioSubclass,
+        }.get(type(annotation.value))
+
+
+class NDClassification:
+
+    @staticmethod
+    def to_common(
+        annotation: "NDClassificationType"
+    ) -> Union[ClassificationAnnotation, VideoClassificationAnnotation]:
+        common = ClassificationAnnotation(value=annotation.to_common(),
+                                          schema_id=annotation.schema_id,
+                                          extra={'uuid': annotation.uuid})
+        if getattr(annotation, 'frames', None) is None:
+            return [common]
+        results = []
+        for frame in annotation.frames:
+            for idx in range(frame.start, frame.end + 1, 1):
+                results.append(
+                    VideoClassificationAnnotation(frame=idx, **common.dict()))
+        return results
+
+    @classmethod
+    def from_common(
+        cls, annotation: Union[ClassificationAnnotation,
+                               VideoClassificationAnnotation],
+        data: Union[VideoData, TextData, RasterData]
+    ) -> Union[NDTextSubclass, NDChecklistSubclass, NDRadioSubclass]:
+        classify_obj = cls.lookup_classification(annotation)
+        if classify_obj is None:
+            raise TypeError(
+                f"Unable to convert object to MAL format. `{type(annotation.value)}`"
+            )
+        return classify_obj.from_common(annotation.value, annotation.schema_id,
+                                        annotation.extra, data)
+
+    @staticmethod
+    def lookup_classification(
+        annotation: Union[ClassificationAnnotation,
+                          VideoClassificationAnnotation]
+    ) -> Union[NDText, NDChecklist, NDRadio]:
+        if isinstance(annotation, Dropdown):
+            raise TypeError("Dropdowns are not supported for MAL")
+        return {
+            Text: NDText,
+            Checklist: NDChecklist,
+            Radio: NDRadio,
+            Dropdown: NDChecklist,
+        }.get(type(annotation.value))
+
+
+NDSubclassificationType = Union[NDRadioSubclass, NDChecklistSubclass,
+                                NDTextSubclass]
+NDClassificationType = Union[NDRadio, NDChecklist, NDText]

labelbox/data/serialization/ndjson/converter.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Any, Dict, Generator, Iterable
+
+from ...annotation_types.collection import LabelCollection, LabelGenerator
+from .label import NDLabel
+
+logger = logging.getLogger(__name__)
+
+
+class NDJsonConverter:
+
+    @staticmethod
+    def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator:
+        """
+        Converts ndjson data (prediction import format) into the common labelbox format.
+
+        Args:
+            json_data: An iterable representing the ndjson data
+        Returns:
+            LabelGenerator containing the ndjson data.
+        """
+        data = NDLabel(**{'annotations': json_data})
+        return data.to_common()
+
+    @staticmethod
+    def serialize(
+            labels: LabelCollection) -> Generator[Dict[str, Any], None, None]:
+        """
+        Converts a labelbox common object to the labelbox ndjson format (prediction import format)
+
+        Note that this function might fail for objects that are not supported by mal.
+        Not all edge cases are handling by custom exceptions, if you get a cryptic pydantic error message it is probably due to this.
+        We will continue to improve the error messages and add helper functions to deal with this.
+
+        Args:
+            labels: Either a LabelList or a LabelGenerator
+        Returns:
+            A generator for accessing the ndjson representation of the data
+        """
+        for example in NDLabel.from_common(labels):
+            yield example.dict(by_alias=True)