OAK-T support (#41)

* Support for having multiple viewers open at the same time.

* Initial OAK-T support, TODO: Pass the string unit over.

* Pylints and pass the unit through to the viewer.

* Fix bug where viewer required a restart after running OAK-T...

* Fixed up the OAK-T support. Survives multiple viewer restarts, supports the Thermal person detection model

* py-lints

* RM get-pip

crates/re_log_types/src/component_types/mod.rs

@@ -65,7 +65,7 @@ pub use scalar::{Scalar, ScalarPlotProps};
 pub use size::Size3D;
 pub use tensor::{
     DecodedTensor, Tensor, TensorCastError, TensorData, TensorDataMeaning, TensorDimension,
-    TensorId,
+    TensorId, TensorColormap
 };
 #[cfg(feature = "image")]
 pub use tensor::{TensorImageLoadError, TensorImageSaveError};

crates/re_log_types/src/component_types/tensor.rs

@@ -342,6 +342,20 @@ pub enum TensorDataMeaning {
     Depth,
 }
 
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ArrowField, ArrowSerialize, ArrowDeserialize)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+#[arrow_field(type = "dense")]
+pub enum TensorColormap {
+    None, // Not sure how to serialize an Option with pyarrow.
+    Grayscale,
+    Inferno,
+    Magma,
+    Plasma,
+    #[default]
+    Turbo,
+    Viridis
+}
+
 /// A Multi-dimensional Tensor.
 ///
 /// All clones are shallow.
@@ -400,6 +414,10 @@ pub struct Tensor {
 
     /// Reciprocal scale of meter unit for depth images
     pub meter: Option<f32>,
+
+    pub colormap: TensorColormap,
+
+    pub unit: Option<String>,
 }
 
 impl Tensor {
@@ -626,13 +644,17 @@ macro_rules! tensor_type {
                         data: TensorData::$variant(Vec::from(slice).into()),
                         meaning: TensorDataMeaning::Unknown,
                         meter: None,
+                        colormap: TensorColormap::None,
+                        unit: None,
                     }),
                     None => Ok(Tensor {
                         tensor_id: TensorId::random(),
                         shape,
                         data: TensorData::$variant(view.iter().cloned().collect::<Vec<_>>().into()),
                         meaning: TensorDataMeaning::Unknown,
                         meter: None,
+                        colormap: TensorColormap::None,
+                        unit: None,
                     }),
                 }
             }
@@ -658,6 +680,8 @@ macro_rules! tensor_type {
                         data: TensorData::$variant(value.into_raw_vec().into()),
                         meaning: TensorDataMeaning::Unknown,
                         meter: None,
+                        colormap: TensorColormap::None,
+                        unit: None,
                     })
                     .ok_or(TensorCastError::NotContiguousStdOrder)
             }
@@ -749,14 +773,18 @@ impl Tensor {
         shape: Vec<TensorDimension>,
         data: TensorData,
         meaning: TensorDataMeaning,
-        meter: Option<f32>
+        meter: Option<f32>,
+        colormap: TensorColormap,
+        unit: Option<String>,
     ) -> Self {
         Self {
             tensor_id,
             shape,
             data,
             meaning,
             meter,
+            colormap,
+            unit,
         }
     }
 }
@@ -796,6 +824,8 @@ impl Tensor {
             data: TensorData::JPEG(jpeg_bytes.into()),
             meaning: TensorDataMeaning::Unknown,
             meter: None,
+            colormap: TensorColormap::None,
+            unit: None,
         })
     }
 
@@ -1036,6 +1066,8 @@ impl DecodedTensor {
             data,
             meaning: TensorDataMeaning::Unknown,
             meter: None,
+            colormap: TensorColormap::None,
+            unit: None,
         };
         Ok(DecodedTensor(tensor))
     }
@@ -1121,6 +1153,8 @@ fn test_ndarray() {
         data: TensorData::U16(vec![1, 2, 3, 4].into()),
         meaning: TensorDataMeaning::Unknown,
         meter: None,
+        colormap: TensorColormap::None,
+        unit: None,
     };
     let a0: ndarray::ArrayViewD<'_, u16> = (&t0).try_into().unwrap();
     dbg!(a0); // NOLINT
@@ -1144,6 +1178,8 @@ fn test_arrow() {
             data: TensorData::U16(vec![1, 2, 3, 4].into()),
             meaning: TensorDataMeaning::Unknown,
             meter: Some(1000.0),
+            colormap: TensorColormap::None,
+            unit: None,
         },
         Tensor {
             tensor_id: TensorId(std::default::Default::default()),
@@ -1154,6 +1190,8 @@ fn test_arrow() {
             data: TensorData::F32(vec![1.23, 2.45].into()),
             meaning: TensorDataMeaning::Unknown,
             meter: None,
+            colormap: TensorColormap::None,
+            unit: None,
         }
     ];
 

crates/re_sdk/src/lib.rs

@@ -78,7 +78,7 @@ pub mod components {
         MeshFormat, MeshId, Pinhole, Point2D, Point3D, Quaternion, Radius, RawMesh3D, Rect2D,
         Rigid3, Scalar, ScalarPlotProps, Size3D, Tensor, TensorData, TensorDataMeaning,
         TensorDimension, TensorId, TextEntry, Transform, Vec2D, Vec3D, Vec4D,
-        ViewCoordinates,
+        ViewCoordinates, TensorColormap
     };
 }
 

crates/re_viewer/pipeline.json

@@ -0,0 +1 @@
+{"assetStorage": [], "assets": {"map": {}}, "pipeline": {"connections": [{"node1Id": 3, "node1Output": "out", "node1OutputGroup": "", "node2Id": 2, "node2Input": "in", "node2InputGroup": ""}, {"node1Id": 0, "node1Output": "raw", "node1OutputGroup": "", "node2Id": 1, "node2Input": "in", "node2InputGroup": ""}], "globalProperties": {"calibData": null, "cameraTuningBlobSize": null, "cameraTuningBlobUri": "", "leonCssFrequencyHz": 700000000.0, "leonMssFrequencyHz": 700000000.0, "pipelineName": null, "pipelineVersion": null, "sippBufferSize": 18432, "sippDmaBufferSize": 16384, "xlinkChunkSize": -1}, "nodes": [[0, {"id": 0, "ioInfo": [[["", "video"], {"blocking": false, "group": "", "id": 11, "name": "video", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "still"], {"blocking": false, "group": "", "id": 9, "name": "still", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "isp"], {"blocking": false, "group": "", "id": 8, "name": "isp", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "preview"], {"blocking": false, "group": "", "id": 10, "name": "preview", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "raw"], {"blocking": false, "group": "", "id": 7, "name": "raw", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "frameEvent"], {"blocking": false, "group": "", "id": 6, "name": "frameEvent", "queueSize": 8, "type": 0, "waitForMessage": false}], [["", "inputConfig"], {"blocking": false, "group": "", "id": 5, "name": "inputConfig", "queueSize": 8, "type": 3, "waitForMessage": false}], [["", "inputControl"], {"blocking": true, "group": "", "id": 4, "name": "inputControl", "queueSize": 8, "type": 3, "waitForMessage": false}]], "name": "Camera", "properties": {"boardSocket": 4, "calibAlpha": null, "cameraName": "", "colorOrder": 0, "fp16": false, "fps": 30.0, "imageOrientation": -1, "initialControl": {"aeLockMode": false, "aeMaxExposureTimeUs": 0, "aeRegion": {"height": 0, "priority": 0, "width": 0, "x": 0, "y": 0}, "afRegion": {"height": 0, "priority": 0, "width": 0, "x": 0, "y": 0}, "antiBandingMode": 0, "autoFocusMode": 3, "awbLockMode": false, "awbMode": 0, "brightness": 0, "captureIntent": 0, "chromaDenoise": 0, "cmdMask": 0, "contrast": 0, "controlMode": 0, "effectMode": 0, "expCompensation": 0, "expManual": {"exposureTimeUs": 0, "frameDurationUs": 0, "sensitivityIso": 0}, "frameSyncMode": 0, "lensPosAutoInfinity": 0, "lensPosAutoMacro": 0, "lensPosition": 0, "lensPositionRaw": 0.0, "lowPowerNumFramesBurst": 0, "lowPowerNumFramesDiscard": 0, "lumaDenoise": 0, "saturation": 0, "sceneMode": 0, "sharpness": 0, "strobeConfig": {"activeLevel": 0, "enable": 0, "gpioNumber": 0}, "strobeTimings": {"durationUs": 0, "exposureBeginOffsetUs": 0, "exposureEndOffsetUs": 0}, "wbColorTemp": 0}, "interleaved": true, "isp3aFps": 0, "ispScale": {"horizDenominator": 0, "horizNumerator": 0, "vertDenominator": 0, "vertNumerator": 0}, "numFramesPoolIsp": 3, "numFramesPoolPreview": 4, "numFramesPoolRaw": 3, "numFramesPoolStill": 4, "numFramesPoolVideo": 4, "previewHeight": 300, "previewKeepAspectRatio": true, "previewWidth": 300, "rawPacked": null, "resolutionHeight": -1, "resolutionWidth": -1, "sensorCropX": -1.0, "sensorCropY": -1.0, "sensorType": -1, "stillHeight": -1, "stillWidth": -1, "videoHeight": -1, "videoWidth": -1, "warpMeshHeight": 0, "warpMeshSource": -1, "warpMeshStepHeight": 32, "warpMeshStepWidth": 32, "warpMeshUri": "", "warpMeshWidth": 0}}], [1, {"id": 1, "ioInfo": [[["", "in"], {"blocking": true, "group": "", "id": 3, "name": "in", "queueSize": 8, "type": 3, "waitForMessage": true}]], "name": "XLinkOut", "properties": {"maxFpsLimit": -1.0, "metadataOnly": false, "streamName": "thermal_cam"}}], [2, {"id": 2, "ioInfo": [[["", "in"], {"blocking": true, "group": "", "id": 2, "name": "in", "queueSize": 8, "type": 3, "waitForMessage": true}]], "name": "XLinkOut", "properties": {"maxFpsLimit": -1.0, "metadataOnly": false, "streamName": "sys_logger"}}], [3, {"id": 3, "ioInfo": [[["", "out"], {"blocking": false, "group": "", "id": 1, "name": "out", "queueSize": 8, "type": 0, "waitForMessage": false}]], "name": "SystemLogger", "properties": {"rateHz": 0.10000000149011612}}]]}}

crates/re_viewer/src/depthai/depthai.rs

@@ -122,6 +122,7 @@ pub enum ImuKind {
 #[derive(serde::Serialize, serde::Deserialize, Clone, Copy, PartialEq, Debug)]
 #[allow(non_camel_case_types)]
 pub enum CameraSensorResolution {
+    THE_256X192,
     THE_400_P,
     THE_480_P,
     THE_720_P,
@@ -142,6 +143,7 @@ pub enum CameraSensorResolution {
 impl fmt::Display for CameraSensorResolution {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
+            Self::THE_256X192 => write!(f, "256x192"),
             Self::THE_400_P => write!(f, "400p"),
             Self::THE_480_P => write!(f, "480p"),
             Self::THE_720_P => write!(f, "720p"),
@@ -322,7 +324,7 @@ pub struct DeviceConfig {
     pub depth_enabled: bool, // Much easier to have an explicit bool for checkbox
     #[serde(default = "StereoDepthConfig::default_as_option")]
     pub depth: Option<StereoDepthConfig>,
-    pub ai_model: AiModel,
+    pub ai_model: Option<AiModel>,
     #[serde(skip)]
     pub dot_brightness: u32,
     #[serde(skip)]
@@ -336,7 +338,7 @@ impl Default for DeviceConfig {
             cameras: Vec::new(),
             depth_enabled: true,
             depth: Some(StereoDepthConfig::default()),
-            ai_model: AiModel::default(),
+            ai_model: None,
             dot_brightness: 0,
             flood_brightness: 0,
         }
@@ -370,7 +372,7 @@ impl From<&DeviceProperties> for DeviceConfig {
             })
             .collect();
         config.depth = Option::<StereoDepthConfig>::from(props);
-        config.ai_model = AiModel::from(props);
+        config.ai_model = Option::<AiModel>::from(props);
         config
     }
 }
@@ -498,7 +500,10 @@ pub struct AiModel {
 
 impl Default for AiModel {
     fn default() -> Self {
-        default_neural_networks()[2].clone()
+        match default_neural_networks()[2].clone() {
+            Some(model) => model,
+            None => panic!("Default neural network not found!")
+        }
     }
 }
 
@@ -512,16 +517,16 @@ impl AiModel {
     }
 }
 
-impl From<&DeviceProperties> for AiModel {
+impl From<&DeviceProperties> for Option<AiModel> {
     fn from(props: &DeviceProperties) -> Self {
-        let mut model = Self::default();
+        let mut model = AiModel::default();
 
         if let Some(cam) = props.cameras.iter().find(|cam| cam.is_color_camera()) {
             model.camera = cam.board_socket;
-        } else if let Some(cam) = props.cameras.first() {
-            model.camera = cam.board_socket;
+            Some(model)
+        } else {
+            None
         }
-        model
     }
 }
 
@@ -582,8 +587,8 @@ pub struct State {
     pub backend_comms: BackendCommChannel,
     #[serde(skip)]
     poll_instant: Option<Instant>,
-    #[serde(default = "default_neural_networks")]
-    pub neural_networks: Vec<AiModel>,
+    #[serde(skip, default = "default_neural_networks")]
+    pub neural_networks: Vec<Option<AiModel>>,
     #[serde(skip)]
     update_timeout_timer: Option<Instant>,
     #[serde(skip, default = "bool_true")]
@@ -597,29 +602,34 @@ fn all_subscriptions() -> Vec<ChannelId> {
 }
 
 #[inline]
-fn default_neural_networks() -> Vec<AiModel> {
+fn default_neural_networks() -> Vec<Option<AiModel>> {
     vec![
-        AiModel::none(),
-        AiModel {
+        None,
+        Some(AiModel {
             path: String::from("yolov8n_coco_640x352"),
             display_name: String::from("Yolo V8"),
             camera: CameraBoardSocket::CAM_A,
-        },
-        AiModel {
+        }),
+        Some(AiModel {
             path: String::from("mobilenet-ssd"),
             display_name: String::from("MobileNet SSD"),
             camera: CameraBoardSocket::CAM_A,
-        },
-        AiModel {
+        }),
+        Some(AiModel {
             path: String::from("face-detection-retail-0004"),
             display_name: String::from("Face Detection"),
             camera: CameraBoardSocket::CAM_A,
-        },
-        AiModel {
+        }),
+        Some(AiModel {
             path: String::from("age-gender-recognition-retail-0013"),
             display_name: String::from("Age gender recognition"),
             camera: CameraBoardSocket::CAM_A,
-        },
+        }),
+        Some(AiModel {
+            path: String::from("yolov6n_thermal_people_256x192"),
+            display_name: String::from("Thermal Person Detection"),
+            camera: CameraBoardSocket::CAM_E
+        }),
     ]
 }
 

crates/re_viewer/src/gpu_bridge/tensor_to_gpu.rs

@@ -8,7 +8,7 @@ use bytemuck::{ allocation::pod_collect_to_vec, cast_slice, Pod };
 use egui::util::hash;
 use wgpu::TextureFormat;
 
-use re_log_types::component_types::{ DecodedTensor, Tensor, TensorData };
+use re_log_types::component_types::{ DecodedTensor, Tensor, TensorData, TensorColormap };
 use re_renderer::{
     renderer::{ ColorMapper, ColormappedTexture, TextureEncoding },
     resource_managers::Texture2DCreationDesc,
@@ -148,8 +148,16 @@ fn color_tensor_to_gpu(
         encoding != Some(TextureEncoding::Nv12) &&
         re_renderer::texture_info::num_texture_components(texture_format) == 1
     {
-        // Single-channel images = luminance = grayscale
-        Some(ColorMapper::Function(re_renderer::Colormap::Grayscale))
+        match tensor.inner().colormap {
+            TensorColormap::Grayscale => Some(ColorMapper::Function(re_renderer::Colormap::Grayscale)),
+            TensorColormap::Viridis => Some(ColorMapper::Function(re_renderer::Colormap::Viridis)),
+            TensorColormap::Plasma => Some(ColorMapper::Function(re_renderer::Colormap::Plasma)),
+            TensorColormap::Inferno => Some(ColorMapper::Function(re_renderer::Colormap::Inferno)),
+            TensorColormap::Magma => Some(ColorMapper::Function(re_renderer::Colormap::Magma)),
+            TensorColormap::Turbo => Some(ColorMapper::Function(re_renderer::Colormap::Turbo)),
+            // Single-channel images = luminance = grayscale
+            TensorColormap::None => Some(ColorMapper::Function(re_renderer::Colormap::Grayscale)),
+        }
     } else {
         None
     };

crates/re_viewer/src/ui/data_ui/image.rs

@@ -220,7 +220,7 @@ pub fn tensor_summary_ui_grid_contents(
     tensor: &Tensor,
     tensor_stats: &TensorStats
 ) {
-    let Tensor { tensor_id: _, shape, data, meaning, meter } = tensor;
+    let Tensor { tensor_id: _, shape, data, meaning, meter, colormap: _, unit: _ } = tensor;
 
     re_ui
         .grid_left_hand_label(ui, "Data type")
@@ -565,7 +565,20 @@ fn tensor_pixel_value_ui(
                             _ => unreachable!("NV12 should only contain u8"),
                         }
                     }),
-                _ => tensor.get(&[y, x]).map(|v| format!("Val: {v}")),
+                _ => {
+                    match &tensor.unit {
+                        Some(unit) => {
+                            if tensor.dtype().is_float() {
+                                tensor.get(&[y, x]).map(|v| format!("Val: {v:.2} {unit}"))
+                            } else {
+                                tensor.get(&[y, x]).map(|v| format!("Val: {v} {unit}"))
+                            }
+                        },
+                        None => {
+                            tensor.get(&[y, x]).map(|v| format!("Val: {v}"))
+                        }
+                    }
+                }
             }
         3 =>
             match tensor.real_shape().as_slice()[2].size {