feat: Add CSV ingestion support for Frames

aperturedb/DataModels.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated, List
-from typing import ClassVar
+from typing import ClassVar, Optional
 from uuid import uuid4
 from aperturedb.Query import ObjectType, PropertyType, RangeType
 
@@ -73,6 +73,8 @@ class PolygonDataModel(IdentityDataModel):
 class FrameDataModel(IdentityDataModel):
     """Frame data model for ApertureDB.
     """
+    url: Annotated[Optional[str], Field(
+        title="URL", description="URL to file, http, s3 or gs resource")] = None
     type = ObjectType.FRAME
 
 

aperturedb/Entities.py

@@ -278,6 +278,7 @@ def get_blob(self, entity) -> Any:
 def load_entities_registry(custom_entities: List[str] = None) -> dict:
     from aperturedb.Polygons import Polygons
     from aperturedb.Images import Images
+    from aperturedb.Frames import Frames
     from aperturedb.Blobs import Blobs
     from aperturedb.BoundingBoxes import BoundingBoxes
     from aperturedb.Videos import Videos
@@ -287,6 +288,7 @@ def load_entities_registry(custom_entities: List[str] = None) -> dict:
     known_entities = {
         ObjectType.POLYGON.value: Polygons,
         ObjectType.IMAGE.value: Images,
+        ObjectType.FRAME.value: Frames,
         ObjectType.VIDEO.value: Videos,
         ObjectType.BOUNDING_BOX.value: BoundingBoxes,
         ObjectType.BLOB.value: Blobs,

aperturedb/FrameDataCSV.py

@@ -0,0 +1,19 @@
+from aperturedb.ImageDataCSV import ImageDataCSV
+from aperturedb.Query import ObjectType
+
+
+class FrameDataCSV(ImageDataCSV):
+    """
+    **Helper class to ingest Frame data from a CSV file.**
+
+    This class extends ImageDataCSV and sets the insertion command to "AddFrame",
+    allowing frame files to be batch ingested from CSVs just like images.
+    """
+    command = "AddFrame"
+
+    def get_indices(self):
+        return {
+            "entity": {
+                ObjectType.FRAME.value: self.get_indexed_properties()
+            }
+        }

aperturedb/Frames.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from aperturedb.Images import Images
+from aperturedb.Query import ObjectType
+
+
+class Frames(Images):
+    """
+    **The python wrapper of frame images in ApertureDB.**
+
+    Frames in ApertureDB are quite similar to images and so
+    are modeled in python as a subclass.
+
+
+    Args:
+        client: The database connector, perhaps as returned by `CommonLibrary.create_connector`
+    """
+    db_object = ObjectType.FRAME
+
+    def __init__(self, client, batch_size=100, response=None, **kwargs):
+        super().__init__(
+            client, batch_size=batch_size, response=response, **kwargs)

aperturedb/ImageDataCSV.py

@@ -171,6 +171,8 @@ class ImageDataCSV(CSVParser.CSVParser, ImageDataProcessor):
     id would be only inserted if it does not already exist in the database.
     :::
     """
+    command = "AddImage"
+
 
     def __init__(self, filename: str, check_image: bool = True, n_download_retries: int = 3, **kwargs):
 
@@ -199,8 +201,6 @@ def __init__(self, filename: str, check_image: bool = True, n_download_retries:
         self.relative_path_prefix = os.path.dirname(self.filename) \
             if self.source_type == HEADER_PATH and self.blobs_relative_to_csv else ""
 
-        self.command = "AddImage"
-
     def getitem(self, idx):
         idx = self.df.index.start + idx
 

aperturedb/Images.py

@@ -3,7 +3,10 @@
 """
 
 from __future__ import annotations
-from typing import Any, Dict, Iterable, List, Tuple, Union
+from typing import Any, Dict, Iterable, List, Tuple, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aperturedb.Frames import Frames
 import cv2
 import math
 import numpy as np
@@ -1003,18 +1006,13 @@ def get_properties(self, prop_list: Iterable[str] = []) -> Dict[str, Any]:
         return return_dictionary
 
 
-class Frames(Images):
-    """
-    **The python wrapper of frame images in ApertureDB.**
+# Shim for backward compatibility
+def __getattr__(name: str):
+    if name == "Frames":
+        from aperturedb.Frames import Frames
+        globals()[name] = Frames
+        return Frames
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
-    Frames in ApertureDB are quite similar to images and so
-    are modeled in python as a subclass.
 
 
-    Args:
-        client: The database connector, perhaps as returned by `CommonLibrary.create_connector`
-    """
-    db_object = ObjectType.FRAME
-
-    def __init__(self, client, batch_size=100, response=None, **kwargs):
-        super().__init__(client, batch_size=batch_size, response=response, **kwargs)

aperturedb/cli/ingest.py

@@ -194,6 +194,7 @@ def from_csv(filepath: Annotated[str, typer.Argument(
     Ingest data from a pre generated CSV file.
     """
     from aperturedb.ImageDataCSV import ImageDataCSV
+    from aperturedb.FrameDataCSV import FrameDataCSV
     from aperturedb.BBoxDataCSV import BBoxDataCSV
     from aperturedb.EntityDataCSV import EntityDataCSV
     from aperturedb.BlobDataCSV import BlobDataCSV
@@ -210,6 +211,7 @@ def from_csv(filepath: Annotated[str, typer.Argument(
         IngestType.DESCRIPTOR: DescriptorDataCSV,
         IngestType.DESCRIPTORSET: DescriptorSetDataCSV,
         IngestType.ENTITY: EntityDataCSV,
+        IngestType.FRAME: FrameDataCSV,
         IngestType.IMAGE: ImageDataCSV,
         IngestType.POLYGON: PolygonDataCSV,
         IngestType.VIDEO: VideoDataCSV

test/test_DataModels.py

@@ -0,0 +1,15 @@
+from aperturedb.DataModels import FrameDataModel
+from aperturedb.Query import ObjectType
+from pydantic import ValidationError
+import pytest
+
+def test_FrameDataModel():
+    # url is optional to avoid breaking existing users
+    frame_no_url = FrameDataModel()
+    assert frame_no_url.url is None
+
+    # Verify that when url is provided, the model instantiates correctly
+    frame = FrameDataModel(url="http://example.com/frame.jpg")
+    assert frame.url == "http://example.com/frame.jpg"
+    assert frame.type == ObjectType.FRAME
+    assert frame.id is not None  # Should have a default UUID generated

test/test_Entities.py

@@ -0,0 +1,9 @@
+from aperturedb.Entities import load_entities_registry
+from aperturedb.Query import ObjectType
+from aperturedb.Frames import Frames
+import pytest
+
+def test_load_entities_registry_frames():
+    registry = load_entities_registry()
+    assert ObjectType.FRAME.value in registry
+    assert registry[ObjectType.FRAME.value] is Frames

test/test_FrameDataCSV.py

@@ -0,0 +1,26 @@
+import tempfile
+import os
+from aperturedb.FrameDataCSV import FrameDataCSV
+from aperturedb.Query import ObjectType
+
+
+def test_FrameDataCSV_command():
+    with tempfile.NamedTemporaryFile(
+        suffix=".csv", mode="w", delete=False
+    ) as f:
+        f.write("url,id\nhttp://example.com/frame.jpg,1\n")
+
+    try:
+        # We don't actually need the image since check_image=False
+        frame_data = FrameDataCSV(f.name, check_image=False)
+
+        cmd = frame_data.command
+        assert cmd == "AddFrame", f"Expected AddFrame, got {cmd}"
+
+        indices = frame_data.get_indices()
+        assert "entity" in indices
+        frame_type = ObjectType.FRAME.value
+        assert frame_type in indices["entity"]
+        assert indices["entity"][frame_type] == frame_data.get_indexed_properties()
+    finally:
+        os.remove(f.name)

test/test_Frames.py

@@ -0,0 +1,26 @@
+from aperturedb.Frames import Frames
+from aperturedb.Query import ObjectType
+
+
+class MockClient:
+    def __init__(self):
+        pass
+
+
+def test_Frames_init():
+    client = MockClient()
+    frame = Frames(client)
+    assert frame.client == client
+    assert frame.db_object == ObjectType.FRAME
+
+
+def test_Frames_backward_compatibility_import():
+    # Verify that importing Frames from aperturedb.Images resolves correctly
+    import aperturedb.Images
+    from aperturedb.Images import Frames as ImagesFrames
+
+    assert ImagesFrames is Frames
+
+    # Verify it is cached in the module's globals
+    assert "Frames" in aperturedb.Images.__dict__
+    assert aperturedb.Images.__dict__["Frames"] is Frames

test/test_cli_ingest.py

@@ -0,0 +1,19 @@
+from unittest.mock import patch, MagicMock
+from aperturedb.cli.ingest import from_csv, IngestType
+
+
+def test_from_csv_frame_type():
+    import aperturedb.FrameDataCSV
+
+    with patch("aperturedb.cli.ingest._process_data") as mock_process_data:
+        mock_process_data.return_value = None
+
+        with patch.object(aperturedb.FrameDataCSV, "FrameDataCSV") as mock_csv_class:
+            mock_data = MagicMock()
+            mock_csv_class.return_value = mock_data
+
+            from_csv(filepath="dummy.csv",
+                     ingest_type=IngestType.FRAME, sample_count=5)
+
+            mock_csv_class.assert_called_once_with(
+                "dummy.csv", use_dask=False, blobs_relative_to_csv=True)