diff --git a/arkindex_worker/worker/transcription.py b/arkindex_worker/worker/transcription.py index 99335804163d7113fdede234fcd06cd42ed1c550..5ddab34c86b950a9ec8796947ed93091e2af106a 100644 --- a/arkindex_worker/worker/transcription.py +++ b/arkindex_worker/worker/transcription.py @@ -227,6 +227,8 @@ class TranscriptionMixin(object): Required. Confidence score between 0 and 1. orientation ([TextOrientation][arkindex_worker.worker.transcription.TextOrientation]) Optional. Orientation of the transcription's text. + element_confidence (float) + Optional. Confidence score of the element between 0 and 1. :returns: A list of dicts as returned by the ``CreateElementTranscriptions`` API endpoint. """ @@ -278,6 +280,12 @@ class TranscriptionMixin(object): assert all( isinstance(coord, (int, float)) for point in polygon for coord in point ), f"Transcription at index {index} in transcriptions: polygon points should be lists of two numbers" + + element_confidence = transcription.get("element_confidence") + assert element_confidence is None or ( + isinstance(element_confidence, float) and 0 <= element_confidence <= 1 + ), f"Transcription at index {index} in transcriptions: element_confidence should be either null or a float in [0..1] range" + if self.is_read_only: logger.warning( "Cannot create transcriptions as this worker is in read-only mode" @@ -327,6 +335,7 @@ class TranscriptionMixin(object): "image_id": element.image_id, "polygon": transcription["polygon"], "worker_run_id": self.worker_run_id, + "confidence": transcription.get("element_confidence"), } ) diff --git a/tests/test_elements_worker/test_transcriptions.py b/tests/test_elements_worker/test_transcriptions.py index d69b6012aaade42ef2cc2824a872bc000a51b85f..70bebc9084220a29ad7cf13352654b2a7694a8da 100644 --- a/tests/test_elements_worker/test_transcriptions.py +++ b/tests/test_elements_worker/test_transcriptions.py @@ -22,6 +22,7 @@ TRANSCRIPTIONS_SAMPLE = [ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]], "confidence": 0.75, "text": "first", + "element_confidence": 0.75, }, { "polygon": [[1000, 300], [1200, 300], [1200, 500], [1000, 500]], @@ -1244,6 +1245,29 @@ def test_create_element_transcriptions_wrong_transcriptions(mock_elements_worker == "Transcription at index 1 in transcriptions: orientation shouldn't be null and should be of type TextOrientation" ) + with pytest.raises(AssertionError) as e: + mock_elements_worker.create_element_transcriptions( + element=elt, + sub_element_type="page", + transcriptions=[ + { + "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]], + "confidence": 0.75, + "text": "The", + }, + { + "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]], + "confidence": 0.75, + "text": "word", + "element_confidence": "not a confidence", + }, + ], + ) + assert ( + str(e.value) + == "Transcription at index 1 in transcriptions: element_confidence should be either null or a float in [0..1] range" + ) + def test_create_element_transcriptions_api_error(responses, mock_elements_worker): elt = Element({"id": "12341234-1234-1234-1234-123412341234"}) @@ -1326,6 +1350,7 @@ def test_create_element_transcriptions(responses, mock_elements_worker): "confidence": 0.75, "text": "first", "orientation": TextOrientation.HorizontalLeftToRight.value, + "element_confidence": 0.75, }, { "polygon": [[1000, 300], [1200, 300], [1200, 500], [1000, 500]], @@ -1411,6 +1436,7 @@ def test_create_element_transcriptions_with_cache( "confidence": 0.75, "text": "first", "orientation": TextOrientation.HorizontalLeftToRight.value, + "element_confidence": 0.75, }, { "polygon": [[1000, 300], [1200, 300], [1200, 500], [1000, 500]], @@ -1454,6 +1480,7 @@ def test_create_element_transcriptions_with_cache( type="page", polygon="[[0, 0], [2000, 0], [2000, 3000], [0, 3000]]", worker_run_id=UUID("56785678-5678-5678-5678-567856785678"), + confidence=0.75, ), ] assert list(CachedTranscription.select()) == [