use std::collections::BTreeSet;
use std::sync::{Arc, LazyLock};
use re_chunk::{Chunk, ChunkResult};
use re_log_types::{ArrowMsg, EntityPath, LogMsg, RecordingId, StoreId, TimePoint};
mod load_file;
mod loader_archetype;
mod loader_directory;
mod loader_rrd;
mod loader_urdf;
#[cfg(not(target_arch = "wasm32"))]
pub mod lerobot;
#[cfg(not(target_arch = "wasm32"))]
pub mod loader_lerobot;
pub mod loader_mcap;
#[cfg(not(target_arch = "wasm32"))]
mod loader_external;
pub use self::load_file::load_from_file_contents;
pub use self::loader_archetype::ArchetypeLoader;
pub use self::loader_directory::DirectoryLoader;
pub use self::loader_mcap::McapLoader;
pub use self::loader_rrd::RrdLoader;
pub use self::loader_urdf::{UrdfDataLoader, UrdfTree, joint_transform as urdf_joint_transform};
#[cfg(not(target_arch = "wasm32"))]
pub use self::{
load_file::load_from_path,
loader_external::{
EXTERNAL_DATA_LOADER_INCOMPATIBLE_EXIT_CODE, EXTERNAL_DATA_LOADER_PREFIX, ExternalLoader,
iter_external_loaders,
},
loader_lerobot::LeRobotDatasetLoader,
};
pub mod external {
pub use urdf_rs;
}
pub const FOXGLOVE_LENSES_IDENTIFIER: &str = "foxglove";
pub const URDF_DECODER_IDENTIFIER: &str = "urdf";
pub fn supported_mcap_decoder_identifiers(
raw_fallback_enabled: bool,
) -> Vec<re_mcap::DecoderIdentifier> {
let mut identifiers = re_mcap::DecoderRegistry::all_builtin(raw_fallback_enabled)
.all_identifiers()
.into_iter()
.map(re_mcap::DecoderIdentifier::from)
.collect::<BTreeSet<_>>();
identifiers.extend([
re_mcap::DecoderIdentifier::from(FOXGLOVE_LENSES_IDENTIFIER),
re_mcap::DecoderIdentifier::from(URDF_DECODER_IDENTIFIER),
]);
identifiers.into_iter().collect()
}
#[derive(Debug, Clone)]
pub struct DataLoaderSettings {
pub application_id: Option<re_log_types::ApplicationId>,
pub recording_id: RecordingId,
pub opened_store_id: Option<StoreId>,
pub force_store_info: bool,
pub entity_path_prefix: Option<EntityPath>,
pub timepoint: Option<TimePoint>,
pub follow: bool,
pub timestamp_offset_ns: Option<i64>,
pub timeline_type: re_log_types::TimeType,
}
impl DataLoaderSettings {
#[inline]
pub fn recommended(recording_id: impl Into<RecordingId>) -> Self {
Self {
recording_id: recording_id.into(),
application_id: None,
opened_store_id: None,
force_store_info: false,
entity_path_prefix: None,
timepoint: None,
follow: false,
timestamp_offset_ns: None,
timeline_type: re_log_types::TimeType::TimestampNs,
}
}
pub fn recommended_store_id(&self) -> StoreId {
StoreId::recording(
self.application_id
.clone()
.unwrap_or_else(re_log_types::ApplicationId::random),
self.recording_id.clone(),
)
}
pub fn opened_store_id_or_recommended(&self) -> StoreId {
self.opened_store_id
.clone()
.unwrap_or_else(|| self.recommended_store_id())
}
pub fn to_cli_args(&self) -> Vec<String> {
let Self {
application_id,
recording_id,
opened_store_id,
force_store_info: _,
entity_path_prefix,
timepoint,
follow: _,
timestamp_offset_ns: _,
timeline_type: _,
} = self;
let mut args = Vec::new();
if let Some(application_id) = application_id {
args.extend(["--application-id".to_owned(), format!("{application_id}")]);
}
args.extend(["--recording-id".to_owned(), format!("{recording_id}")]);
if let Some(opened_store_id) = opened_store_id {
args.extend([
"--opened-application-id".to_owned(),
format!("{}", opened_store_id.application_id()),
]);
args.extend([
"--opened-recording-id".to_owned(),
format!("{}", opened_store_id.recording_id()),
]);
}
if let Some(entity_path_prefix) = entity_path_prefix {
args.extend([
"--entity-path-prefix".to_owned(),
format!("{entity_path_prefix}"),
]);
}
if let Some(timepoint) = timepoint {
if timepoint.is_static() {
args.push("--timeless".to_owned()); args.push("--static".to_owned());
}
for (timeline, cell) in timepoint.iter() {
match cell.typ() {
re_log_types::TimeType::Sequence => {
args.extend([
"--time_sequence".to_owned(),
format!("{timeline}={}", cell.value),
]);
args.extend([
"--sequence".to_owned(),
format!("{timeline}={}", cell.value),
]);
}
re_log_types::TimeType::DurationNs => {
args.extend([
"--time_duration_nanos".to_owned(),
format!("{timeline}={}", cell.value),
]);
args.extend(["--time".to_owned(), format!("{timeline}={}", cell.value)]);
}
re_log_types::TimeType::TimestampNs => {
args.extend([
"--time_duration_nanos".to_owned(),
format!("{timeline}={}", cell.value),
]);
args.extend([
"--sequence".to_owned(),
format!("{timeline}={}", cell.value),
]);
}
}
}
}
args
}
}
pub type DataLoaderName = String;
pub trait DataLoader: Send + Sync {
fn name(&self) -> DataLoaderName;
#[cfg(not(target_arch = "wasm32"))]
fn load_from_path(
&self,
settings: &DataLoaderSettings,
path: std::path::PathBuf,
tx: crossbeam::channel::Sender<LoadedData>,
) -> Result<(), DataLoaderError>;
fn load_from_file_contents(
&self,
settings: &DataLoaderSettings,
filepath: std::path::PathBuf,
contents: std::borrow::Cow<'_, [u8]>,
tx: crossbeam::channel::Sender<LoadedData>,
) -> Result<(), DataLoaderError>;
}
#[derive(thiserror::Error, Debug)]
pub enum DataLoaderError {
#[cfg(not(target_arch = "wasm32"))]
#[error(transparent)]
IO(#[from] std::io::Error),
#[error(transparent)]
Arrow(#[from] arrow::error::ArrowError),
#[error(transparent)]
Chunk(#[from] re_chunk::ChunkError),
#[error(transparent)]
Decode(#[from] re_log_encoding::DecodeError),
#[error("No data-loader support for {0:?}")]
Incompatible(std::path::PathBuf),
#[error(transparent)]
Mcap(#[from] ::mcap::McapError),
#[error("{}", re_error::format(.0))]
Other(#[from] anyhow::Error),
}
impl DataLoaderError {
#[inline]
pub fn is_path_not_found(&self) -> bool {
match self {
#[cfg(not(target_arch = "wasm32"))]
Self::IO(err) => err.kind() == std::io::ErrorKind::NotFound,
_ => false,
}
}
#[inline]
pub fn is_incompatible(&self) -> bool {
matches!(self, Self::Incompatible { .. })
}
}
#[derive(Debug)]
pub enum LoadedData {
Chunk(DataLoaderName, re_log_types::StoreId, Chunk),
ArrowMsg(DataLoaderName, re_log_types::StoreId, ArrowMsg),
LogMsg(DataLoaderName, LogMsg),
}
impl LoadedData {
#[inline]
pub fn data_loader_name(&self) -> &DataLoaderName {
match self {
Self::Chunk(name, ..) | Self::ArrowMsg(name, ..) | Self::LogMsg(name, ..) => name,
}
}
#[inline]
pub fn into_log_msg(self) -> ChunkResult<LogMsg> {
match self {
Self::Chunk(_name, store_id, chunk) => {
Ok(LogMsg::ArrowMsg(store_id, chunk.to_arrow_msg()?))
}
Self::ArrowMsg(_name, store_id, msg) => Ok(LogMsg::ArrowMsg(store_id, msg)),
Self::LogMsg(_name, msg) => Ok(msg),
}
}
pub fn into_chunk(self) -> Option<Chunk> {
match self {
Self::Chunk(_name, _store_id, chunk) => Some(chunk),
Self::ArrowMsg(_name, _store_id, arrow_msg) => Chunk::from_arrow_msg(&arrow_msg).ok(),
Self::LogMsg(_name, msg) => match msg {
LogMsg::ArrowMsg(_store_id, arrow_msg) => Chunk::from_arrow_msg(&arrow_msg).ok(),
LogMsg::SetStoreInfo { .. } | LogMsg::BlueprintActivationCommand { .. } => None,
},
}
}
}
static BUILTIN_LOADERS: LazyLock<Vec<Arc<dyn DataLoader>>> = LazyLock::new(|| {
vec![
Arc::new(RrdLoader) as Arc<dyn DataLoader>,
Arc::new(ArchetypeLoader),
Arc::new(DirectoryLoader),
Arc::new(McapLoader::default()),
#[cfg(not(target_arch = "wasm32"))]
Arc::new(LeRobotDatasetLoader),
#[cfg(not(target_arch = "wasm32"))]
Arc::new(ExternalLoader),
Arc::new(UrdfDataLoader),
]
});
#[inline]
pub fn iter_loaders() -> impl Iterator<Item = Arc<dyn DataLoader>> {
BUILTIN_LOADERS
.clone()
.into_iter()
.chain(CUSTOM_LOADERS.read().clone())
}
static CUSTOM_LOADERS: LazyLock<parking_lot::RwLock<Vec<Arc<dyn DataLoader>>>> =
LazyLock::new(parking_lot::RwLock::default);
#[inline]
pub fn register_custom_data_loader(loader: impl DataLoader + 'static) {
CUSTOM_LOADERS.write().push(Arc::new(loader));
}
#[inline]
pub(crate) fn extension(path: &std::path::Path) -> String {
path.extension()
.unwrap_or_default()
.to_ascii_lowercase()
.to_string_lossy()
.to_string()
}
pub const SUPPORTED_IMAGE_EXTENSIONS: &[&str] = &[
"avif", "bmp", "dds", "exr", "farbfeld", "ff", "gif", "hdr", "ico", "jpeg", "jpg", "pam",
"pbm", "pgm", "png", "ppm", "tga", "tif", "tiff", "webp",
];
pub const SUPPORTED_DEPTH_IMAGE_EXTENSIONS: &[&str] = &["rvl", "png"];
pub const SUPPORTED_VIDEO_EXTENSIONS: &[&str] = &["mp4"];
pub const SUPPORTED_MESH_EXTENSIONS: &[&str] = &["glb", "gltf", "obj", "stl", "dae"];
pub const SUPPORTED_POINT_CLOUD_EXTENSIONS: &[&str] = &["ply"];
pub const SUPPORTED_RERUN_EXTENSIONS: &[&str] = &["rbl", "rrd"];
pub const SUPPORTED_THIRD_PARTY_FORMATS: &[&str] = &["mcap", "urdf"];
pub const SUPPORTED_TEXT_EXTENSIONS: &[&str] = &["txt", "md"];
pub fn supported_extensions() -> impl Iterator<Item = &'static str> {
SUPPORTED_RERUN_EXTENSIONS
.iter()
.chain(SUPPORTED_THIRD_PARTY_FORMATS)
.chain(SUPPORTED_IMAGE_EXTENSIONS)
.chain(SUPPORTED_DEPTH_IMAGE_EXTENSIONS)
.chain(SUPPORTED_VIDEO_EXTENSIONS)
.chain(SUPPORTED_MESH_EXTENSIONS)
.chain(SUPPORTED_POINT_CLOUD_EXTENSIONS)
.chain(SUPPORTED_TEXT_EXTENSIONS)
.copied()
}
pub fn is_supported_file_extension(extension: &str) -> bool {
re_log::debug_assert!(
!extension.starts_with('.'),
"Expected extension without period, but got {extension:?}"
);
let extension = extension.to_lowercase();
supported_extensions().any(|ext| ext == extension)
}
pub fn detect_format_from_bytes(bytes: &[u8]) -> Option<String> {
let media_type = re_sdk_types::components::MediaType::guess_from_data(bytes)?;
media_type.file_extension().map(|e| e.to_owned())
}
pub fn content_type_to_extension(content_type: &str) -> Option<String> {
let mime = content_type.split(';').next()?.trim();
if mime == "application/octet-stream" {
return None;
}
let media_type = re_sdk_types::components::MediaType(mime.to_owned().into());
media_type.file_extension().map(|e| e.to_owned())
}
#[test]
fn test_supported_extensions() {
assert!(is_supported_file_extension("rrd"));
assert!(is_supported_file_extension("mcap"));
assert!(is_supported_file_extension("png"));
assert!(is_supported_file_extension("urdf"));
}
#[test]
fn test_supported_mcap_decoder_identifiers() {
let identifiers = supported_mcap_decoder_identifiers(true);
let as_strings = identifiers
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>();
assert!(as_strings.contains(&FOXGLOVE_LENSES_IDENTIFIER.to_owned()));
assert!(as_strings.contains(&URDF_DECODER_IDENTIFIER.to_owned()));
assert!(as_strings.contains(&"raw".to_owned()));
assert!(as_strings.contains(&"protobuf".to_owned()));
assert!(as_strings.contains(&"ros2msg".to_owned()));
let unique = as_strings.iter().collect::<std::collections::BTreeSet<_>>();
assert_eq!(as_strings.len(), unique.len());
}
#[test]
fn test_detect_format_from_bytes() {
assert_eq!(
detect_format_from_bytes(b"RRF2xxxxx").as_deref(),
Some("rrd")
);
assert_eq!(
detect_format_from_bytes(b"RRF0xxxxx").as_deref(),
Some("rrd")
);
assert_eq!(
detect_format_from_bytes(&[0x89, 0x4D, 0x43, 0x41, 0x50, 0x30, 0x0D, 0x0A]).as_deref(),
Some("mcap")
);
assert_eq!(
detect_format_from_bytes(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]).as_deref(),
Some("png")
);
assert_eq!(
detect_format_from_bytes(&[0xFF, 0xD8, 0xFF, 0xE0]).as_deref(),
Some("jpg")
);
assert_eq!(
detect_format_from_bytes(b"glTFxxxx").as_deref(),
Some("glb")
);
assert_eq!(
detect_format_from_bytes(b"ply\nxxx").as_deref(),
Some("ply")
);
assert_eq!(detect_format_from_bytes(b"unknown").as_deref(), None);
assert_eq!(detect_format_from_bytes(b"").as_deref(), None);
}
#[test]
fn test_content_type_to_extension() {
assert_eq!(
content_type_to_extension("image/png").as_deref(),
Some("png")
);
assert_eq!(
content_type_to_extension("image/png; charset=utf-8").as_deref(),
Some("png")
);
assert_eq!(
content_type_to_extension("image/jpeg").as_deref(),
Some("jpg")
);
assert_eq!(
content_type_to_extension("video/mp4").as_deref(),
Some("mp4")
);
assert_eq!(
content_type_to_extension("model/gltf-binary").as_deref(),
Some("glb")
);
assert_eq!(
content_type_to_extension("application/x-rerun").as_deref(),
Some("rrd")
);
assert_eq!(
content_type_to_extension("application/octet-stream").as_deref(),
None
);
}