feat: Improve annotation process for nutrition_extraction insight

robotoff/insights/annotate.py

@@ -28,6 +28,7 @@
     update_expiration_date,
     update_quantity,
 )
+from robotoff.prediction.utils import get_image_rotation, get_nutrition_table_prediction
 from robotoff.products import get_image_id, get_product
 from robotoff.types import InsightAnnotation, InsightType, JSONType, NutrientData
 from robotoff.utils import get_logger
@@ -734,6 +735,7 @@ def process_annotation(
         if is_vote:
             return CANNOT_VOTE_RESULT
 
+        insight_updated = False
         # The annotator can change the nutrient values to fix the model errors
         if data is not None:
             try:
@@ -749,22 +751,32 @@ def process_annotation(
             # user
             insight.data["annotation"] = validated_nutrients.model_dump()
             insight.data["was_updated"] = True
-            insight.save()
+            insight_updated = True
         else:
             validated_nutrients = NutrientData.model_validate(insight.data)
             validated_nutrients = cls.add_default_unit(validated_nutrients)
 
             insight.data["annotation"] = validated_nutrients.model_dump()
             insight.data["was_updated"] = False
+            insight_updated = True
+
+        product_id = insight.get_product_id()
+        product = get_product(product_id, ["code", "images", "lang"])
+
+        if product is None:
+            return MISSING_PRODUCT_RESULT
+
+        if insight_updated:
             insight.save()
 
         save_nutrients(
-            product_id=insight.get_product_id(),
+            product_id=product_id,
             nutrient_data=validated_nutrients,
             insight_id=insight.id,
             auth=auth,
             is_vote=is_vote,
         )
+        cls.select_nutrition_image(insight, product, auth)
         return UPDATED_ANNOTATION_RESULT
 
     @classmethod
@@ -788,6 +800,80 @@ def validate_data(cls, data: JSONType) -> NutrientData:
             raise ValidationError("missing 'nutrients' field")
         return NutrientData.model_validate(data)
 
+    @classmethod
+    def select_nutrition_image(
+        cls,
+        insight: ProductInsight,
+        product: JSONType,
+        auth: OFFAuthentication | None = None,
+    ) -> None:
+        """If the insight is validated, select the source image as nutrition image.
+
+        We fetch the image orientation from the `predictions` table and the prediction
+        of the nutrition table detector from the `image_predictions` table to know the
+        rotation angle and the bounding box of the nutrition table.
+        If any of these predictions are missing, we just select the image without any
+        rotation or crop bounding box.
+
+        :param insight: the original `nutrient_extraction` insight
+        :param product: the product data
+        :param auth: the user authentication data
+        """
+
+        if insight.source_image is None:
+            return None
+
+        image_id = get_image_id(insight.source_image)
+        images = product.get("images", {})
+        image_meta: JSONType | None = images.get(image_id)
+
+        if not image_id or not image_meta:
+            return None
+
+        # Use the language of the product. This field should always be available,
+        # but we provide a default value just in case.
+        lang = product.get("lang", "en")
+        image_key = f"nutrition_{lang}"
+        # We don't want to select the nutrition image if one has already been
+        # selected
+        if image_key in images:
+            return None
+
+        rotation = get_image_rotation(insight.source_image)
+
+        nutrition_table_detections = get_nutrition_table_prediction(
+            insight.source_image, threshold=0.5
+        )
+        bounding_box = None
+        # Only crop according to the model predicted bounding box if there is exactly
+        # one nutrition table detected
+        if nutrition_table_detections and len(nutrition_table_detections) == 1:
+            bounding_box = nutrition_table_detections[0]["bounding_box"]
+
+        crop_bounding_box: tuple[float, float, float, float] | None = None
+        if bounding_box:
+            rotation = rotation or 0
+            # convert crop bounding box to the format expected by Product
+            # Opener
+            image_size = image_meta["sizes"]["full"]
+            width = image_size["w"]
+            height = image_size["h"]
+            crop_bounding_box = convert_crop_bounding_box(
+                bounding_box, width, height, rotation
+            )
+
+        product_id = insight.get_product_id()
+        select_rotate_image(
+            product_id=product_id,
+            image_id=image_id,
+            image_key=image_key,
+            rotate=rotation,
+            crop_bounding_box=crop_bounding_box,
+            auth=auth,
+            is_vote=False,
+            insight_id=insight.id,
+        )
+
 
 ANNOTATOR_MAPPING: dict[str, Type] = {
     InsightType.packager_code.name: PackagerCodeAnnotator,

robotoff/prediction/nutrition_extraction.py

@@ -569,6 +569,11 @@ def match_nutrient_value(
             )
         ) and (
             value in ("08", "09")
+            or (
+                len(value) > 2
+                and "." not in value
+                and (value.endswith("8") or (value.endswith("9")))
+            )
             or (value.endswith("8") and "." in value and not value.endswith(".8"))
             or (value.endswith("9") and "." in value and not value.endswith(".9"))
         ):

robotoff/prediction/ocr/image_lang.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, TypedDict, Union
 
 from openfoodfacts.ocr import OCRResult
 
@@ -9,6 +9,11 @@
 PREDICTOR_VERSION = "1"
 
 
+class ImageLangDataType(TypedDict):
+    count: dict[str, int]
+    percent: dict[str, float]
+
+
 def get_image_lang(ocr_result: Union[OCRResult, str]) -> list[Prediction]:
     if isinstance(ocr_result, str):
         return []

robotoff/prediction/utils.py

@@ -0,0 +1,74 @@
+from robotoff.models import ImageModel, ImagePrediction, Prediction
+from robotoff.prediction.ocr.image_lang import ImageLangDataType
+from robotoff.types import ObjectDetectionModel, PredictionType
+
+
+def get_image_rotation(source_image: str) -> int | None:
+    """Return the image rotation of the image, by fetching the associated
+    `image_orientation` prediction from the DB.
+
+    The image orientation is represented by a rotation angle in degrees:
+    - 0: upright
+    - 90: left
+    - 180: upside down
+    - 270: right
+
+    If no prediction is found, return None.
+
+    :param source_image: the source image of the prediction
+    :return: the rotation angle of the image, or None if no prediction is found
+    """
+    image_orientation_prediction = Prediction.get_or_none(
+        Prediction.type == PredictionType.image_orientation,
+        Prediction.source_image == source_image,
+    )
+
+    if image_orientation_prediction is None:
+        return None
+
+    return image_orientation_prediction.data["rotation"]
+
+
+def get_image_lang(source_image: str) -> ImageLangDataType | None:
+    """Return the name of the language detected in the image, by fetching the
+    associated `image_lang` prediction from the DB.
+
+    If no prediction is found, return None.
+
+    :param source_image: the source image of the prediction
+    :return: the name of the language detected in the image, or None if no prediction
+        is found
+    """
+    image_lang_prediction = Prediction.get_or_none(
+        Prediction.type == PredictionType.image_lang,
+        Prediction.source_image == source_image,
+    )
+
+    if image_lang_prediction is None:
+        return None
+
+    return image_lang_prediction.data
+
+
+def get_nutrition_table_prediction(
+    source_image: str, threshold: float = 0.5
+) -> list | None:
+    """Return the nutrition table prediction associated with the image.
+
+    :param source_image: the source image of the prediction
+    :return: the nutrition table prediction associated with the image
+    """
+    image_prediction = (
+        ImagePrediction.select()
+        .join(ImageModel)
+        .where(
+            ImagePrediction.model_name == ObjectDetectionModel.nutrition_table.name,
+            ImageModel.source_image == source_image,
+        )
+    ).get_or_none()
+
+    if image_prediction is None:
+        return None
+
+    objects = image_prediction.data["objects"]
+    return [obj for obj in objects if obj["score"] >= threshold]