New tool to export information about every image in a VRS file [POC]

Summary:
This is  a POC API + experimentation/debug CLI tool.
The tool demonstrate that we can extract the information necessary to read every image in a VRS file without going through a RecordFileReader. Use the API and tool in the next diff to use the data to extract frames.

To be clear: the tool is just for experimentation, and maybe debugging purposes. For production use, we expect to use only the API, in C++ or through a Python bindings TBD.

Differential Revision: D64584016

fbshipit-source-id: da200d7414e2d9f2be842cc91e157020912b45d3

vrs/utils/ImageIndexer.cpp

@@ -0,0 +1,126 @@
+/*
+ * 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 <memory>
+#include <ostream>
+#include <vector>
+
+#define DEFAULT_LOG_CHANNEL "ImageIndexer"
+#include <logging/Log.h>
+#include <logging/Verify.h>
+
+#include <vrs/FileFormat.h>
+#include <vrs/os/Utils.h>
+#include <vrs/utils/FilteredFileReader.h>
+#include <vrs/utils/PixelFrame.h>
+#include <vrs/utils/VideoRecordFormatStreamPlayer.h>
+
+#include "ImageIndexer.h"
+
+using namespace std;
+using namespace vrs;
+using namespace vrs::utils;
+
+namespace {
+
+class ImageOffsetWriter : public utils::VideoRecordFormatStreamPlayer {
+ public:
+  explicit ImageOffsetWriter(vector<DirectImageReference>& images, int& counter, int& compressed)
+      : images_{images}, counter_{counter}, compressed_{compressed} {}
+
+  bool processRecordHeader(const CurrentRecord& record, DataReference& outDataReference) override {
+    recordStartOffset_ = record.reader->getFileOffset();
+    recordDiskSize_ = record.reader->getUnreadDiskBytes();
+    return RecordFormatStreamPlayer::processRecordHeader(record, outDataReference);
+  }
+
+  bool onImageRead(const CurrentRecord& r, size_t /* id x*/, const ContentBlock& cb) override {
+    if (r.reader->getCompressionType() == CompressionType::None) {
+      images_.emplace_back(r.reader->getFileOffset(), cb.getBlockSize(), cb.image().asString());
+      counter_++;
+    } else {
+      images_.emplace_back(
+          recordStartOffset_,
+          recordDiskSize_,
+          cb.image().asString(),
+          r.reader->getCompressionType(),
+          r.recordSize - r.reader->getUnreadBytes(),
+          cb.getBlockSize());
+      counter_++, compressed_++;
+    }
+    utils::PixelFrame frame;
+    XR_VERIFY(readFrame(frame, r, cb));
+    return true;
+  }
+
+ private:
+  vector<DirectImageReference>& images_;
+  int& counter_;
+  int& compressed_;
+  int64_t recordStartOffset_{-1};
+  uint32_t recordDiskSize_{0};
+};
+
+} // namespace
+
+namespace vrs::utils {
+
+int indexImages(FilteredFileReader& reader, vector<DirectImageReference>& outImages) {
+  outImages.clear();
+
+  int uncompressed = 0;
+  int compressed = 0;
+
+  vector<unique_ptr<StreamPlayer>> streamPlayers;
+  bool hasImages = false;
+  size_t maxCounter = 0;
+  map<StreamId, size_t> imageOffsets;
+  for (auto id : reader.filter.streams) {
+    if (reader.reader.mightContainImages(id)) {
+      XR_LOGI("Found {} - {}...", id.getNumericName(), id.getTypeName());
+      auto player = make_unique<ImageOffsetWriter>(outImages, uncompressed, compressed);
+      reader.reader.setStreamPlayer(id, player.get());
+      streamPlayers.emplace_back(std::move(player));
+      maxCounter += reader.reader.getRecordCount(id, Record::Type::DATA);
+      hasImages = true;
+    }
+  }
+  if (!hasImages) {
+    XR_LOGW("No image stream found in the file");
+    return 0;
+  }
+  outImages.reserve(maxCounter);
+  reader.iterateSafe();
+
+  if (compressed == 0) {
+    XR_LOGI("Found {} frames, none compressed!", uncompressed);
+  } else {
+    XR_LOGI("Found {} frames, {} compressed!", uncompressed, compressed);
+  }
+
+  return 0;
+}
+
+int indexImages(const string& path, vector<DirectImageReference>& outImages) {
+  FilteredFileReader reader;
+  int status = reader.setSource(path);
+  if (status != 0 || (status = reader.openFile()) != 0) {
+    return status;
+  }
+  return indexImages(reader, outImages);
+}
+
+} // namespace vrs::utils

vrs/utils/ImageIndexer.h

@@ -0,0 +1,78 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <vrs/utils/FilteredFileReader.h>
+
+namespace vrs::utils {
+
+/// Helper class to reference images in a VRS file.
+struct DirectImageReference {
+  DirectImageReference(int64_t dataOffset, uint32_t dataSize, string imageFormat)
+      : dataOffset{dataOffset}, dataSize{dataSize}, imageFormat{std::move(imageFormat)} {}
+
+  DirectImageReference(
+      int64_t dataOffset,
+      uint32_t dataSize,
+      string imageFormat,
+      CompressionType compressionType,
+      uint32_t compressedOffset,
+      uint32_t compressedLength)
+      : dataOffset{dataOffset},
+        dataSize{dataSize},
+        imageFormat{std::move(imageFormat)},
+        compressionType{compressionType},
+        compressedOffset{compressedOffset},
+        compressedLength{compressedLength} {}
+
+  void setCompression(
+      CompressionType _compressionType,
+      uint32_t _compressedOffset,
+      uint32_t _compressedLength) {
+    this->compressionType = _compressionType;
+    this->compressedOffset = _compressedOffset;
+    this->compressedLength = _compressedLength;
+  }
+
+  bool operator==(const DirectImageReference& other) const {
+    return dataOffset == other.dataOffset && dataSize == other.dataSize &&
+        imageFormat == other.imageFormat && compressionType == other.compressionType &&
+        compressedOffset == other.compressedOffset && compressedLength == other.compressedLength;
+  }
+
+  int64_t dataOffset{};
+  uint32_t dataSize{};
+  string imageFormat;
+  CompressionType compressionType{CompressionType::None};
+  uint32_t compressedOffset{};
+  uint32_t compressedLength{};
+};
+
+/// Get the list of references for all the images found in a VRS file.
+/// @param path: path to the VRS file to index.
+/// @param outImages: on exit, a list of image references.
+/// @return 0 on success, or an error code.
+int indexImages(const string& path, vector<DirectImageReference>& outImages);
+
+/// Get a list of references for the images found in a VRS file, using a FilteredFileReader
+/// that may restrict the streams and time range considered. Useful for a CLI tool.
+/// @param reader: an open FilteredFileReader, with or without filters.
+/// @param outImages: on exit, a list of image references.
+/// @return 0 on success, or an error code.
+int indexImages(FilteredFileReader& reader, vector<DirectImageReference>& outImages);
+
+} // namespace vrs::utils

vrs/utils/test/ImageIndexerLoaderTest.cpp

@@ -0,0 +1,50 @@
+/*
+ * 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 <gtest/gtest.h>
+
+#include <TestDataDir/TestDataDir.h>
+
+#include <vrs/os/Utils.h>
+#include <vrs/utils/ImageIndexer.h>
+
+using namespace std;
+using namespace vrs;
+using namespace vrs::utils;
+
+struct ImageIndexerLoaderTest : testing::Test {
+  string kRgbFile = os::pathJoin(coretech::getTestDataDir(), "VRS_Files/rgb8.vrs");
+  string kJpgFile = os::pathJoin(coretech::getTestDataDir(), "VRS_Files/jpg.vrs");
+};
+
+TEST_F(ImageIndexerLoaderTest, ImageIndexerLoaderTest) {
+  vector<DirectImageReference> readImages;
+
+  ASSERT_EQ(indexImages(kRgbFile, readImages), 0);
+  const string format = "raw/1224x1024/pixel=rgb8/stride=3672";
+  vector<DirectImageReference> expectedImages = {
+      {2251, 2105916, format, CompressionType::Zstd, 52, 3760128},
+      {2108199, 2106944, format, CompressionType::Zstd, 52, 3760128},
+      {4215175, 2106022, format, CompressionType::Zstd, 52, 3760128},
+  };
+  EXPECT_EQ(readImages, expectedImages);
+
+  ASSERT_EQ(indexImages(kJpgFile, readImages), 0);
+  expectedImages = {
+      {6046, 1985655, "jpg", CompressionType::None, 0, 0},
+  };
+  EXPECT_EQ(readImages, expectedImages);
+}