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        data_type: crate::TensorDataType,
58        order: crate::TensorDimOrder,
59        dims: &[usize],
60    ) -> Option<&crate::Tensor> {
61        unsafe {
62            let res = ffi::gst_tensor_meta_get_typed_tensor(
63                self.as_mut_ptr(),
64                id.into_glib(),
65                data_type.into_glib(),
66                order.into_glib(),
67                dims.len(),
68                dims.as_ptr(),
69            );
70            if res.is_null() {
71                None
72            } else {
73                // FIXME: This is not ideal but otherwise we can't return a reference safely
74                self.as_slice()
75                    .iter()
76                    .find(|t| std::ptr::eq(t.as_ptr(), res))
77            }
78        }
79    }
80}
81
82unsafe impl MetaAPI for TensorMeta {
83    type GstType = ffi::GstTensorMeta;
84
85    #[doc(alias = "gst_tensor_meta_api_get_type")]
86    #[inline]
87    fn meta_api() -> glib::Type {
88        unsafe { from_glib(ffi::gst_tensor_meta_api_get_type()) }
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use crate::*;
95
96    #[test]
97    fn build_tensor_meta() {
98        gst::init().unwrap();
99
100        let mut buf = gst::Buffer::new();
101
102        let mut tmeta = TensorMeta::add(buf.make_mut());
103
104        let tensor = Tensor::new_simple(
105            glib::Quark::from_str("me"),
106            TensorDataType::Int16,
107            gst::Buffer::with_size(2 * 3 * 4 * 5).unwrap(),
108            TensorDimOrder::RowMajor,
109            &[3, 4, 5],
110        );
111
112        let tptr = tensor.as_ptr();
113
114        tmeta.set([tensor].into());
115
116        let tensors = tmeta.as_slice();
117
118        assert_eq!(tensors.len(), 1);
119
120        // Check that it's the same tensor
121        assert_eq!(tptr, tensors[0].as_ptr());
122        assert_eq!(tensors[0].dims_order(), TensorDimOrder::RowMajor);
123        assert_eq!(tensors[0].dims().len(), 3);
124        assert_eq!(tensors[0].dims()[0], 3);
125
126        assert_eq!(tmeta.as_slice().len(), 1);
127
128        tmeta.as_mut_slice();
129    }
130}
gstreamer_analytics/tensor_meta.rs

gstreamer_analytics/
tensor_meta.rs
