Refactor MaxImageDimensionSubscriber

Cargo.lock

@@ -5888,6 +5888,7 @@ dependencies = [
  "mimalloc",
  "nohash-hasher",
  "once_cell",
+ "re_arrow2",
  "re_chunk_store",
  "re_data_ui",
  "re_entity_db",
@@ -5902,6 +5903,7 @@ dependencies = [
  "re_tracing",
  "re_types",
  "re_ui",
+ "re_video",
  "re_viewer_context",
  "re_viewport_blueprint",
  "serde",
@@ -6306,6 +6308,7 @@ dependencies = [
  "re_types",
  "re_types_core",
  "re_ui",
+ "re_video",
  "rfd",
  "serde",
  "slotmap",

crates/store/re_data_loader/src/loader_archetype.rs

@@ -39,8 +39,10 @@ impl DataLoader for ArchetypeLoader {
 
         re_tracing::profile_function!(filepath.display().to_string());
 
-        let contents = std::fs::read(&filepath)
-            .with_context(|| format!("Failed to read file {filepath:?}"))?;
+        let contents = {
+            re_tracing::profile_scope!("fs::read");
+            std::fs::read(&filepath).with_context(|| format!("Failed to read file {filepath:?}"))?
+        };
         let contents = std::borrow::Cow::Owned(contents);
 
         self.load_from_file_contents(settings, filepath, contents, tx)

crates/store/re_types/src/components/media_type_ext.rs

@@ -228,6 +228,11 @@ impl MediaType {
     pub fn is_image(&self) -> bool {
         self.as_str().starts_with("image/")
     }
+
+    /// Returns `true` if this is an video media type.
+    pub fn is_video(&self) -> bool {
+        self.as_str().starts_with("video/")
+    }
 }
 
 impl std::fmt::Display for MediaType {

crates/store/re_video/src/demux/mod.rs

@@ -75,6 +75,7 @@ impl VideoData {
     /// TODO(andreas, jan): This should not copy the data, but instead store slices into a shared buffer.
     /// at the very least the should be a way to extract only metadata.
     pub fn load_from_bytes(data: &[u8], media_type: &str) -> Result<Self, VideoLoadError> {
+        re_tracing::profile_function!();
         match media_type {
             "video/mp4" => Self::load_mp4(data),
 
@@ -98,6 +99,12 @@ impl VideoData {
         std::time::Duration::from_nanos(self.duration.into_nanos(self.timescale) as _)
     }
 
+    /// Natural width and height of the video
+    #[inline]
+    pub fn dimensions(&self) -> [u32; 2] {
+        [self.width(), self.height()]
+    }
+
     /// Natural width of the video.
     #[inline]
     pub fn width(&self) -> u32 {

crates/store/re_video/src/demux/mp4.rs

@@ -6,7 +6,11 @@ use crate::{Time, Timescale};
 
 impl VideoData {
     pub fn load_mp4(bytes: &[u8]) -> Result<Self, VideoLoadError> {
-        let mp4 = re_mp4::Mp4::read_bytes(bytes)?;
+        re_tracing::profile_function!();
+        let mp4 = {
+            re_tracing::profile_scope!("Mp4::read_bytes");
+            re_mp4::Mp4::read_bytes(bytes)?
+        };
 
         let mp4_tracks = mp4.tracks().iter().map(|(k, t)| (*k, t.kind)).collect();
 

crates/viewer/re_renderer/src/video/mod.rs

@@ -5,7 +5,7 @@ use std::{collections::hash_map::Entry, ops::Range, sync::Arc};
 use ahash::HashMap;
 use parking_lot::Mutex;
 
-use re_video::VideoLoadError;
+use re_video::VideoData;
 
 use crate::{resource_managers::GpuTexture2D, RenderContext};
 
@@ -160,19 +160,17 @@ impl Video {
     /// - `video/mp4`
     pub fn load(
         debug_name: String,
-        data: &[u8],
-        media_type: &str,
+        data: Arc<VideoData>,
         decode_hw_acceleration: DecodeHardwareAcceleration,
-    ) -> Result<Self, VideoLoadError> {
-        let data = Arc::new(re_video::VideoData::load_from_bytes(data, media_type)?);
+    ) -> Self {
         let decoders = Mutex::new(HashMap::default());
 
-        Ok(Self {
+        Self {
             debug_name,
             data,
             decoders,
             decode_hw_acceleration,
-        })
+        }
     }
 
     /// The video data

crates/viewer/re_space_view_spatial/Cargo.toml

@@ -33,13 +33,15 @@ re_renderer = { workspace = true, features = [
   "import-obj",
   "import-stl",
 ] }
-re_types = { workspace = true, features = ["ecolor", "glam", "image"] }
+re_space_view.workspace = true
 re_tracing.workspace = true
+re_types = { workspace = true, features = ["ecolor", "glam", "image"] }
 re_ui.workspace = true
+re_video.workspace = true
 re_viewer_context.workspace = true
 re_viewport_blueprint.workspace = true
-re_space_view.workspace = true
 
+arrow2.workspace = true
 ahash.workspace = true
 anyhow.workspace = true
 bitflags.workspace = true

crates/viewer/re_space_view_spatial/src/max_image_dimension_subscriber.rs

@@ -1,14 +1,14 @@
 use ahash::HashMap;
+use arrow2::array::Array;
 use nohash_hasher::IntMap;
 use once_cell::sync::OnceCell;
 
 use re_chunk_store::{ChunkStore, ChunkStoreSubscriber, ChunkStoreSubscriberHandle};
 use re_log_types::{EntityPath, StoreId};
 use re_types::{
-    archetypes::EncodedImage,
     components::{Blob, ImageFormat, MediaType},
     external::image,
-    Archetype, Loggable,
+    Loggable,
 };
 
 #[derive(Debug, Clone, Default)]
@@ -17,11 +17,12 @@ pub struct MaxDimensions {
     pub height: u32,
 }
 
+/// The size of the largest image and/or video at a given entity path.
 #[derive(Default, Debug, Clone)]
 pub struct MaxImageDimensions(IntMap<EntityPath, MaxDimensions>);
 
 impl MaxImageDimensions {
-    /// Accesses the image dimension information for a given store
+    /// Accesses the image/video dimension information for a given store
     pub fn access<T>(
         store_id: &StoreId,
         f: impl FnOnce(&IntMap<EntityPath, MaxDimensions>) -> T,
@@ -76,6 +77,7 @@ impl ChunkStoreSubscriber for MaxImageDimensionSubscriber {
                 continue;
             }
 
+            // Handle `Image`, `DepthImage`, `SegmentationImage`…
             if let Some(all_dimensions) = event.diff.chunk.components().get(&ImageFormat::name()) {
                 for new_dim in all_dimensions.iter().filter_map(|array| {
                     array
@@ -89,62 +91,75 @@ impl ChunkStoreSubscriber for MaxImageDimensionSubscriber {
                         .entry(event.diff.chunk.entity_path().clone())
                         .or_default();
 
-                    max_dim.height = max_dim.height.max(new_dim.height);
                     max_dim.width = max_dim.width.max(new_dim.width);
+                    max_dim.height = max_dim.height.max(new_dim.height);
                 }
             }
 
-            // TODO(jleibs): Image blob/mediatypes should have their own component
-            // Is there a more canonical way to check the indicators for a chunk?
-            if event
-                .diff
-                .chunk
-                .components()
-                .get(&EncodedImage::indicator().name())
-                .is_some()
+            // Handle `ImageEncoded`, `AssetVideo`…
+            let blobs = event.diff.chunk.iter_component_arrays(&Blob::name());
+            let media_types = event.diff.chunk.iter_component_arrays(&MediaType::name());
+            for (blob, media_type) in
+                itertools::izip!(blobs, media_types.map(Some).chain(std::iter::repeat(None)))
             {
-                let media_types = event.diff.chunk.iter_component_arrays(&MediaType::name());
-                let blobs = event.diff.chunk.iter_component_arrays(&Blob::name());
-                for (media, blob) in media_types.zip(blobs) {
-                    let Ok(media) = MediaType::from_arrow_opt(media.as_ref()) else {
-                        continue;
-                    };
-                    let Ok(blob) = Blob::from_arrow_opt(blob.as_ref()) else {
-                        continue;
-                    };
-                    if let (media, Some(Some(blob))) = (media.first(), blob.first()) {
-                        let image_bytes = blob.0.as_slice();
-
-                        let mut reader = image::io::Reader::new(std::io::Cursor::new(image_bytes));
-
-                        if let Some(Some(media)) = media {
-                            if let Some(format) = image::ImageFormat::from_mime_type(&media.0) {
-                                reader.set_format(format);
-                            }
-                        }
-
-                        if reader.format().is_none() {
-                            if let Ok(format) = image::guess_format(image_bytes) {
-                                // Weirdly enough, `reader.decode` doesn't do this for us.
-                                reader.set_format(format);
-                            }
-                        }
-
-                        if let Ok((width, height)) = reader.into_dimensions() {
-                            let max_dim = self
-                                .max_dimensions
-                                .entry(event.store_id.clone())
-                                .or_default()
-                                .0
-                                .entry(event.diff.chunk.entity_path().clone())
-                                .or_default();
-
-                            max_dim.height = max_dim.height.max(height);
-                            max_dim.width = max_dim.width.max(width);
-                        }
-                    }
+                if let Some([width, height]) = size_from_blob(blob.as_ref(), media_type.as_deref())
+                {
+                    let max_dim = self
+                        .max_dimensions
+                        .entry(event.store_id.clone())
+                        .or_default()
+                        .0
+                        .entry(event.diff.chunk.entity_path().clone())
+                        .or_default();
+                    max_dim.width = max_dim.width.max(width);
+                    max_dim.height = max_dim.height.max(height);
                 }
             }
         }
     }
 }
+
+fn size_from_blob(blob: &dyn Array, media_type: Option<&dyn Array>) -> Option<[u32; 2]> {
+    re_tracing::profile_function!();
+
+    let blob = Blob::from_arrow_opt(blob).ok()?.first()?.clone()?;
+
+    let media_type: Option<MediaType> = media_type
+        .and_then(|media_type| MediaType::from_arrow_opt(media_type).ok())
+        .and_then(|list| list.first().cloned())
+        .flatten();
+
+    let media_type = MediaType::or_guess_from_data(media_type, &blob)?;
+
+    if media_type.is_image() {
+        re_tracing::profile_scope!("image");
+
+        let image_bytes = blob.0.as_slice();
+
+        let mut reader = image::io::Reader::new(std::io::Cursor::new(image_bytes));
+
+        if let Some(format) = image::ImageFormat::from_mime_type(&media_type.0) {
+            reader.set_format(format);
+        }
+
+        if reader.format().is_none() {
+            if let Ok(format) = image::guess_format(image_bytes) {
+                // Weirdly enough, `reader.decode` doesn't do this for us.
+                reader.set_format(format);
+            }
+        }
+
+        reader.into_dimensions().ok().map(|size| size.into())
+    } else if media_type.is_video() {
+        re_tracing::profile_scope!("video");
+        if true {
+            None // TODO(#7821): Use the VideoCache here so we make sure we only load each video ONCE
+        } else {
+            re_video::VideoData::load_from_bytes(&blob, &media_type)
+                .ok()
+                .map(|video| video.dimensions())
+        }
+    } else {
+        None
+    }
+}

crates/viewer/re_viewer_context/Cargo.toml

@@ -35,6 +35,7 @@ re_tracing.workspace = true
 re_types = { workspace = true, features = ["ecolor", "glam", "image"] }
 re_types_core.workspace = true
 re_ui.workspace = true
+re_video.workspace = true
 
 ahash.workspace = true
 anyhow.workspace = true

crates/viewer/re_viewer_context/src/cache/video_cache.rs

@@ -43,7 +43,7 @@ impl VideoCache {
         media_type: Option<&MediaType>,
         hw_acceleration: DecodeHardwareAcceleration,
     ) -> Arc<Result<Video, VideoLoadError>> {
-        re_tracing::profile_function!();
+        re_tracing::profile_function!(&debug_name);
 
         // In order to avoid loading the same video multiple times with
         // known and unknown media type, we have to resolve the media type before
@@ -63,8 +63,8 @@ impl VideoCache {
             .or_default()
             .entry(inner_key)
             .or_insert_with(|| {
-                let video =
-                    Video::load(debug_name, video_data, media_type.as_str(), hw_acceleration);
+                let video = re_video::VideoData::load_from_bytes(video_data, &media_type)
+                    .map(|data| Video::load(debug_name, Arc::new(data), hw_acceleration));
                 Entry {
                     used_this_frame: AtomicBool::new(true),
                     video: Arc::new(video),