rskit_dataset/
transform.rs

1//! Transform trait and built-in transforms.
2
3use crate::{DataItem, DatasetItem, DatasetLimits};
4use rskit_errors::AppResult;
5
6#[cfg(feature = "image-transform")]
7use crate::DataPayload;
8#[cfg(feature = "image-transform")]
9use rskit_errors::{AppError, ErrorCode};
10#[cfg(feature = "image-transform")]
11use rskit_validation::Validator;
12
13/// Protocol for data transforms from item type `I` to item type `O`.
14pub trait Transform<I: DatasetItem, O: DatasetItem>: Send + Sync {
15    /// Stable transform name.
16    fn name(&self) -> &str;
17    /// Apply the transform, returning `Ok(None)` when the item is filtered out.
18    fn apply(&self, item: I, limits: &DatasetLimits) -> AppResult<Option<O>>;
19}
20
21/// Resize images to a fixed size and re-encode as JPEG.
22#[cfg(feature = "image-transform")]
23#[derive(Clone)]
24pub struct ResizeTransform {
25    width: u32,
26    height: u32,
27    quality: u8,
28    name: String,
29}
30
31#[cfg(feature = "image-transform")]
32impl ResizeTransform {
33    /// Create a resize transform after validating dimensions.
34    pub fn new(width: u32, height: u32) -> AppResult<Self> {
35        Validator::new()
36            .min_value("width", width, 1)
37            .min_value("height", height, 1)
38            .validate()?;
39        Ok(Self {
40            width,
41            height,
42            quality: 85,
43            name: format!("resize-{width}x{height}"),
44        })
45    }
46
47    /// Set JPEG quality after validating the accepted encoder range.
48    #[must_use = "builder methods return the updated transform"]
49    pub fn with_quality(mut self, quality: u8) -> AppResult<Self> {
50        Validator::new()
51            .in_range("quality", quality, 1, 100)
52            .validate()?;
53        self.quality = quality;
54        Ok(self)
55    }
56}
57
58#[cfg(feature = "image-transform")]
59impl Transform<DataItem, DataItem> for ResizeTransform {
60    fn name(&self) -> &str {
61        &self.name
62    }
63
64    fn apply(&self, item: DataItem, limits: &DatasetLimits) -> AppResult<Option<DataItem>> {
65        use image::ImageReader;
66        use std::io::Cursor;
67
68        let bytes = item.payload().read_bytes_bounded(limits)?;
69        let img = ImageReader::new(Cursor::new(&bytes))
70            .with_guessed_format()
71            .map_err(|error| {
72                AppError::new(
73                    ErrorCode::InvalidInput,
74                    format!("image format detection failed: {error}"),
75                )
76            })?
77            .decode()
78            .map_err(|error| {
79                AppError::new(
80                    ErrorCode::InvalidInput,
81                    format!("image decode failed: {error}"),
82                )
83            })?;
84
85        let resized = img.resize_exact(
86            self.width,
87            self.height,
88            image::imageops::FilterType::Lanczos3,
89        );
90
91        let mut buf = Vec::new();
92        let mut cursor = Cursor::new(&mut buf);
93        image::codecs::jpeg::JpegEncoder::new_with_quality(&mut cursor, self.quality)
94            .encode_image(&resized)
95            .map_err(|error| {
96                AppError::new(ErrorCode::Internal, format!("image encode failed: {error}"))
97            })?;
98
99        Ok(Some(
100            item.try_with_payload(DataPayload::bytes(buf, limits)?, limits)?
101                .with_extension(".jpg"),
102        ))
103    }
104}
105
106#[cfg(all(test, feature = "image-transform"))]
107mod tests {
108    use image::{ImageBuffer, ImageFormat, Rgb};
109    use std::io::Cursor;
110
111    use super::*;
112    use crate::{Label, MediaType};
113
114    fn png_bytes(width: u32, height: u32) -> Vec<u8> {
115        let image = ImageBuffer::from_pixel(width, height, Rgb([10_u8, 20, 30]));
116        let mut bytes = Vec::new();
117        image
118            .write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
119            .unwrap();
120        bytes
121    }
122
123    #[test]
124    fn resize_transform_validates_dimensions_and_quality() {
125        assert!(ResizeTransform::new(0, 10).is_err());
126        assert!(ResizeTransform::new(10, 0).is_err());
127
128        let transform = ResizeTransform::new(4, 5).unwrap();
129        assert_eq!(transform.name(), "resize-4x5");
130        assert!(transform.clone().with_quality(0).is_err());
131        assert!(transform.clone().with_quality(101).is_err());
132        assert_eq!(transform.with_quality(90).unwrap().quality, 90);
133    }
134
135    #[test]
136    fn resize_transform_reencodes_image_and_preserves_metadata() {
137        let limits = DatasetLimits::default();
138        let item = DataItem::new(png_bytes(8, 6), Label::Real, MediaType::Image, "fixture")
139            .unwrap()
140            .with_extension(".png")
141            .with_source_offset(42);
142        let transform = ResizeTransform::new(3, 2)
143            .unwrap()
144            .with_quality(80)
145            .unwrap();
146
147        let output = transform.apply(item, &limits).unwrap().unwrap();
148        let bytes = output.payload().read_bytes_bounded(&limits).unwrap();
149        let decoded = image::load_from_memory(&bytes).unwrap();
150
151        assert_eq!(decoded.width(), 3);
152        assert_eq!(decoded.height(), 2);
153        assert_eq!(output.extension, ".jpg");
154        assert_eq!(output.source_offset(), Some(42));
155        assert_eq!(output.source_name, "fixture");
156    }
157
158    #[test]
159    fn resize_transform_rejects_non_image_payloads() {
160        let item = DataItem::new(
161            b"not an image".to_vec(),
162            Label::AiGenerated,
163            MediaType::Image,
164            "bad",
165        )
166        .unwrap();
167
168        let error = ResizeTransform::new(2, 2)
169            .unwrap()
170            .apply(item, &DatasetLimits::default())
171            .unwrap_err();
172
173        assert_eq!(error.code(), ErrorCode::InvalidInput);
174    }
175}
rskit_dataset/transform.rs

rskit_dataset/
transform.rs
