gstreamer_analytics/
tensor_meta.rs

1// Take a look at the license at the top of the repository in the LICENSE file.
2
3use glib::translate::*;
4use gst::prelude::*;
5
6use crate::ffi;
7use crate::Tensor;
8
9#[repr(transparent)]
10#[doc(alias = "GstTensorMeta")]
11pub struct TensorMeta(ffi::GstTensorMeta);
12
13unsafe impl Send for TensorMeta {}
14unsafe impl Sync for TensorMeta {}
15
16impl TensorMeta {
17    #[doc(alias = "gst_buffer_add_tensor_meta")]
18    pub fn add(buffer: &mut gst::BufferRef) -> gst::MetaRefMut<Self, gst::meta::Standalone> {
19        skip_assert_initialized!();
20
21        unsafe {
22            let meta_ptr = ffi::gst_buffer_add_tensor_meta(buffer.as_mut_ptr());
23            Self::from_mut_ptr(buffer, meta_ptr)
24        }
25    }
26
27    #[doc(alias = "gst_tensor_meta_set")]
28    pub fn set(&mut self, tensors: glib::Slice<Tensor>) {
29        unsafe {
30            ffi::gst_tensor_meta_set(self.as_mut_ptr(), tensors.len() as u32, tensors.into_raw());
31        }
32    }
33
34    #[doc(alias = "gst_tensor_meta_get_index_from_id")]
35    pub fn index_from_id(&self, id: glib::Quark) -> i32 {
36        unsafe { ffi::gst_tensor_meta_get_index_from_id(self.as_mut_ptr(), id.into_glib()) }
37    }
38
39    pub fn as_slice(&self) -> &[Tensor] {
40        unsafe { glib::Slice::from_glib_borrow_num(self.0.tensors, self.0.num_tensors) }
41    }
42
43    pub fn as_mut_slice(&mut self) -> &mut [Tensor] {
44        unsafe { glib::Slice::from_glib_borrow_num_mut(self.0.tensors, self.0.num_tensors) }
45    }
46
47    unsafe fn as_mut_ptr(&self) -> *mut ffi::GstTensorMeta {
48        mut_override(&self.0)
49    }
50
51    #[cfg(feature = "v1_28")]
52    #[cfg_attr(docsrs, doc(cfg(feature = "v1_28")))]
53    #[doc(alias = "gst_tensor_meta_get_typed_tensor")]
54    pub fn typed_tensor(
55        &self,
56        id: glib::Quark,
57        order: crate::TensorDimOrder,
58        num_dims: usize,
59        data_type: crate::TensorDataType,
60        data: &gst::BufferRef,
61    ) -> Option<&crate::Tensor> {
62        unsafe {
63            let res = ffi::gst_tensor_meta_get_typed_tensor(
64                self.as_mut_ptr(),
65                id.into_glib(),
66                order.into_glib(),
67                num_dims,
68                data_type.into_glib(),
69                mut_override(data.as_ptr()),
70            );
71            if res.is_null() {
72                None
73            } else {
74                // FIXME: This is not ideal but otherwise we can't return a reference safely
75                self.as_slice()
76                    .iter()
77                    .find(|t| std::ptr::eq(t.as_ptr(), res))
78            }
79        }
80    }
81}
82
83unsafe impl MetaAPI for TensorMeta {
84    type GstType = ffi::GstTensorMeta;
85
86    #[doc(alias = "gst_tensor_meta_api_get_type")]
87    #[inline]
88    fn meta_api() -> glib::Type {
89        unsafe { from_glib(ffi::gst_tensor_meta_api_get_type()) }
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use crate::*;
96
97    #[test]
98    fn build_tensor_meta() {
99        gst::init().unwrap();
100
101        let mut buf = gst::Buffer::new();
102
103        let mut tmeta = TensorMeta::add(buf.make_mut());
104
105        let tensor = Tensor::new_simple(
106            glib::Quark::from_str("me"),
107            TensorDataType::Int16,
108            gst::Buffer::with_size(2 * 3 * 4 * 5).unwrap(),
109            TensorDimOrder::RowMajor,
110            &[3, 4, 5],
111        );
112
113        let tptr = tensor.as_ptr();
114
115        tmeta.set([tensor].into());
116
117        let tensors = tmeta.as_slice();
118
119        assert_eq!(tensors.len(), 1);
120
121        // Check that it's the same tensor
122        assert_eq!(tptr, tensors[0].as_ptr());
123        assert_eq!(tensors[0].dims_order(), TensorDimOrder::RowMajor);
124        assert_eq!(tensors[0].dims().len(), 3);
125        assert_eq!(tensors[0].dims()[0], 3);
126
127        assert_eq!(tmeta.as_slice().len(), 1);
128
129        tmeta.as_mut_slice();
130    }
131}