Add stream_path to avm_frame.proto
diff --git a/tools/extract_proto/avm_frame.proto b/tools/extract_proto/avm_frame.proto
index abb1852..a448d44 100644
--- a/tools/extract_proto/avm_frame.proto
+++ b/tools/extract_proto/avm_frame.proto
@@ -238,13 +238,14 @@
string stream_hash = 1;
string stream_name = 2;
float frame_rate = 3;
- // Note: these are present both here and in FrameParms. For most streams we
+ // Note: these are present both here and in FrameParams. For most streams we
// care about, the frame dimensions will be the same across every frame. For
// streams with variable-sized frames, these fields can be omitted.
int32 width = 4;
int32 height = 5;
string avm_version = 6;
map<string, string> encoder_args = 7;
+ string stream_path = 8;
}
message TipFrameParams {
diff --git a/tools/extract_proto/extract_proto.cc b/tools/extract_proto/extract_proto.cc
index a06982f..2dffd87 100644
--- a/tools/extract_proto/extract_proto.cc
+++ b/tools/extract_proto/extract_proto.cc
@@ -93,6 +93,10 @@
"Print progress as each frame is decoded.");
ABSL_FLAG(bool, ignore_y4m_header, false,
"Ignore dimensions and colorspace in Y4M header.");
+ABSL_FLAG(int, preserve_stream_path_depth, 0,
+ "Preserve this many levels of directory hierarchy of the stream's "
+ "path. -1 will preserve the entire absolute path. 0 will keep only "
+ "the filename. 1 will keep just the direct parent, etc...");
namespace {
constexpr std::string_view kY4mFrameMarker = "FRAME\n";
@@ -130,6 +134,7 @@
std::ifstream *orig_yuv_file;
int orig_yuv_bit_depth;
std::filesystem::path stream_path;
+ int preserve_stream_path_depth;
int decode_count;
bool is_y4m_file;
std::optional<Y4mHeader> y4m_header;
@@ -894,6 +899,23 @@
aom_codec_control(&ctx->codec, AV1_SET_INSPECTION_CALLBACK, &ii);
}
+void PopulateStreamPath(ExtractProtoContext *ctx) {
+ auto absolute_path = std::filesystem::absolute(ctx->stream_path);
+ std::filesystem::path relative_path;
+ if (ctx->preserve_stream_path_depth == -1) {
+ relative_path = absolute_path;
+ } else {
+ int index = std::distance(absolute_path.begin(), absolute_path.end());
+ for (const auto &part : absolute_path) {
+ index -= 1;
+ if (index <= ctx->preserve_stream_path_depth) {
+ relative_path /= part;
+ }
+ }
+ }
+ ctx->stream_params->set_stream_path(relative_path.make_preferred());
+}
+
absl::Status OpenStream(ExtractProtoContext *ctx) {
ctx->reader = aom_video_reader_open(ctx->stream_path.c_str());
if (ctx->reader == nullptr) {
@@ -909,6 +931,7 @@
LOG(INFO) << "Using " << aom_codec_iface_name(decoder);
ctx->stream_params->set_avm_version(aom_codec_iface_name(decoder));
ctx->stream_params->set_stream_name(ctx->stream_path.filename());
+ PopulateStreamPath(ctx);
if (aom_codec_dec_init(&ctx->codec, decoder, nullptr, 0)) {
return absl::InternalError("Failed to initialize decoder.");
}
@@ -1034,6 +1057,8 @@
.orig_yuv_file = have_orig_yuv ? &orig_yuv_file : nullptr,
.orig_yuv_bit_depth = orig_yuv_bit_depth,
.stream_path = stream_path,
+ .preserve_stream_path_depth =
+ absl::GetFlag(FLAGS_preserve_stream_path_depth),
.decode_count = 0,
.is_y4m_file = is_y4m_file,
.y4m_header = y4m_header,