Add DataLayout support for pyvrs VRSWriter (#356)

Summary:
Pull Request resolved: #356

Pull Request resolved: #355

* **rename of zstd**: needed because OSS version only expose ZSTD_MEDIUM via pybind. internal version still works b.c. it exposes both ZSTD_MEDIUM and Zmedium.

Differential Revision: D95732614

Author: danielyan86129

.github/workflows/deploy.yml

@@ -54,14 +54,14 @@ 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: Install cibuildwheel
         run: |
-          python3 -m pip install cibuildwheel==2.17.0
+          python3 -m pip install cibuildwheel==2.22.0
 
       - name: Build wheels for CPython
         run: |

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)