feat: add nutrition extractor

pyproject.toml

@@ -83,6 +83,8 @@ duckdb = "~1.0.0"
 google-cloud-storage = "~2.14.0"
 pandas = "~2.2.2"
 pyarrow = "~17.0.0"
+# Used for CLI pretty print
+rich = "~13.9.2"
 
 [tool.poetry.dependencies.sentry-sdk]
 version = ">=1.14,<2.9"

robotoff/cli/main.py

@@ -93,12 +93,19 @@ def create_redis_update(
 
     get_logger()
     client = get_redis_client()
+    flavor_to_product_type = {
+        "off": "food",
+        "obf": "beauty",
+        "opff": "petfood",
+        "opf": "product",
+    }
     event = {
         "code": barcode,
         "flavor": flavor,
         "user_id": user_id,
         "action": action,
         "comment": comment,
+        "product_type": flavor_to_product_type[flavor],
     }
 
     diffs: JSONType
@@ -602,6 +609,45 @@ def run_object_detection_model(
             )
 
 
+@app.command()
+def run_nutrition_extraction(
+    image_url: str = typer.Argument(
+        ..., help="URL of the image to run nutrition extraction on"
+    ),
+    triton_uri: Optional[str] = typer.Option(
+        None,
+        help="URI of the Triton Inference Server to use. If not provided, the default value from settings is used.",
+    ),
+) -> None:
+    """Run nutrition extraction on a product image.
+
+    The image URL should be an Open Food Facts image URL, e.g.
+    https://images.openfoodfacts.org/images/products/327/408/000/5003/3.jpg
+
+    The OCR JSON is expected to be available at the same URL with a `.json`
+    extension, e.g.
+    https://images.openfoodfacts.org/images/products/327/408/000/5003/3.json
+
+    Prediction is printed to stdout.
+    """
+    from typing import cast
+
+    from openfoodfacts.ocr import OCRResult
+    from PIL import Image
+    from rich import print as pprint
+
+    from robotoff.images import get_image_from_url
+    from robotoff.prediction.nutrition_extraction import predict
+
+    image = cast(Image.Image, get_image_from_url(image_url))
+    ocr_result = cast(OCRResult, OCRResult.from_url(image_url.replace(".jpg", ".json")))
+    prediction = predict(image, ocr_result, triton_uri=triton_uri)
+    if prediction is not None:
+        pprint(prediction)
+    else:
+        pprint("No prediction")
+
+
 @app.command()
 def init_elasticsearch() -> None:
     """This command is used for index creation."""

robotoff/insights/annotate.py

@@ -660,24 +660,6 @@ def process_annotation(
         return UPDATED_ANNOTATION_RESULT
 
 
-class NutritionTableStructureAnnotator(InsightAnnotator):
-    @classmethod
-    def process_annotation(
-        cls,
-        insight: ProductInsight,
-        data: Optional[dict] = None,
-        auth: Optional[OFFAuthentication] = None,
-        is_vote: bool = False,
-    ) -> AnnotationResult:
-        insight.data["annotation"] = data
-        insight.save()
-        return SAVED_ANNOTATION_RESULT
-
-    @classmethod
-    def is_data_required(cls) -> bool:
-        return True
-
-
 class IngredientSpellcheckAnnotator(InsightAnnotator):
     @classmethod
     def process_annotation(
@@ -687,7 +669,8 @@ def process_annotation(
         auth: Optional[OFFAuthentication] = None,
         is_vote: bool = False,
     ) -> AnnotationResult:
-        # Possibility for the annotator to change the spellcheck correction if data is provided
+        # Possibility for the annotator to change the spellcheck correction if data is
+        # provided
         if data is not None:
             annotation = data.get("annotation")
             if not annotation or len(data) > 1:
@@ -720,7 +703,6 @@ def process_annotation(
     InsightType.store.name: StoreAnnotator,
     InsightType.packaging.name: PackagingAnnotator,
     InsightType.nutrition_image.name: NutritionImageAnnotator,
-    InsightType.nutrition_table_structure.name: NutritionTableStructureAnnotator,
     InsightType.is_upc_image.name: UPCImageAnnotator,
     InsightType.ingredient_spellcheck.name: IngredientSpellcheckAnnotator,
 }

robotoff/insights/importer.py

@@ -1524,6 +1524,33 @@ def _keep_prediction(
         )
 
 
+class NutrientExtractionImporter(InsightImporter):
+    @staticmethod
+    def get_type() -> InsightType:
+        return InsightType.nutrient_extraction
+
+    @classmethod
+    def get_required_prediction_types(cls) -> set[PredictionType]:
+        return {PredictionType.nutrient_extraction}
+
+    @classmethod
+    def generate_candidates(
+        cls,
+        product: Optional[Product],
+        predictions: list[Prediction],
+        product_id: ProductIdentifier,
+    ) -> Iterator[ProductInsight]:
+        for prediction in predictions:
+            yield ProductInsight(**prediction.to_dict())
+
+    @classmethod
+    def is_conflicting_insight(
+        cls, candidate: ProductInsight, reference: ProductInsight
+    ) -> bool:
+        # Only one insight per product
+        return True
+
+
 class PackagingElementTaxonomyException(Exception):
     pass
 
@@ -1860,6 +1887,7 @@ def import_product_predictions(
     UPCImageImporter,
     NutritionImageImporter,
     IngredientSpellcheckImporter,
+    NutrientExtractionImporter,
 ]
 
 

robotoff/off.py

@@ -68,8 +68,16 @@ def get_username(self) -> Optional[str]:
         return None
 
 
-def get_source_from_url(ocr_url: str) -> str:
-    url_path = urlparse(ocr_url).path
+def get_source_from_url(url: str) -> str:
+    """Get the `source_image` field from an image or OCR URL.
+
+    It's the path of the image or OCR JSON file, but without the `/images/products`
+    prefix. It always ends with `.jpg`, whather it's an image or an OCR JSON file.
+
+    :param url: the URL of the image or OCR JSON file
+    :return: the source image path
+    """
+    url_path = urlparse(url).path
 
     if url_path.startswith("/images/products"):
         url_path = url_path[len("/images/products") :]