Add AriaGen2 image data layout bindings (#359)

.github/workflows/build-and-test.yml

@@ -12,8 +12,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        # UPDATED: Replaced macos-13 with macos-latest
-        os: [ubuntu-latest, macos-latest]
+        os: [ubuntu-latest, macos-14, macos-15]
       fail-fast: false
 
     steps:
@@ -54,6 +53,31 @@ jobs:
           fi
           python -m pip install -U pip
           python -m pip install numpy typing dataclasses pytest parameterized Pillow
+
+      - name: Pre-download virtualenv to avoid GitHub rate limits
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          VIRTUALENV_VERSION=20.21.1
+          if [ "$RUNNER_OS" == "macOS" ]; then
+            CACHE_DIR="$HOME/Library/Caches/cibuildwheel"
+          else
+            CACHE_DIR="$HOME/.cache/cibuildwheel"
+          fi
+          mkdir -p "$CACHE_DIR"
+          DEST="$CACHE_DIR/virtualenv-${VIRTUALENV_VERSION}.pyz"
+          if [ ! -f "$DEST" ]; then
+            # Use GitHub API (authenticated, 5000 req/hr) instead of raw blob URL (unauthenticated, 60 req/hr)
+            curl -sL --retry 5 --retry-delay 10 --retry-all-errors \
+              -H "Authorization: token ${GH_TOKEN}" \
+              -H "Accept: application/vnd.github.v3.raw" \
+              -o "$DEST" \
+              "https://api.github.com/repos/pypa/get-virtualenv/contents/public/virtualenv.pyz?ref=${VIRTUALENV_VERSION}"
+          fi
+          ls -la "$DEST"
+
+      - name: Install cibuildwheel
+        run: |
           python -m pip install cibuildwheel==2.17.0
 
       - name: Build wheels for CPython
@@ -68,6 +92,8 @@ jobs:
           CIBW_BEFORE_BUILD_LINUX: bash scripts/install-manylinux-deps.sh
           CIBW_BEFORE_BUILD_MACOS: bash scripts/install-macos-deps.sh
           CIBW_SKIP: "*-manylinux_i686 *musllinux*"
+          # Target macOS 14 (Sonoma) so wheels work on 14+.
+          CIBW_ENVIRONMENT_MACOS: MACOSX_DEPLOYMENT_TARGET=14.0
 
       - name: Build and Test Python
         shell: bash

.github/workflows/deploy.yml

@@ -35,9 +35,9 @@ jobs:
           python3 -m pip install setuptools wheel twine
           python3 setup.py sdist
 
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v4
         with:
-          name: python-package-distributions
+          name: python-package-distributions-sdist
           path: dist/*.tar.gz
 
   build-wheels:
@@ -54,11 +54,33 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.9'
+          python-version: '3.11'
       - name: Upgrade pip
         run: |
           python3 -m pip install --upgrade pip
 
+      - name: Pre-download virtualenv to avoid GitHub rate limits
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          VIRTUALENV_VERSION=20.21.1
+          if [ "$RUNNER_OS" == "macOS" ]; then
+            CACHE_DIR="$HOME/Library/Caches/cibuildwheel"
+          else
+            CACHE_DIR="$HOME/.cache/cibuildwheel"
+          fi
+          mkdir -p "$CACHE_DIR"
+          DEST="$CACHE_DIR/virtualenv-${VIRTUALENV_VERSION}.pyz"
+          if [ ! -f "$DEST" ]; then
+            # Use GitHub API (authenticated, 5000 req/hr) instead of raw blob URL (unauthenticated, 60 req/hr)
+            curl -sL --retry 5 --retry-delay 10 --retry-all-errors \
+              -H "Authorization: token ${GH_TOKEN}" \
+              -H "Accept: application/vnd.github.v3.raw" \
+              -o "$DEST" \
+              "https://api.github.com/repos/pypa/get-virtualenv/contents/public/virtualenv.pyz?ref=${VIRTUALENV_VERSION}"
+          fi
+          ls -la "$DEST"
+
       - name: Install cibuildwheel
         run: |
           python3 -m pip install cibuildwheel==2.17.0
@@ -73,9 +95,11 @@ jobs:
           CIBW_BEFORE_BUILD_LINUX: bash scripts/install-manylinux-deps.sh
           CIBW_BEFORE_BUILD_MACOS: bash scripts/install-macos-deps.sh
           CIBW_SKIP: "*-manylinux_i686 *musllinux*"
-      - uses: actions/upload-artifact@v2
+          # Target macOS 14 (Sonoma) so wheels work on 14+.
+          CIBW_ENVIRONMENT_MACOS: MACOSX_DEPLOYMENT_TARGET=14.0
+      - uses: actions/upload-artifact@v4
         with:
-          name: python-package-distributions
+          name: python-package-distributions-${{ matrix.os }}
           path: dist
 
   publish-to-pypi:
@@ -84,13 +108,13 @@ jobs:
       - build-wheels
       - create-sdist
     steps:
-    - name: Download wheels from previous jobs
+    - name: Download all artifacts
       # by default this will download all artifacts
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
-        name: python-package-distributions
         # PyPI publish action uploads everything under dist/* by default
         path: dist
+        merge-multiple: true
 
     - name: Display the list of artifacts
       run: ls -R dist

csrc/writer/VRSWriter.cpp

@@ -40,6 +40,7 @@
 #include "StreamFactory.h"
 
 // Open source DataLayout definitions
+#include "datalayouts/AriaGen2ImageDataLayout.h"
 #include "datalayouts/SampleDataLayout.h"
 
 namespace py = pybind11;
@@ -74,6 +75,16 @@ void VRSWriter::init() {
       "sample_with_image", createSampleStreamWithImage);
   StreamFactory::getInstance().registerStreamCreationFunction(
       "sample_with_multiple_data_layout", createSampleStreamWithMultipleDataLayout);
+  // Aria Gen2 camera streams with correct RecordableTypeId + H.265 content block
+  StreamFactory::getInstance().registerFlavoredStreamCreationFunction(
+      "aria_gen2_rgb_camera", [](const string& flavor) {
+        return createAriaGen2ImageStream(
+            flavor, RecordableTypeId::RgbCameraRecordableClass, "H.265");
+      });
+  StreamFactory::getInstance().registerFlavoredStreamCreationFunction(
+      "aria_gen2_slam_camera", [](const string& flavor) {
+        return createAriaGen2ImageStream(flavor, RecordableTypeId::SlamCameraData, "H.265");
+      });
   /// Register open source stream writers (end)
 
 #if IS_VRS_FB_INTERNAL()

csrc/writer/datalayouts/AriaGen2ImageDataLayout.cpp

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "AriaGen2ImageDataLayout.h"
+
+#include <vrs/RecordFormat.h>
+
+#include "../PyRecordable.h"
+
+using namespace vrs;
+
+namespace pyvrs {
+
+std::unique_ptr<PyStream> createAriaGen2ImageStream(
+    const std::string& flavor,
+    RecordableTypeId typeId,
+    const std::string& codec) {
+  constexpr uint32_t kVersion = 2;
+
+  auto configurationRecordFormat = std::make_unique<PyRecordFormat>(
+      Record::Type::CONFIGURATION,
+      kVersion,
+      std::make_unique<ImageSensorConfigurationLayout>(/*allocateVideoFields=*/true));
+
+  auto dataContentBlocks = codec == "H.265"
+      ? std::vector<ContentBlock>{ContentBlock("H.265", ImageContentBlockSpec::kQualityUndefined)}
+      : std::vector<ContentBlock>{ContentBlock(ImageFormat::RAW)};
+
+  auto dataRecordFormat = std::make_unique<PyRecordFormat>(
+      Record::Type::DATA,
+      kVersion,
+      std::make_unique<ImageDataLayout>(/*allocateVideoFields=*/true),
+      dataContentBlocks);
+
+  return std::make_unique<PyStream>(
+      typeId, flavor, std::move(configurationRecordFormat), std::move(dataRecordFormat));
+}
+
+} // namespace pyvrs

csrc/writer/datalayouts/AriaGen2ImageDataLayout.h

@@ -0,0 +1,135 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Aria Gen2 camera DataLayouts for OSS pyvrs writer.
+// Vendored from arvr/libraries/visiontypes/vrs/data_layouts/ImageDataLayout.h
+// with namespace changed from visiontypes::detail to pyvrs.
+
+#pragma once
+
+#include <cstdint>
+
+#include <vrs/DataLayout.h>
+#include <vrs/DataLayoutConventions.h>
+#include <vrs/DataPieces.h>
+
+namespace pyvrs {
+
+using vrs::OptionalDataPieces;
+using vrs::datalayout_conventions::ImageSpecType;
+
+// Additional fields to enable in ImageSensorConfigurationLayout when data was
+// encoded as video before recording.
+struct VideoConfigurationFields {
+  vrs::DataPieceString videoCodecName{vrs::datalayout_conventions::kImageCodecName};
+};
+
+struct ImageSensorConfigurationLayout : public vrs::AutoDataLayout {
+  static constexpr uint32_t kVersion = 2;
+
+  explicit ImageSensorConfigurationLayout(bool allocateVideoFields = false)
+      : videoConfigurationFields(allocateVideoFields) {}
+
+  vrs::DataPieceString deviceType{"device_type"};
+  vrs::DataPieceString deviceVersion{"device_version"};
+  vrs::DataPieceString deviceSerial{"device_serial"};
+
+  vrs::DataPieceValue<std::uint32_t> cameraId{"camera_id"};
+  vrs::DataPieceValue<std::uint32_t> streamType{"stream_type"};
+  vrs::DataPieceValue<std::uint32_t> streamIndex{"stream_index"};
+
+  vrs::DataPieceString sensorModel{"sensor_model"};
+  vrs::DataPieceString sensorSerial{"sensor_serial"};
+
+  vrs::DataPieceValue<double> nominalRateHz{"nominal_rate"};
+
+  vrs::DataPieceValue<ImageSpecType> imageWidth{vrs::datalayout_conventions::kImageWidth};
+  vrs::DataPieceValue<ImageSpecType> imageHeight{vrs::datalayout_conventions::kImageHeight};
+  vrs::DataPieceValue<ImageSpecType> imageStride{vrs::datalayout_conventions::kImageStride};
+  vrs::DataPieceValue<ImageSpecType> imageStride2{vrs::datalayout_conventions::kImageStride2};
+  vrs::DataPieceValue<ImageSpecType> pixelFormat{vrs::datalayout_conventions::kImagePixelFormat};
+  vrs::DataPieceValue<ImageSpecType> plane2OffsetRows{"image_plane_2_offset_rows"};
+  vrs::DataPieceValue<ImageSpecType> plane3OffsetRows{"image_plane_3_offset_rows"};
+
+  vrs::DataPieceValue<std::uint32_t> imageOrientation{"image_orientation"};
+  vrs::DataPieceValue<std::uint32_t> shutterDirection{"shutter_direction"};
+
+  vrs::DataPieceValue<double> exposureDurationMin{"exposure_duration.min"};
+  vrs::DataPieceValue<double> exposureDurationMax{"exposure_duration.max"};
+
+  vrs::DataPieceValue<double> gainMin{"gain.min"};
+  vrs::DataPieceValue<double> gainMax{"gain.max"};
+
+  vrs::DataPieceValue<double> gammaFactor{"gamma_factor"};
+
+  vrs::DataPieceString factoryCalibration{"factory_calibration"};
+  vrs::DataPieceString onlineCalibration{"online_calibration"};
+
+  vrs::DataPieceString description{"description"};
+
+  vrs::DataPieceString cameraMuxModeName{"camera_mux_mode_name"};
+
+  const OptionalDataPieces<VideoConfigurationFields> videoConfigurationFields;
+
+  vrs::AutoDataLayoutEnd end;
+};
+
+// Additional fields to enable in ImageDataLayout when data was encoded as video
+// before recording.
+struct VideoDataFields {
+  vrs::DataPieceValue<double> keyFrameTimestamp{
+      vrs::datalayout_conventions::kImageKeyFrameTimeStamp};
+  vrs::DataPieceValue<ImageSpecType> keyFrameIndex{
+      vrs::datalayout_conventions::kImageKeyFrameIndex};
+};
+
+struct ImageDataLayout : public vrs::AutoDataLayout {
+  static constexpr uint32_t kVersion = 2;
+
+  explicit ImageDataLayout(bool allocateVideoFields = false)
+      : videoDataFields(allocateVideoFields) {}
+
+  vrs::DataPieceValue<std::uint64_t> groupId{"group_id"};
+  vrs::DataPieceValue<std::uint64_t> groupMask{"group_mask"};
+  vrs::DataPieceValue<std::uint64_t> streamIndexMask{"stream_index_mask"};
+  vrs::DataPieceValue<std::uint64_t> frameNumber{"frame_number"};
+  vrs::DataPieceValue<std::uint32_t> frameTag{"frame_tag"};
+  vrs::DataPieceValue<double> exposureDuration{"exposure_duration_s"};
+  vrs::DataPieceValue<double> gain{"gain"};
+  vrs::DataPieceValue<double> readoutDurationSeconds{"readout_duration_s"};
+  vrs::DataPieceValue<std::int64_t> captureTimestampNs{"capture_timestamp_ns"};
+  vrs::DataPieceValue<std::int64_t> captureTimestampInProcessingClockDomainNs{
+      "capture_timestamp_in_processing_clock_domain_ns"};
+  vrs::DataPieceValue<std::int64_t> arrivalTimestampNs{"arrival_timestamp_ns"};
+  vrs::DataPieceValue<std::int64_t> processingStartTimestampNs{"processing_start_timestamp_ns"};
+  vrs::DataPieceValue<double> temperature{"temperature_deg_c"};
+  vrs::DataPieceVector<uint8_t> imageMetadata{"image_metadata"};
+
+  const OptionalDataPieces<VideoDataFields> videoDataFields;
+
+  vrs::DataPieceValue<double> focusDistanceMm{"focus_distance_mm", -1.0};
+
+  vrs::AutoDataLayoutEnd end;
+};
+
+class PyStream;
+
+std::unique_ptr<PyStream> createAriaGen2ImageStream(
+    const std::string& flavor,
+    vrs::RecordableTypeId typeId,
+    const std::string& codec = "H.265");
+
+} // namespace pyvrs

pyvrs/__init__.py

@@ -35,12 +35,15 @@
     recordable_type_id_name,
     RecordableId,
     RecordableTypeId,
+    RecordFormat,
     records_checksum,
     RecordType,
+    Stream,
     StreamNotFoundError,
     TimestampNotFoundError,
     verbatim_checksum,
     VRSRecord,
+    Writer,
 )
 
 from .reader import AsyncVRSReader, SyncVRSReader
@@ -69,10 +72,13 @@
     "recordable_type_id_name",
     "RecordableId",
     "RecordableTypeId",
+    "RecordFormat",
     "records_checksum",
     "RecordType",
+    "Stream",
     "StreamNotFoundError",
     "TimestampNotFoundError",
     "verbatim_checksum",
     "VRSRecord",
+    "Writer",
 ]

pyvrs/writer.py

@@ -57,7 +57,7 @@ def create_stream(
         self,
         name: str,
         flavor: str = "",
-        compression: CompressionPreset = CompressionPreset.Zmedium,
+        compression: CompressionPreset = CompressionPreset.ZSTD_MEDIUM,
     ) -> "VRSStream":
         if len(flavor) > 0:
             return VRSStream(
@@ -121,7 +121,7 @@ def __init__(
         self,
         stream: Stream,
         writer: VRSWriter,
-        compression: CompressionPreset = CompressionPreset.Zmedium,
+        compression: CompressionPreset = CompressionPreset.ZSTD_MEDIUM,
     ) -> None:
         self.stream = stream
         self.stream.setCompression(compression)