Text orientation in base worker

closes #91 (closed)

arkindex_worker/worker/transcription.py

 # -*- coding: utf-8 -*-
 
+from enum import Enum
+
 from peewee import IntegrityError
 
 from arkindex_worker import logger
@@ -7,8 +9,17 @@ from arkindex_worker.cache import CachedElement, CachedTranscription
 from arkindex_worker.models import Element
 
 
+class TextOrientation(Enum):
+    HorizontalLeftToRight = "horizontal-lr"
+    HorizontalRightToLeft = "horizontal-rl"
+    VerticalRightToLeft = "vertical-rl"
+    VerticalLeftToRight = "vertical-lr"
+
+
 class TranscriptionMixin(object):
-    def create_transcription(self, element, text, score):
+    def create_transcription(
+        self, element, text, score, orientation=TextOrientation.HorizontalLeftToRight
+    ):
         """
         Create a transcription on the given element through the API.
         """
@@ -18,7 +29,9 @@ class TranscriptionMixin(object):
         assert text and isinstance(
             text, str
         ), "text shouldn't be null and should be of type str"
-
+        assert orientation and isinstance(
+            orientation, TextOrientation
+        ), "orientation shouldn't be null and should be of type TextOrientation"
         assert (
             isinstance(score, float) and 0 <= score <= 1
         ), "score shouldn't be null and should be a float in [0..1] range"
@@ -36,6 +49,7 @@ class TranscriptionMixin(object):
                 "text": text,
                 "worker_version": self.worker_version_id,
                 "score": score,
+                "orientation": orientation.value,
             },
         )
 
@@ -50,6 +64,7 @@ class TranscriptionMixin(object):
                         "element_id": element.id,
                         "text": created["text"],
                         "confidence": created["confidence"],
+                        "orientation": created["orientation"],
                         "worker_version_id": self.worker_version_id,
                     }
                 ]
@@ -86,6 +101,13 @@ class TranscriptionMixin(object):
                 score is not None and isinstance(score, float) and 0 <= score <= 1
             ), f"Transcription at index {index} in transcriptions: score shouldn't be null and should be a float in [0..1] range"
 
+            orientation = transcription.get(
+                "orientation", TextOrientation.HorizontalLeftToRight
+            )
+            assert orientation and isinstance(
+                orientation, TextOrientation
+            ), f"Transcription at index {index} in transcriptions: orientation shouldn't be null and should be of type TextOrientation"
+
         created_trs = self.request(
             "CreateTranscriptions",
             body={
@@ -106,6 +128,7 @@ class TranscriptionMixin(object):
                         "element_id": created_tr["element_id"],
                         "text": created_tr["text"],
                         "confidence": created_tr["confidence"],
+                        "orientation": created_tr["orientation"],
                         "worker_version_id": self.worker_version_id,
                     }
                     for created_tr in created_trs
@@ -143,6 +166,14 @@ class TranscriptionMixin(object):
                 score is not None and isinstance(score, float) and 0 <= score <= 1
             ), f"Transcription at index {index} in transcriptions: score shouldn't be null and should be a float in [0..1] range"
 
+            orientation = transcription.get(
+                "orientation", TextOrientation.HorizontalLeftToRight
+            )
+            assert orientation and isinstance(
+                orientation, TextOrientation
+            ), f"Transcription at index {index} in transcriptions: orientation shouldn't be null and should be of type TextOrientation"
+            transcription["orientation"] = orientation
+
             polygon = transcription.get("polygon")
             assert polygon and isinstance(
                 polygon, list
@@ -162,13 +193,23 @@ class TranscriptionMixin(object):
             )
             return
 
+        sent_transcriptions = [
+            {
+                "text": transcription["text"],
+                "score": transcription["score"],
+                "orientation": transcription["orientation"].value,
+                "polygon": transcription["polygon"],
+            }
+            for transcription in transcriptions
+        ]
+
         annotations = self.request(
             "CreateElementTranscriptions",
             id=element.id,
             body={
                 "element_type": sub_element_type,
                 "worker_version": self.worker_version_id,
-                "transcriptions": transcriptions,
+                "transcriptions": sent_transcriptions,
                 "return_elements": True,
             },
         )
@@ -216,6 +257,7 @@ class TranscriptionMixin(object):
                         "element_id": annotation["element_id"],
                         "text": transcription["text"],
                         "confidence": transcription["score"],
+                        "orientation": transcription["orientation"],
                         "worker_version_id": self.worker_version_id,
                     }
                 )