diff --git a/mindee/parsing/v2/field/base_field.py b/mindee/parsing/v2/field/base_field.py index 27393ab3..8217f40c 100644 --- a/mindee/parsing/v2/field/base_field.py +++ b/mindee/parsing/v2/field/base_field.py @@ -1,10 +1,33 @@ from typing import List, Optional -from mindee.parsing.v2.field.dynamic_field import DynamicField +from mindee.parsing.common.string_dict import StringDict +from mindee.parsing.v2.field.dynamic_field import DynamicField, FieldType +from mindee.parsing.v2.field.field_confidence import FieldConfidence +from mindee.parsing.v2.field.field_location import FieldLocation class BaseField(DynamicField): """Field with base information.""" - locations: List - confidence: Optional[str] + locations: List[FieldLocation] + confidence: Optional[FieldConfidence] + + def __init__( + self, field_type: FieldType, raw_response: StringDict, indent_level: int = 0 + ) -> None: + super().__init__(field_type, indent_level) + self._indent_level = indent_level + + self.confidence = None + self.locations = [] + + if "confidence" in raw_response and raw_response["confidence"] is not None: + try: + self.confidence = FieldConfidence(raw_response["confidence"]) + except ValueError: + self.confidence = None + + if "locations" in raw_response: + self.locations = [] + for location in raw_response["locations"]: + self.locations.append(FieldLocation(location)) diff --git a/mindee/parsing/v2/field/field_confidence.py b/mindee/parsing/v2/field/field_confidence.py new file mode 100644 index 00000000..0c3bb1b6 --- /dev/null +++ b/mindee/parsing/v2/field/field_confidence.py @@ -0,0 +1,10 @@ +from enum import Enum + + +class FieldConfidence(str, Enum): + """Confidence level of a field as returned by the V2 API.""" + + CERTAIN = "Certain" + HIGH = "High" + MEDIUM = "Medium" + LOW = "Low" diff --git a/mindee/parsing/v2/field/field_location.py b/mindee/parsing/v2/field/field_location.py new file mode 100644 index 00000000..9c23e547 --- /dev/null +++ b/mindee/parsing/v2/field/field_location.py @@ -0,0 +1,31 @@ +from typing import Optional + +from mindee.geometry import Polygon +from mindee.parsing.common.string_dict import StringDict + + +class FieldLocation: + """Location of a field.""" + + def __init__(self, server_response: StringDict) -> None: + """ + Initialize FieldLocation from server response. + + :param server_response: Raw server response. + """ + self.polygon: Optional[Polygon] = None + self.page: Optional[int] = None + + if "polygon" in server_response and server_response["polygon"] is not None: + self.polygon = Polygon(server_response["polygon"]) + + if "page" in server_response and isinstance(server_response["page"], int): + self.page = server_response["page"] + + def __str__(self) -> str: + """ + String representation. + + :return: String representation of the field location. + """ + return str(self.polygon) if self.polygon else "" diff --git a/mindee/parsing/v2/field/list_field.py b/mindee/parsing/v2/field/list_field.py index ce99b389..d26733f3 100644 --- a/mindee/parsing/v2/field/list_field.py +++ b/mindee/parsing/v2/field/list_field.py @@ -1,6 +1,7 @@ from typing import List from mindee.parsing.common.string_dict import StringDict +from mindee.parsing.v2.field.base_field import BaseField from mindee.parsing.v2.field.dynamic_field import ( DynamicField, FieldType, @@ -8,14 +9,14 @@ ) -class ListField(DynamicField): +class ListField(BaseField): """List field containing multiple fields.""" items: List[DynamicField] """Items contained in the list.""" def __init__(self, raw_response: StringDict, indent_level: int = 0): - super().__init__(FieldType.LIST, indent_level) + super().__init__(FieldType.LIST, raw_response, indent_level) self.items = [] for item in raw_response["items"]: diff --git a/mindee/parsing/v2/field/object_field.py b/mindee/parsing/v2/field/object_field.py index 5a751e34..06e90f76 100644 --- a/mindee/parsing/v2/field/object_field.py +++ b/mindee/parsing/v2/field/object_field.py @@ -11,7 +11,7 @@ class ObjectField(BaseField): """Fields contained in the object.""" def __init__(self, raw_response: StringDict, indent_level: int = 0): - super().__init__(FieldType.OBJECT, indent_level) + super().__init__(FieldType.OBJECT, raw_response, indent_level) inner_fields = raw_response.get("fields", raw_response) self.fields = InferenceResultFields(inner_fields, self._indent_level + 1) diff --git a/mindee/parsing/v2/field/simple_field.py b/mindee/parsing/v2/field/simple_field.py index c9cbd296..ce764f98 100644 --- a/mindee/parsing/v2/field/simple_field.py +++ b/mindee/parsing/v2/field/simple_field.py @@ -11,7 +11,7 @@ class SimpleField(BaseField): value: Union[str, float, bool, None] def __init__(self, raw_response: StringDict, indent_level: int = 0): - super().__init__(FieldType.SIMPLE, indent_level) + super().__init__(FieldType.SIMPLE, raw_response, indent_level) value = raw_response.get("value", None) if isinstance(value, int) and not isinstance(raw_response.get("value"), bool): self.value = float(value) diff --git a/tests/data b/tests/data index 632af76d..f0175f0e 160000 --- a/tests/data +++ b/tests/data @@ -1 +1 @@ -Subproject commit 632af76d6eefe551cdeb2b7fa2f325cecec1b66f +Subproject commit f0175f0ee644b57b409e6ad7e1c030f28fbe57ef diff --git a/tests/v2/test_inference_response.py b/tests/v2/test_inference_response.py index 61857b75..5afa7ab5 100644 --- a/tests/v2/test_inference_response.py +++ b/tests/v2/test_inference_response.py @@ -4,6 +4,7 @@ import pytest +from mindee.parsing.v2.field.field_confidence import FieldConfidence from mindee.parsing.v2.field.list_field import ListField from mindee.parsing.v2.field.object_field import ObjectField from mindee.parsing.v2.field.simple_field import SimpleField @@ -185,3 +186,42 @@ def test_full_inference_response(): assert inference_result.inference.file.mime_type == "image/jpeg" assert not inference_result.inference.file.alias assert not inference_result.inference.result.options + + +@pytest.mark.v2 +def test_field_locations_and_confidence() -> None: + """ + Validate that the first location polygon for the ``date`` field is correctly + deserialized together with the associated confidence level. + """ + json_sample, _ = _get_product_samples( + "financial_document", "complete_with_coordinates" + ) + + inference_result = InferenceResponse(json_sample) + + date_field: SimpleField = inference_result.inference.result.fields.date + + assert date_field.locations, "date field should expose locations" + loc0 = date_field.locations[0] + assert loc0 is not None + assert loc0.page == 0 + + polygon = loc0.polygon + assert polygon is not None + assert len(polygon[0]) == 2 + + assert polygon[0][0] == 0.948979073166918 + assert polygon[0][1] == 0.23097924535067715 + + assert polygon[1][0] == 0.85422 + assert polygon[1][1] == 0.230072 + + assert polygon[2][0] == 0.8540899268330819 + assert polygon[2][1] == 0.24365775464932288 + + assert polygon[3][0] == 0.948849 + assert polygon[3][1] == 0.244565 + + assert date_field.confidence == FieldConfidence.MEDIUM + assert str(date_field.confidence.value) == "Medium"