Fix: Add transformers for videos and bounding boxes

.github/workflows/pr.yaml

@@ -17,7 +17,9 @@ jobs:
     steps:
 
     - name: Cleanup previous run
-      run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb"
+      run: |
+        docker ps -q -a --filter "name=${{ runner.name }}_pr" | xargs -r docker rm -f || true
+        docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb" || true
       continue-on-error: true
 
     - uses: actions/checkout@v3
@@ -80,7 +82,9 @@ jobs:
     steps:
 
     - name: Cleanup previous run
-      run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb"
+      run: |
+        docker ps -q -a --filter "name=${{ runner.name }}_${{ github.run_id }}" | xargs -r docker rm -f || true
+        docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb" || true
       continue-on-error: true
 
     - uses: actions/checkout@v3
@@ -101,6 +105,7 @@ jobs:
       uses: google-github-actions/setup-gcloud@v2
     - name: Build tests on pytorch GPU image
       run: |
+        docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb" || true
         rm -rf docker/pytorch-gpu/aperturedata
         mkdir -p docker/pytorch-gpu/aperturedata
         cp -r aperturedb pyproject.toml README.md test docker/pytorch-gpu/aperturedata

aperturedb/transformers/__init__.py

@@ -0,0 +1,36 @@
+def __getattr__(name):
+    if name == "Transformer":
+        from .transformer import Transformer
+        val = Transformer
+    elif name == "CommonProperties":
+        from .common_properties import CommonProperties
+        val = CommonProperties
+    elif name == "ImageProperties":
+        from .image_properties import ImageProperties
+        val = ImageProperties
+    elif name == "VideoProperties":
+        from .video_properties import VideoProperties
+        val = VideoProperties
+    elif name == "BoundingBoxProperties":
+        from .bounding_box_properties import BoundingBoxProperties
+        val = BoundingBoxProperties
+    elif name == "FacenetPyTorchEmbeddings":
+        from .facenet_pytorch_embeddings import FacenetPyTorchEmbeddings
+        val = FacenetPyTorchEmbeddings
+    elif name == "CLIPPyTorchEmbeddings":
+        from .clip_pytorch_embeddings import CLIPPyTorchEmbeddings
+        val = CLIPPyTorchEmbeddings
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    globals()[name] = val
+    return val
+
+
+__all__ = [
+    "Transformer",
+    "CommonProperties",
+    "ImageProperties",
+    "VideoProperties",
+    "BoundingBoxProperties",
+]

aperturedb/transformers/bounding_box_properties.py

@@ -0,0 +1,41 @@
+from aperturedb.transformers.transformer import Transformer
+from aperturedb.Subscriptable import Subscriptable
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BoundingBoxProperties(Transformer):
+    """
+    This transformer applies static annotation metadata (like annotation_source
+    and annotation_mode) to bounding boxes and polygons.
+    """
+
+    def __init__(self, data: Subscriptable, **kwargs) -> None:
+        super().__init__(data, **kwargs)
+        self.annotation_source = kwargs.get("annotation_source", None)
+        self.annotation_mode = kwargs.get("annotation_mode", None)
+
+    def getitem(self, subscript):
+        if self.annotation_source is None and self.annotation_mode is None:
+            return self.data[subscript]
+
+        x = self.data[subscript]
+        # Iterate over current transaction commands to handle variable annotation counts
+        for cmd_dict in x[0]:
+            try:
+                if not isinstance(cmd_dict, dict):
+                    continue
+                for cmd_name in ["AddBoundingBox", "AddPolygon"]:
+                    if cmd_name in cmd_dict:
+                        src_properties = cmd_dict[cmd_name].setdefault(
+                            "properties", {})
+                        if self.annotation_source is not None:
+                            src_properties["annotation_source"] = self.annotation_source
+                        if self.annotation_mode is not None:
+                            src_properties["annotation_mode"] = self.annotation_mode
+            except Exception:
+                logger.exception(
+                    "Error applying bounding box properties", stack_info=True)
+
+        return x

aperturedb/transformers/clip.py

@@ -1,7 +1,5 @@
 
 import logging
-import numpy as np
-from PIL import Image
 
 logger = logging.getLogger(__name__)
 
@@ -13,12 +11,14 @@
 """
 
 try:
+    import numpy as np
+    from PIL import Image
     import clip
     import torch
     import cv2
-except ImportError:
+except ImportError as e:
     logger.critical(error_message)
-    exit(1)
+    raise ImportError(error_message) from e
 
 descriptor_set = "ViT-B/16"
 device = "cuda" if torch.cuda.is_available() else "cpu"

aperturedb/transformers/clip_pytorch_embeddings.py

@@ -1,8 +1,11 @@
 import hashlib
+import logging
 from aperturedb.Subscriptable import Subscriptable
 from aperturedb.transformers.transformer import Transformer
 from .clip import generate_embedding, descriptor_set
 
+logger = logging.getLogger(__name__)
+
 
 class CLIPPyTorchEmbeddings(Transformer):
     """
@@ -19,38 +22,82 @@ def __init__(self, data: Subscriptable, **kwargs) -> None:
         self.search_set_name = kwargs.pop(
             "search_set_name", descriptor_set)
         super().__init__(data, **kwargs)
-
-        # Let's sample some data to figure out the descriptorset we need.
-        if len(self._add_image_index) > 0:
-            sample = generate_embedding(self.data[0][1][0])
-            utils = self.get_utils()
-            utils.add_descriptorset(
-                self.search_set_name, dim=len(sample) // 4, metric=["CS"])
+        self._descriptorset_initialized = False
 
     def getitem(self, subscript):
         x = self.data[subscript]
 
-        for ic in self._add_image_index:
-            serialized = generate_embedding(x[1][ic])
-            # If the image already has an image_sha256, we use it.
-            image_sha256 = x[0][ic]["AddImage"].get("properties", {}).get(
-                "adb_image_sha256", None)
-            if not image_sha256:
-                image_sha256 = hashlib.sha256(x[1][ic]).hexdigest()
-            x[1].append(serialized)
-            x[0].append(
-                {
-                    "AddDescriptor": {
-                        "set": self.search_set_name,
-                        "properties": {
-                            "image_sha256": image_sha256,
-                        },
-                        "if_not_found": {
-                            "image_sha256": ["==", image_sha256],
-                        },
-                        "connect": {
-                            "ref": x[0][ic]["AddImage"]["_ref"]
+        blob_index = 0
+        new_descriptors = []
+        new_blobs = []
+
+        for cmd_dict in x[0]:
+            cmd_name = None
+            if isinstance(cmd_dict, dict) and len(cmd_dict) > 0:
+                cmd_name = next(iter(cmd_dict.keys()))
+
+            if cmd_name == "AddImage":
+                try:
+                    blob = x[1][blob_index]
+                except IndexError:
+                    logger.warning(
+                        f"Missing blob for AddImage at index {blob_index}. Stopping blob processing for this transaction.")
+                    break
+
+                serialized = None
+
+                if (
+                    not getattr(self, "_descriptorset_initialized", False)
+                    and isinstance(cmd_dict["AddImage"], dict)
+                    and "_ref" in cmd_dict["AddImage"]
+                ):
+                    try:
+                        serialized = generate_embedding(blob)
+                        dim = len(serialized) // 4
+                        utils = self.get_utils()
+                        success = utils.add_descriptorset(
+                            self.search_set_name, dim=dim, metric=["CS"])
+                        if success or self.search_set_name in utils.get_descriptorset_list():
+                            self._descriptorset_initialized = True
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to initialize descriptorset: {e}", exc_info=True)
+
+                # If the image already has an image_sha256, we use it.
+                if (
+                    getattr(self, "_descriptorset_initialized", False)
+                    and isinstance(cmd_dict["AddImage"], dict)
+                    and "_ref" in cmd_dict["AddImage"]
+                ):
+                    try:
+                        if serialized is None:
+                            serialized = generate_embedding(blob)
+                        image_sha256 = cmd_dict["AddImage"].get("properties", {}).get(
+                            "adb_image_sha256", None)
+                        if not image_sha256:
+                            image_sha256 = hashlib.sha256(blob).hexdigest()
+                        new_blobs.append(serialized)
+                        desc_cmd = {
+                            "AddDescriptor": {
+                                "set": self.search_set_name,
+                                "properties": {
+                                    "image_sha256": image_sha256,
+                                },
+                                "if_not_found": {
+                                    "image_sha256": ["==", image_sha256],
+                                },
+                                "connect": {
+                                    "ref": cmd_dict["AddImage"]["_ref"]
+                                }
+                            }
                         }
-                    }
-                })
+                        new_descriptors.append(desc_cmd)
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to generate embedding or descriptor: {e}", exc_info=True)
+            if cmd_name in ["AddImage", "AddDescriptor", "AddVideo", "AddBlob"]:
+                blob_index += 1
+
+        x[0].extend(new_descriptors)
+        x[1].extend(new_blobs)
         return x

aperturedb/transformers/common_properties.py

@@ -26,20 +26,34 @@ def __init__(self, data: Subscriptable, **kwargs) -> None:
         self.adb_timestamp = kwargs.get("adb_timestamp", None)
         self.adb_main_object = kwargs.get("adb_main_object", None)
 
+    def _apply_common_properties(self, properties: dict):
+        if self.adb_data_source is not None:
+            properties["adb_data_source"] = self.adb_data_source
+        if self.adb_timestamp is not None:
+            properties["adb_timestamp"] = self.adb_timestamp
+        if self.adb_main_object is not None:
+            properties["adb_main_object"] = self.adb_main_object
+
     def getitem(self, subscript):
+        if (
+            self.adb_data_source is None and
+            self.adb_timestamp is None and
+            self.adb_main_object is None
+        ):
+            return self.data[subscript]
+
         x = self.data[subscript]
-        try:
-            # x is a transaction that has an add_image command and a blob
-            for ic in self._add_image_index:
-                src_properties = x[0][ic]["AddImage"]["properties"]
-                # Set the static properties, if explicitly set
-                if self.adb_data_source:
-                    src_properties["adb_data_source"] = self.adb_data_source
-                if self.adb_timestamp:
-                    src_properties["adb_timestamp"] = self.adb_timestamp
-                if self.adb_main_object:
-                    src_properties["adb_main_object"] = self.adb_main_object
-        except Exception as e:
-            logger.exception(e.with_traceback(), stack_info=True)
+        for cmd_dict in x[0]:
+            try:
+                if not isinstance(cmd_dict, dict):
+                    continue
+                for cmd_name in ["AddImage", "AddVideo", "AddBoundingBox", "AddPolygon"]:
+                    if cmd_name in cmd_dict:
+                        src_properties = cmd_dict[cmd_name].setdefault(
+                            "properties", {})
+                        self._apply_common_properties(src_properties)
+            except Exception:
+                logger.exception(
+                    "Error applying common properties", stack_info=True)
 
         return x

aperturedb/transformers/facenet.py

@@ -12,9 +12,9 @@
 try:
     from facenet_pytorch import MTCNN, InceptionResnetV1
     import torch
-except ImportError:
+except ImportError as e:
     logger.critical(error_message)
-    exit(1)
+    raise ImportError(error_message) from e
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')