isaac_sim_arrow/camera/
info.rs

1// SPDX-License-Identifier: MPL-2.0
2//! Arrow encoder and decoder for the camera-info (calibration metadata) channel.
3use std::sync::{Arc, OnceLock};
4
5use arrow::array::{
6    Array, ArrayRef, Float32Array, Float64Array, Int32Array, Int64Array, ListArray, StringArray,
7    StructArray,
8};
9use arrow::buffer::OffsetBuffer;
10use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
11use arrow::record_batch::RecordBatch;
12
13/// Borrowed view of camera calibration metadata, used as input to [`to_record_batch`].
14#[allow(missing_docs)]
15pub struct CameraInfo<'a> {
16    pub frame_id: &'a str,
17    pub distortion_model: &'a str,
18    pub projection_type: &'a str,
19    pub k: &'a [f64],
20    pub r: &'a [f64],
21    pub p: &'a [f64],
22    pub distortion: &'a [f32],
23    pub width: i32,
24    pub height: i32,
25    pub timestamp_ns: i64,
26}
27
28/// Owned variant returned by [`from_struct_array`].
29#[derive(Debug, Clone, PartialEq)]
30#[allow(missing_docs)]
31pub struct CameraInfoOwned {
32    pub frame_id: String,
33    pub distortion_model: String,
34    pub projection_type: String,
35    pub k: Vec<f64>,
36    pub r: Vec<f64>,
37    pub p: Vec<f64>,
38    pub distortion: Vec<f32>,
39    pub width: i32,
40    pub height: i32,
41    pub timestamp_ns: i64,
42}
43
44/// Stable Arrow schema for a `CameraInfo` record batch.
45pub fn schema() -> SchemaRef {
46    static SCHEMA: OnceLock<SchemaRef> = OnceLock::new();
47    SCHEMA
48        .get_or_init(|| {
49            Arc::new(Schema::new(vec![
50                Field::new("frame_id", DataType::Utf8, false),
51                Field::new("distortion_model", DataType::Utf8, false),
52                Field::new("projection_type", DataType::Utf8, false),
53                Field::new(
54                    "k",
55                    DataType::List(Arc::new(Field::new("item", DataType::Float64, false))),
56                    false,
57                ),
58                Field::new(
59                    "r",
60                    DataType::List(Arc::new(Field::new("item", DataType::Float64, false))),
61                    false,
62                ),
63                Field::new(
64                    "p",
65                    DataType::List(Arc::new(Field::new("item", DataType::Float64, false))),
66                    false,
67                ),
68                Field::new(
69                    "distortion",
70                    DataType::List(Arc::new(Field::new("item", DataType::Float32, false))),
71                    false,
72                ),
73                Field::new("width", DataType::Int32, false),
74                Field::new("height", DataType::Int32, false),
75                Field::new("timestamp_ns", DataType::Int64, false),
76            ]))
77        })
78        .clone()
79}
80
81fn list_f64(values: &[f64]) -> ListArray {
82    let inner = Float64Array::from_iter_values(values.iter().copied());
83    let offsets = OffsetBuffer::from_lengths([values.len()]);
84    ListArray::new(
85        Arc::new(Field::new("item", DataType::Float64, false)),
86        offsets,
87        Arc::new(inner),
88        None,
89    )
90}
91
92fn list_f32(values: &[f32]) -> ListArray {
93    let inner = Float32Array::from_iter_values(values.iter().copied());
94    let offsets = OffsetBuffer::from_lengths([values.len()]);
95    ListArray::new(
96        Arc::new(Field::new("item", DataType::Float32, false)),
97        offsets,
98        Arc::new(inner),
99        None,
100    )
101}
102
103/// Encode a `CameraInfo` frame as a single-row `RecordBatch` matching [`schema`].
104///
105/// # Example
106///
107/// ```
108/// use isaac_sim_arrow::camera::info::{CameraInfo, to_record_batch};
109/// let k = [500.0_f64, 0.0, 320.0, 0.0, 500.0, 240.0, 0.0, 0.0, 1.0];
110/// let r = [1.0_f64, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
111/// let p = [500.0_f64, 0.0, 320.0, 0.0, 0.0, 500.0, 240.0, 0.0, 0.0, 0.0, 1.0, 0.0];
112/// let d = [0.0_f32; 5];
113/// let info = CameraInfo { frame_id: "cam", distortion_model: "plumb_bob",
114///     projection_type: "pinhole", k: &k, r: &r, p: &p, distortion: &d,
115///     width: 640, height: 480, timestamp_ns: 1 };
116/// let batch = to_record_batch(&info).unwrap();
117/// assert_eq!(batch.num_rows(), 1);
118/// assert_eq!(batch.num_columns(), 10);
119/// ```
120pub fn to_record_batch(info: &CameraInfo) -> Result<RecordBatch, arrow::error::ArrowError> {
121    let columns: Vec<ArrayRef> = vec![
122        Arc::new(StringArray::from(vec![info.frame_id])),
123        Arc::new(StringArray::from(vec![info.distortion_model])),
124        Arc::new(StringArray::from(vec![info.projection_type])),
125        Arc::new(list_f64(info.k)),
126        Arc::new(list_f64(info.r)),
127        Arc::new(list_f64(info.p)),
128        Arc::new(list_f32(info.distortion)),
129        Arc::new(Int32Array::from_iter_values(std::iter::once(info.width))),
130        Arc::new(Int32Array::from_iter_values(std::iter::once(info.height))),
131        Arc::new(Int64Array::from_iter_values(std::iter::once(
132            info.timestamp_ns,
133        ))),
134    ];
135    RecordBatch::try_new(schema(), columns)
136}
137
138/// Decode the first row of a `StructArray` into a heap-owned `CameraInfoOwned`.
139///
140/// # Example
141///
142/// ```
143/// use arrow::array::StructArray;
144/// use isaac_sim_arrow::camera::info::{CameraInfo, to_record_batch, from_struct_array};
145/// let k = [500.0_f64, 0.0, 320.0, 0.0, 500.0, 240.0, 0.0, 0.0, 1.0];
146/// let r = [1.0_f64, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
147/// let p = [500.0_f64, 0.0, 320.0, 0.0, 0.0, 500.0, 240.0, 0.0, 0.0, 0.0, 1.0, 0.0];
148/// let d = [0.0_f32; 5];
149/// let info = CameraInfo { frame_id: "cam", distortion_model: "plumb_bob",
150///     projection_type: "pinhole", k: &k, r: &r, p: &p, distortion: &d,
151///     width: 640, height: 480, timestamp_ns: 2 };
152/// let batch = to_record_batch(&info).unwrap();
153/// let array = StructArray::from(batch);
154/// let owned = from_struct_array(&array).unwrap();
155/// assert_eq!(owned.frame_id, "cam");
156/// assert_eq!(owned.width, 640);
157/// ```
158pub fn from_struct_array(array: &StructArray) -> Result<CameraInfoOwned, arrow::error::ArrowError> {
159    if array.is_empty() {
160        return Err(arrow::error::ArrowError::InvalidArgumentError(
161            "camera_info struct array is empty".into(),
162        ));
163    }
164    let str_at = |idx: usize, name: &str| -> Result<String, arrow::error::ArrowError> {
165        array
166            .column(idx)
167            .as_any()
168            .downcast_ref::<StringArray>()
169            .ok_or_else(|| {
170                arrow::error::ArrowError::SchemaError(format!("camera_info '{name}' not Utf8"))
171            })
172            .map(|a| a.value(0).to_string())
173    };
174    let f64_list = |idx: usize, name: &str| -> Result<Vec<f64>, arrow::error::ArrowError> {
175        let list = array
176            .column(idx)
177            .as_any()
178            .downcast_ref::<ListArray>()
179            .ok_or_else(|| {
180                arrow::error::ArrowError::SchemaError(format!("camera_info '{name}' not ListArray"))
181            })?;
182        Ok(list
183            .values()
184            .as_any()
185            .downcast_ref::<Float64Array>()
186            .ok_or_else(|| {
187                arrow::error::ArrowError::SchemaError(format!(
188                    "camera_info '{name}' inner not Float64"
189                ))
190            })?
191            .values()
192            .to_vec())
193    };
194    let f32_list = |idx: usize, name: &str| -> Result<Vec<f32>, arrow::error::ArrowError> {
195        let list = array
196            .column(idx)
197            .as_any()
198            .downcast_ref::<ListArray>()
199            .ok_or_else(|| {
200                arrow::error::ArrowError::SchemaError(format!("camera_info '{name}' not ListArray"))
201            })?;
202        Ok(list
203            .values()
204            .as_any()
205            .downcast_ref::<Float32Array>()
206            .ok_or_else(|| {
207                arrow::error::ArrowError::SchemaError(format!(
208                    "camera_info '{name}' inner not Float32"
209                ))
210            })?
211            .values()
212            .to_vec())
213    };
214    Ok(CameraInfoOwned {
215        frame_id: str_at(0, "frame_id")?,
216        distortion_model: str_at(1, "distortion_model")?,
217        projection_type: str_at(2, "projection_type")?,
218        k: f64_list(3, "k")?,
219        r: f64_list(4, "r")?,
220        p: f64_list(5, "p")?,
221        distortion: f32_list(6, "distortion")?,
222        width: array
223            .column(7)
224            .as_any()
225            .downcast_ref::<Int32Array>()
226            .ok_or_else(|| {
227                arrow::error::ArrowError::SchemaError("camera_info 'width' not Int32".into())
228            })?
229            .value(0),
230        height: array
231            .column(8)
232            .as_any()
233            .downcast_ref::<Int32Array>()
234            .ok_or_else(|| {
235                arrow::error::ArrowError::SchemaError("camera_info 'height' not Int32".into())
236            })?
237            .value(0),
238        timestamp_ns: array
239            .column(9)
240            .as_any()
241            .downcast_ref::<Int64Array>()
242            .ok_or_else(|| {
243                arrow::error::ArrowError::SchemaError("camera_info 'timestamp_ns' not Int64".into())
244            })?
245            .value(0),
246    })
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use arrow::array::Array;
253
254    #[test]
255    fn round_trips_through_record_batch() {
256        let k = [500.0_f64, 0.0, 320.0, 0.0, 500.0, 240.0, 0.0, 0.0, 1.0];
257        let r = [1.0_f64, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
258        let p = [
259            500.0_f64, 0.0, 320.0, 0.0, 0.0, 500.0, 240.0, 0.0, 0.0, 0.0, 1.0, 0.0,
260        ];
261        let d = [0.0_f32; 5];
262        let info = CameraInfo {
263            frame_id: "sim_camera",
264            distortion_model: "plumb_bob",
265            projection_type: "pinhole",
266            k: &k,
267            r: &r,
268            p: &p,
269            distortion: &d,
270            width: 640,
271            height: 480,
272            timestamp_ns: 7,
273        };
274        let batch = to_record_batch(&info).expect("convert");
275        assert_eq!(batch.num_rows(), 1);
276        assert_eq!(batch.num_columns(), 10);
277
278        let frame = batch
279            .column(0)
280            .as_any()
281            .downcast_ref::<StringArray>()
282            .expect("frame_id is Utf8");
283        assert_eq!(frame.value(0), "sim_camera");
284
285        let k_col = batch
286            .column(3)
287            .as_any()
288            .downcast_ref::<ListArray>()
289            .expect("k is ListArray");
290        let k_inner = k_col
291            .values()
292            .as_any()
293            .downcast_ref::<Float64Array>()
294            .expect("k inner is Float64");
295        assert_eq!(k_inner.len(), 9);
296        assert_eq!(k_inner.value(0), 500.0);
297
298        let p_col = batch
299            .column(5)
300            .as_any()
301            .downcast_ref::<ListArray>()
302            .expect("p is ListArray");
303        assert_eq!(p_col.values().len(), 12);
304    }
305
306    #[test]
307    fn from_struct_array_round_trips() {
308        let k = [500.0_f64, 0.0, 320.0, 0.0, 500.0, 240.0, 0.0, 0.0, 1.0];
309        let r = [1.0_f64, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
310        let p = [
311            500.0_f64, 0.0, 320.0, 0.0, 0.0, 500.0, 240.0, 0.0, 0.0, 0.0, 1.0, 0.0,
312        ];
313        let d = [0.0_f32; 5];
314        let info = CameraInfo {
315            frame_id: "sim_camera",
316            distortion_model: "plumb_bob",
317            projection_type: "pinhole",
318            k: &k,
319            r: &r,
320            p: &p,
321            distortion: &d,
322            width: 640,
323            height: 480,
324            timestamp_ns: 7,
325        };
326        let batch = to_record_batch(&info).expect("to");
327        let array = StructArray::from(batch);
328        let owned = from_struct_array(&array).expect("from");
329        assert_eq!(owned.frame_id, "sim_camera");
330        assert_eq!(owned.distortion_model, "plumb_bob");
331        assert_eq!(owned.k, k);
332        assert_eq!(owned.p, p);
333        assert_eq!(owned.distortion, d);
334        assert_eq!(owned.width, 640);
335        assert_eq!(owned.timestamp_ns, 7);
336    }
337}
isaac_sim_arrow/camera/info.rs

isaac_sim_arrow/camera/
info.rs
