Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expose utility for logging video frames for an entire video #7421

Merged
merged 22 commits into from
Sep 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -4946,7 +4946,6 @@ dependencies = [
"re_smart_channel",
"re_tracing",
"re_types",
"re_video",
"thiserror",
"walkdir",
]
Expand Down Expand Up @@ -5798,8 +5797,10 @@ dependencies = [
name = "re_video"
version = "0.19.0-alpha.1+dev"
dependencies = [
"itertools 0.13.0",
"mp4",
"ordered-float",
"thiserror",
]

[[package]]
Expand Down Expand Up @@ -6155,6 +6156,7 @@ dependencies = [
"re_arrow2",
"re_log",
"re_sdk",
"re_video",
]

[[package]]
Expand All @@ -6179,6 +6181,7 @@ dependencies = [
"re_log_types",
"re_memory",
"re_sdk",
"re_video",
"re_web_viewer_server",
"re_ws_comms",
"uuid",
Expand Down
2 changes: 1 addition & 1 deletion crates/store/README.md
Original file line number Diff line number Diff line change
@@ -1 +1 @@
Creates related to storing, indexing, trasmitting, and handling data.
Creates related to storing, indexing, transmitting, and handling data.
3 changes: 1 addition & 2 deletions crates/store/re_data_loader/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,7 @@ re_log_types.workspace = true
re_log.workspace = true
re_smart_channel.workspace = true
re_tracing.workspace = true
re_types = { workspace = true, features = ["image"] }
re_video.workspace = true
re_types = { workspace = true, features = ["image", "video"] }

ahash.workspace = true
anyhow.workspace = true
Expand Down
143 changes: 51 additions & 92 deletions crates/store/re_data_loader/src/loader_archetype.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
use re_chunk::{Chunk, RowId};
use re_log_types::{EntityPath, TimeInt, TimePoint};
use re_types::archetypes::VideoFrameReference;
use re_types::archetypes::{AssetVideo, VideoFrameReference};
use re_types::components::VideoTimestamp;
use re_types::Archetype;
use re_types::{components::MediaType, ComponentBatch};

use arrow2::array::{
ListArray as ArrowListArray, NullArray as ArrowNullArray, PrimitiveArray as ArrowPrimitiveArray,
};
use arrow2::array::PrimitiveArray as ArrowPrimitiveArray;
use arrow2::Either;

use crate::{DataLoader, DataLoaderError, LoadedData};
Expand Down Expand Up @@ -220,100 +219,60 @@ fn load_video(
let video_timeline = re_log_types::Timeline::new_temporal("video");
timepoint.insert(video_timeline, re_log_types::TimeInt::new_temporal(0));

let media_type = MediaType::guess_from_path(filepath);

// TODO(andreas): Video frame reference generation should be available as a utility from the SDK.

let video = if media_type.as_ref().map(|v| v.as_str()) == Some("video/mp4") {
match re_video::load_mp4(&contents) {
Ok(video) => Some(video),
Err(err) => {
re_log::warn!("Failed to load video asset {filepath:?}: {err}");
None
}
let video_asset = AssetVideo::new(contents);

let video_frame_reference_chunk = match video_asset.read_frame_timestamps_ns() {
Ok(frame_timestamps_ns) => {
// Time column.
let is_sorted = Some(true);
let time_column_times = ArrowPrimitiveArray::from_slice(&frame_timestamps_ns);
let time_column =
re_chunk::TimeColumn::new(is_sorted, video_timeline, time_column_times);

// VideoTimestamp component column.
let video_timestamps = frame_timestamps_ns
.into_iter()
.map(VideoTimestamp::from_nanoseconds)
.collect::<Vec<_>>();
let video_timestamp_batch = &video_timestamps as &dyn ComponentBatch;
let video_timestamp_list_array = video_timestamp_batch
.to_arrow_list_array()
.map_err(re_chunk::ChunkError::from)?;

// Indicator column.
let video_frame_reference_indicators =
<VideoFrameReference as Archetype>::Indicator::new_array(video_timestamps.len());
let video_frame_reference_indicators_list_array = video_frame_reference_indicators
.to_arrow_list_array()
.map_err(re_chunk::ChunkError::from)?;

Some(Chunk::from_auto_row_ids(
re_chunk::ChunkId::new(),
entity_path.clone(),
std::iter::once((video_timeline, time_column)).collect(),
[
(
VideoFrameReference::indicator().name(),
video_frame_reference_indicators_list_array,
),
(video_timestamp_batch.name(), video_timestamp_list_array),
]
.into_iter()
.collect(),
)?)
}
} else {
re_log::warn!("Video asset {filepath:?} has an unsupported container format.");
None
};

// Log video frame references on the `video` timeline.
let video_frame_reference_chunk = if let Some(video) = video {
let first_timestamp = video
.segments
.first()
.map_or(0, |segment| segment.timestamp.as_nanoseconds());

// Time column.
let is_sorted = Some(true);
let time_column_times =
ArrowPrimitiveArray::<i64>::from_values(video.segments.iter().flat_map(|segment| {
segment
.samples
.iter()
.map(|s| s.timestamp.as_nanoseconds() - first_timestamp)
}));

let time_column = re_chunk::TimeColumn::new(is_sorted, video_timeline, time_column_times);

// VideoTimestamp component column.
let video_timestamps = video
.segments
.iter()
.flat_map(|segment| {
segment.samples.iter().map(|s| {
// TODO(andreas): Use sample indices instead of timestamps once possible.
re_types::components::VideoTimestamp::from_nanoseconds(
s.timestamp.as_nanoseconds(),
)
})
})
.collect::<Vec<_>>();
let video_timestamp_batch = &video_timestamps as &dyn ComponentBatch;
let video_timestamp_list_array = video_timestamp_batch
.to_arrow_list_array()
.map_err(re_chunk::ChunkError::from)?;

// Indicator column.
let video_frame_reference_indicator_datatype = arrow2::datatypes::DataType::Null;
let video_frame_reference_indicator_list_array = ArrowListArray::<i32>::try_new(
ArrowListArray::<i32>::default_datatype(
video_frame_reference_indicator_datatype.clone(),
),
video_timestamp_list_array.offsets().clone(),
Box::new(ArrowNullArray::new(
video_frame_reference_indicator_datatype,
video_timestamps.len(),
)),
None,
)
.map_err(re_chunk::ChunkError::from)?;

Some(Chunk::from_auto_row_ids(
re_chunk::ChunkId::new(),
entity_path.clone(),
std::iter::once((video_timeline, time_column)).collect(),
[
(
VideoFrameReference::indicator().name(),
video_frame_reference_indicator_list_array,
),
(video_timestamp_batch.name(), video_timestamp_list_array),
]
.into_iter()
.collect(),
)?)
} else {
None
Err(err) => {
re_log::warn_once!(
"Failed to read frame timestamps from video asset {filepath:?}: {err}"
);
None
}
};

// Put video asset into its own chunk since it can be fairly large.
let video_asset_chunk = Chunk::builder(entity_path.clone())
.with_archetype(
RowId::new(),
timepoint.clone(),
&re_types::archetypes::AssetVideo::from_file_contents(contents, media_type.clone()),
)
.with_archetype(RowId::new(), timepoint.clone(), &video_asset)
.with_component_batch(RowId::new(), timepoint.clone(), &ExperimentalFeature)
.build()?;

Expand Down
4 changes: 2 additions & 2 deletions crates/store/re_types/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ features = ["all"]
default = ["ecolor"]

## All features except `testing`.
all = ["ecolor", "egui_plot", "glam", "image", "mint", "serde"]
all = ["ecolor", "egui_plot", "glam", "image", "mint", "serde", "video"]

## Enable color conversions.
ecolor = ["dep:ecolor"]
Expand All @@ -39,7 +39,7 @@ glam = ["dep:glam"]
## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG support.
image = ["dep:ecolor", "dep:image"]

## Conversion to/from our video format
## Inspecting video data.
video = ["dep:re_video"]

## Enable (de)serialization using serde.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@ namespace rerun.archetypes;
///
/// In order to display a video, you need to log a [archetypes.VideoFrameReference] for each frame.
///
/// \example archetypes/video_manual_frames title="Video with explicit frames" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
// TODO(#7368): Example and reference to `send_video_frames` API.
/// \example archetypes/video_auto_frames title="Video with automatically determined frames" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
/// \example archetypes/video_manual_frames title="Demonstrates manual use of video frame references" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
// TODO(#7420): update screenshot for manual frames example
table AssetVideo (
"attr.docs.unreleased",
"attr.rerun.experimental"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@ namespace rerun.archetypes;
/// Used to display individual video frames from a [archetypes.AssetVideo].
/// To show an entire video, a fideo frame reference for each frame of the video should be logged.
///
/// \example archetypes/video_manual_frames title="Video with explicit frames" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
// TODO(#7368): Example and reference to `send_video_frames` API.
/// \example archetypes/video_auto_frames title="Video with automatically determined frames" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
/// \example archetypes/video_manual_frames title="Demonstrates manual use of video frame references" image="https://static.rerun.io/video_manual_frames/320a44e1e06b8b3a3161ecbbeae3e04d1ccb9589/1200w.png"
// TODO(#7420): update screenshot for manual frames example
table VideoFrameReference (
"attr.docs.unreleased",
"attr.rerun.experimental"
Expand Down
84 changes: 65 additions & 19 deletions crates/store/re_types/src/archetypes/asset_video.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

13 changes: 13 additions & 0 deletions crates/store/re_types/src/archetypes/asset_video_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,17 @@ impl AssetVideo {
media_type,
}
}

/// Determines the presentation timestamps of all frames inside the video.
///
/// Returned timestamps are in nanoseconds since start and are guaranteed to be monotonically increasing.
#[cfg(feature = "video")]
pub fn read_frame_timestamps_ns(&self) -> Result<Vec<i64>, re_video::VideoLoadError> {
Ok(re_video::VideoData::load_from_bytes(
self.blob.as_slice(),
self.media_type.as_ref().map(|m| m.as_str()),
)?
.frame_timestamps_ns()
.collect())
}
}
Loading
Loading