fix: add logo_annotation.server_type field

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")

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)
 

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)")]

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)

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)