diff --git a/robotoff/cli/logos.py b/robotoff/cli/logos.py index 1d91cbc58a..0ac4f697c9 100644 --- a/robotoff/cli/logos.py +++ b/robotoff/cli/logos.py @@ -67,6 +67,8 @@ def insert_batch(data_path: pathlib.Path, model_name: str, model_version: str) - index=i, score=item["score"], bounding_box=item["bounding_box"], + barcode=image_instance.barcode, + source_image=image_instance.source_image, ) seen_set.add(key) diff --git a/robotoff/models.py b/robotoff/models.py index c8584734ef..b26c62d223 100644 --- a/robotoff/models.py +++ b/robotoff/models.py @@ -160,6 +160,8 @@ class ProductInsight(BaseModel): # Hunger Games) on a subset of products. Each product have 0+ campaign # tags campaign = BinaryJSONField(null=True, index=True, default=list) + # Confidence score of the insight, may be null + confidence = peewee.FloatField(null=True, index=True) class Prediction(BaseModel): @@ -175,6 +177,7 @@ class Prediction(BaseModel): help_text="server domain linked to the insight", index=True ) predictor = peewee.CharField(max_length=100, null=True) + confidence = peewee.FloatField(null=True, index=False) class AnnotationVote(BaseModel): @@ -267,6 +270,8 @@ class LogoAnnotation(BaseModel): username = peewee.TextField(null=True, index=True) completed_at = peewee.DateTimeField(null=True, index=True) nearest_neighbors = BinaryJSONField(null=True) + barcode = peewee.CharField(max_length=100, null=True, index=True) + source_image = peewee.TextField(null=True, index=True) class Meta: constraints = [peewee.SQL("UNIQUE(image_prediction_id, index)")] diff --git a/robotoff/prediction/types.py b/robotoff/prediction/types.py index 048dd60a0a..f326ff3afd 100644 --- a/robotoff/prediction/types.py +++ b/robotoff/prediction/types.py @@ -18,6 +18,7 @@ class Prediction: source_image: Optional[str] = None server_domain: Optional[str] = None id: Optional[int] = None + confidence: Optional[float] = None def to_dict(self) -> dict[str, Any]: return dataclasses.asdict(self, dict_factory=dict_factory) diff --git a/robotoff/workers/tasks/import_image.py b/robotoff/workers/tasks/import_image.py index 9740022e16..a64b9e2da1 100644 --- a/robotoff/workers/tasks/import_image.py +++ b/robotoff/workers/tasks/import_image.py @@ -368,6 +368,8 @@ def run_logo_object_detection( index=i, score=item["score"], bounding_box=item["bounding_box"], + barcode=image_model.barcode, + source_image=image_model.source_image, ) ) logger.info("%s logos found for image %s", len(logos), source_image) diff --git a/scripts/insert_image_predictions.py b/scripts/insert_image_predictions.py index 466f0765af..1cb65df9a5 100644 --- a/scripts/insert_image_predictions.py +++ b/scripts/insert_image_predictions.py @@ -71,6 +71,8 @@ def insert_batch(data_path: pathlib.Path, model_name: str, model_version: str) - index=i, score=item["score"], bounding_box=item["bounding_box"], + barcode=image_instance.barcode, + source_image=image_instance.source_image, ) seen_set.add(key) diff --git a/tests/integration/models_utils.py b/tests/integration/models_utils.py index f6195ec957..9cd1e7adc3 100644 --- a/tests/integration/models_utils.py +++ b/tests/integration/models_utils.py @@ -67,6 +67,7 @@ class Meta: server_domain = factory.LazyFunction(lambda: settings.OFF_SERVER_DOMAIN) automatic_processing = None predictor = None + confidence = None class AnnotationVoteFactory(UuidSequencer, PeeweeModelFactory):