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#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "snake_case")]
14pub enum ImageFormat {
15 Gray,
17 Rgb,
19 Rgba,
21}
22
23impl ImageFormat {
24 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#[derive(Clone, Debug, PartialEq, Serialize)]
40pub struct Image {
41 format: ImageFormat,
42 pixels: Tensor,
43}
44
45impl Image {
46 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 pub fn format(&self) -> ImageFormat {
66 self.format
67 }
68 pub fn pixels(&self) -> &Tensor {
70 &self.pixels
71 }
72 pub fn width(&self) -> usize {
74 self.pixels.shape()[1]
75 }
76 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#[derive(Clone, Debug, PartialEq, Serialize)]
101pub struct Audio {
102 sample_rate: u32,
103 samples: Tensor,
104}
105
106impl Audio {
107 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 pub fn sample_rate(&self) -> u32 {
127 self.sample_rate
128 }
129 pub fn samples(&self) -> &Tensor {
131 &self.samples
132 }
133 pub fn frames(&self) -> usize {
135 self.samples.shape()[0]
136 }
137 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#[derive(Clone, Debug, PartialEq)]
159pub struct VideoFrame {
160 timestamp_ns: u64,
161 image: Image,
162}
163
164impl VideoFrame {
165 pub fn new(timestamp_ns: u64, image: Image) -> Self {
167 Self {
168 timestamp_ns,
169 image,
170 }
171 }
172 pub fn timestamp_ns(&self) -> u64 {
174 self.timestamp_ns
175 }
176 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#[derive(Clone, Debug, PartialEq)]
213pub struct Video {
214 frames: Arc<[VideoFrame]>,
215}
216
217impl Video {
218 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 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}