diff --git a/migrations/006_add_logo_annotation_server_type.py b/migrations/006_add_logo_annotation_server_type.py new file mode 100644 index 0000000000..a50b5f4649 --- /dev/null +++ b/migrations/006_add_logo_annotation_server_type.py @@ -0,0 +1,23 @@ +import peewee as pw +from peewee_migrate import Migrator + + +def migrate(migrator: Migrator, database: pw.Database, *, fake=False): + """Write your migrations here.""" + + migrator.add_fields( + "logo_annotation", + server_type=pw.CharField( + null=True, + max_length=10, + help_text="project associated with the logo annotation, " + "one of 'off', 'obf', 'opff', 'opf', 'off-pro'", + index=False, + ), + ) + + +def rollback(migrator: Migrator, database: pw.Database, *, fake=False): + """Write your rollback migrations here.""" + + migrator.remove_fields("logo_annotation", "server_type") diff --git a/robotoff/cli/logos.py b/robotoff/cli/logos.py index 9fad302dff..b5e5301dd0 100644 --- a/robotoff/cli/logos.py +++ b/robotoff/cli/logos.py @@ -89,6 +89,7 @@ def import_logos( bounding_box=item["bounding_box"], barcode=image_instance.barcode, source_image=image_instance.source_image, + server_type=server_type.name, ) seen_set.add(key) diff --git a/robotoff/models.py b/robotoff/models.py index 6c9f2ec713..64866ccfcb 100644 --- a/robotoff/models.py +++ b/robotoff/models.py @@ -224,8 +224,8 @@ class Prediction(BaseModel): server_type = peewee.CharField( null=False, max_length=10, - help_text="project associated with the insight, " - "one of 'off', 'obf', 'opff', 'opf'", + help_text="project associated with the prediction, " + "one of 'off', 'obf', 'opff', 'opf', 'off-pro'", index=True, default="off", ) @@ -343,6 +343,17 @@ class LogoAnnotation(BaseModel): source_image = peewee.TextField(null=True, index=True) # The logo text extracted from the image using OCR text = peewee.TextField(null=True) + # This is the same as ImageModel.server_type, but we store it here to + # avoid performing a double join with ImageModel table when we need to + # filter logos by server_type + # (LogoAnnotation > ImagePrediction > ImageModel) + server_type = peewee.CharField( + null=True, + max_length=10, + help_text="project associated with the logo annotation, " + "one of 'off', 'obf', 'opff', 'opf', 'off-pro'", + index=False, + ) class Meta: constraints = [peewee.SQL("UNIQUE(image_prediction_id, index)")] diff --git a/robotoff/workers/tasks/import_image.py b/robotoff/workers/tasks/import_image.py index d9fe71f566..fcc43a6adb 100644 --- a/robotoff/workers/tasks/import_image.py +++ b/robotoff/workers/tasks/import_image.py @@ -507,6 +507,7 @@ def run_logo_object_detection( barcode=image_model.barcode, source_image=image_model.source_image, text=text, + server_type=product_id.server_type.name, ) ) 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 6352ba47c7..8c85bb9f44 100644 --- a/scripts/insert_image_predictions.py +++ b/scripts/insert_image_predictions.py @@ -85,6 +85,7 @@ def insert_batch( bounding_box=item["bounding_box"], barcode=image_instance.barcode, source_image=image_instance.source_image, + server_type=server_type.name, ) seen_set.add(key)