Skip to main content

konjure_sdk/data/
media.rs

1use super::{DType, DataError, Tensor};
2use alloc::{
3    string::{String, ToString},
4    sync::Arc,
5    vec::Vec,
6};
7use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeStruct};
8
9/// Channel ordering for a decoded image. Rows are top-to-bottom; channels are
10/// interleaved in the last tensor axis. This is storage metadata, not a promise
11/// of calibrated camera orientation, color space, or camera-to-eye alignment.
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "snake_case")]
14pub enum ImageFormat {
15    /// One luminance channel.
16    Gray,
17    /// Red, green, blue.
18    Rgb,
19    /// Red, green, blue, alpha.
20    Rgba,
21}
22
23impl ImageFormat {
24    /// Required size of the last image tensor axis.
25    pub const fn channels(self) -> usize {
26        match self {
27            Self::Gray => 1,
28            Self::Rgb => 3,
29            Self::Rgba => 4,
30        }
31    }
32}
33
34/// Immutable decoded image in `[height, width, channels]` tensor order.
35///
36/// Accepted channels are U8, U16, F16, F32, F64, or F128. Float samples are
37/// finite but may exceed `[0, 1]` for high dynamic range processing. Encoding,
38/// color conversion, capture, and display belong to explicit host adapters.
39#[derive(Clone, Debug, PartialEq, Serialize)]
40pub struct Image {
41    format: ImageFormat,
42    pixels: Tensor,
43}
44
45impl Image {
46    /// Validates the format and nonempty decoded pixel layout.
47    ///
48    /// # Errors
49    /// Rejects incompatible types, zero dimensions, and wrong rank/channel count.
50    pub fn new(format: ImageFormat, pixels: Tensor) -> Result<Self, DataError> {
51        if pixels.shape().len() != 3
52            || pixels.shape().contains(&0)
53            || pixels.shape()[2] != format.channels()
54            || !matches!(
55                pixels.dtype(),
56                DType::U8 | DType::U16 | DType::F16 | DType::F32 | DType::F64 | DType::F128
57            )
58        {
59            return Err(DataError::InvalidMetadata);
60        }
61        Ok(Self { format, pixels })
62    }
63
64    /// Ordered pixel-channel interpretation.
65    pub fn format(&self) -> ImageFormat {
66        self.format
67    }
68    /// Immutable decoded tensor, whose views may be noncontiguous.
69    pub fn pixels(&self) -> &Tensor {
70        &self.pixels
71    }
72    /// Number of columns.
73    pub fn width(&self) -> usize {
74        self.pixels.shape()[1]
75    }
76    /// Number of rows.
77    pub fn height(&self) -> usize {
78        self.pixels.shape()[0]
79    }
80}
81
82impl<'de> Deserialize<'de> for Image {
83    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
84        #[derive(Deserialize)]
85        #[serde(deny_unknown_fields)]
86        struct Wire {
87            format: ImageFormat,
88            pixels: Tensor,
89        }
90        let wire = Wire::deserialize(deserializer)?;
91        Self::new(wire.format, wire.pixels).map_err(serde::de::Error::custom)
92    }
93}
94
95/// Immutable decoded audio in `[frames, channels]` tensor order.
96///
97/// Samples are finite F32 or F64 values. Their amplitude is not clamped, so
98/// intermediate processing can retain headroom. Zero-frame blocks are valid;
99/// channel count and sample rate still describe their intended decoded layout.
100#[derive(Clone, Debug, PartialEq, Serialize)]
101pub struct Audio {
102    sample_rate: u32,
103    samples: Tensor,
104}
105
106impl Audio {
107    /// Validates a sample block with 1..=32 channels and a 1..=384000 Hz rate.
108    ///
109    /// # Errors
110    /// Rejects unsupported sample types, rank, channel count, or sample rate.
111    pub fn new(sample_rate: u32, samples: Tensor) -> Result<Self, DataError> {
112        if !(1..=384_000).contains(&sample_rate)
113            || samples.shape().len() != 2
114            || !(1..=32).contains(&samples.shape()[1])
115            || !matches!(samples.dtype(), DType::F32 | DType::F64)
116        {
117            return Err(DataError::InvalidMetadata);
118        }
119        Ok(Self {
120            sample_rate,
121            samples,
122        })
123    }
124
125    /// Frames per second, shared by every channel.
126    pub fn sample_rate(&self) -> u32 {
127        self.sample_rate
128    }
129    /// Immutable interleaved sample tensor.
130    pub fn samples(&self) -> &Tensor {
131        &self.samples
132    }
133    /// Number of decoded frames, each containing one sample per channel.
134    pub fn frames(&self) -> usize {
135        self.samples.shape()[0]
136    }
137    /// Number of channels; no speaker-layout inference is performed.
138    pub fn channels(&self) -> usize {
139        self.samples.shape()[1]
140    }
141}
142
143impl<'de> Deserialize<'de> for Audio {
144    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
145        #[derive(Deserialize)]
146        #[serde(deny_unknown_fields)]
147        struct Wire {
148            sample_rate: u32,
149            samples: Tensor,
150        }
151        let wire = Wire::deserialize(deserializer)?;
152        Self::new(wire.sample_rate, wire.samples).map_err(serde::de::Error::custom)
153    }
154}
155
156/// One decoded image with a presentation timestamp in a host-defined timeline.
157/// Timestamps serialize as decimal strings to preserve the full unsigned range.
158#[derive(Clone, Debug, PartialEq)]
159pub struct VideoFrame {
160    timestamp_ns: u64,
161    image: Image,
162}
163
164impl VideoFrame {
165    /// Associates a validated image with a nanosecond timestamp.
166    pub fn new(timestamp_ns: u64, image: Image) -> Self {
167        Self {
168            timestamp_ns,
169            image,
170        }
171    }
172    /// Presentation time, without inferring a wall-clock origin.
173    pub fn timestamp_ns(&self) -> u64 {
174        self.timestamp_ns
175    }
176    /// Decoded frame pixels and their channel layout.
177    pub fn image(&self) -> &Image {
178        &self.image
179    }
180}
181
182impl Serialize for VideoFrame {
183    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
184        let mut frame = serializer.serialize_struct("VideoFrame", 2)?;
185        frame.serialize_field("timestamp_ns", &self.timestamp_ns.to_string())?;
186        frame.serialize_field("image", &self.image)?;
187        frame.end()
188    }
189}
190
191impl<'de> Deserialize<'de> for VideoFrame {
192    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
193        #[derive(Deserialize)]
194        #[serde(deny_unknown_fields)]
195        struct Wire {
196            timestamp_ns: String,
197            image: Image,
198        }
199        let wire = Wire::deserialize(deserializer)?;
200        let timestamp_ns = wire
201            .timestamp_ns
202            .parse()
203            .map_err(serde::de::Error::custom)?;
204        Ok(Self::new(timestamp_ns, wire.image))
205    }
206}
207
208/// Immutable decoded video frames with increasing presentation timestamps.
209///
210/// All frames retain a single image format, shape, and channel type. This is an
211/// inspectable value, not a codec, playback clock, capture stream, or audio muxer.
212#[derive(Clone, Debug, PartialEq)]
213pub struct Video {
214    frames: Arc<[VideoFrame]>,
215}
216
217impl Video {
218    /// Validates a nonempty sequence with strictly increasing timestamps and
219    /// an identical decoded layout for every frame.
220    ///
221    /// # Errors
222    /// Rejects empty sequences, more than one million frames, changing layouts,
223    /// duplicate or decreasing timestamps, and overflow in logical byte counts.
224    pub fn new(frames: Vec<VideoFrame>) -> Result<Self, DataError> {
225        let first = frames.first().ok_or(DataError::InvalidMetadata)?;
226        if frames.len() > 1_000_000 {
227            return Err(DataError::AllocationLimit);
228        }
229        let mut previous = None;
230        let mut total_bytes = 0_usize;
231        for frame in &frames {
232            if previous.is_some_and(|timestamp| timestamp >= frame.timestamp_ns)
233                || frame.image.format != first.image.format
234                || frame.image.pixels.dtype() != first.image.pixels.dtype()
235                || frame.image.pixels.shape() != first.image.pixels.shape()
236            {
237                return Err(DataError::InvalidMetadata);
238            }
239            total_bytes = total_bytes
240                .checked_add(frame.image.pixels.byte_len())
241                .ok_or(DataError::AllocationLimit)?;
242            previous = Some(frame.timestamp_ns);
243        }
244        Ok(Self {
245            frames: Arc::from(frames),
246        })
247    }
248
249    /// Borrows the ordered decoded frames; cloning this value shares the list.
250    pub fn frames(&self) -> &[VideoFrame] {
251        &self.frames
252    }
253}
254
255impl Serialize for Video {
256    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
257        let mut video = serializer.serialize_struct("Video", 1)?;
258        video.serialize_field("frames", self.frames())?;
259        video.end()
260    }
261}
262
263impl<'de> Deserialize<'de> for Video {
264    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
265        #[derive(Deserialize)]
266        #[serde(deny_unknown_fields)]
267        struct Wire {
268            frames: Vec<VideoFrame>,
269        }
270        let wire = Wire::deserialize(deserializer)?;
271        Self::new(wire.frames).map_err(serde::de::Error::custom)
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278
279    fn image() -> Image {
280        Image::new(
281            ImageFormat::Rgb,
282            Tensor::zeros(DType::U8, alloc::vec![2, 3, 3]).unwrap(),
283        )
284        .unwrap()
285    }
286
287    #[test]
288    fn decoded_media_validate_layout_and_wire() {
289        assert!(
290            Image::new(
291                ImageFormat::Rgb,
292                Tensor::zeros(DType::U8, alloc::vec![2, 3, 4]).unwrap()
293            )
294            .is_err()
295        );
296        assert!(
297            Image::new(
298                ImageFormat::Rgb,
299                Tensor::zeros(DType::U8, alloc::vec![0, 3, 3]).unwrap()
300            )
301            .is_err()
302        );
303        assert!(Audio::new(0, Tensor::zeros(DType::F32, alloc::vec![32, 2]).unwrap()).is_err());
304        assert!(
305            Audio::new(
306                48000,
307                Tensor::zeros(DType::I32, alloc::vec![32, 2]).unwrap()
308            )
309            .is_err()
310        );
311        assert!(Audio::new(48000, Tensor::zeros(DType::F32, alloc::vec![0, 2]).unwrap()).is_ok());
312        let video = Video::new(alloc::vec![
313            VideoFrame::new(u64::MAX - 1, image()),
314            VideoFrame::new(u64::MAX, image())
315        ])
316        .unwrap();
317        let wire = serde_json::to_string(&video).unwrap();
318        assert!(wire.contains("\"18446744073709551615\""));
319        assert_eq!(serde_json::from_str::<Video>(&wire).unwrap(), video);
320        assert!(
321            Video::new(alloc::vec![
322                VideoFrame::new(5, image()),
323                VideoFrame::new(5, image())
324            ])
325            .is_err()
326        );
327        assert!(
328            serde_json::from_str::<Image>(
329                r#"{"format":"rgba","pixels":{"dtype":"u8","shape":[1,1,3],"data":["1","2","3"]}}"#
330            )
331            .is_err()
332        );
333    }
334}