exr/image/write/
channels.rs

1//! How to read arbitrary channels and rgb channels.
2
3use crate::prelude::*;
4use crate::io::*;
5use crate::math::*;
6use crate::meta::{header::*, attribute::*};
7use crate::block::*;
8use crate::image::recursive::*;
9use crate::block::samples::*;
10use crate::image::write::samples::*;
11
12use std::marker::PhantomData;
13
14
15/// Enables an image containing this list of channels to be written to a file.
16pub trait WritableChannels<'slf> {
17
18    /// Generate the file meta data for this list of channel
19    fn infer_channel_list(&self) -> ChannelList;
20
21    ///  Generate the file meta data of whether and how resolution levels should be stored in the file
22    fn infer_level_modes(&self) -> (LevelMode, RoundingMode);
23
24    /// The type of temporary writer
25    type Writer: ChannelsWriter;
26
27    /// Create a temporary writer for this list of channels
28    fn create_writer(&'slf self, header: &Header) -> Self::Writer;
29}
30
31/// A temporary writer for a list of channels
32pub trait ChannelsWriter: Sync {
33
34    /// Deliver a block of pixels, containing all channel data, to be stored in the file
35    fn extract_uncompressed_block(&self, header: &Header, block: BlockIndex) -> Vec<u8>; // TODO return uncompressed block?
36}
37
38
39/// Define how to get a pixel from your custom pixel storage.
40/// Can be a closure of type [`Sync + Fn(Vec2<usize>) -> YourPixel`].
41pub trait GetPixel: Sync {
42
43    /// The pixel tuple containing `f32`, `f16`, `u32` and `Sample` values.
44    /// The length of the tuple must match the number of channels in the image.
45    type Pixel;
46
47    /// Inspect a single pixel at the requested position.
48    /// Will be called exactly once for each pixel in the image.
49    /// The position will not exceed the image dimensions.
50    /// Might be called from multiple threads at the same time.
51    fn get_pixel(&self, position: Vec2<usize>) -> Self::Pixel;
52}
53
54impl<F, P> GetPixel for F where F: Sync + Fn(Vec2<usize>) -> P {
55    type Pixel = P;
56    fn get_pixel(&self, position: Vec2<usize>) -> P { self(position) }
57}
58
59impl<'samples, Samples> WritableChannels<'samples> for AnyChannels<Samples>
60    where Samples: 'samples + WritableSamples<'samples>
61{
62    fn infer_channel_list(&self) -> ChannelList {
63        ChannelList::new(self.list.iter().map(|channel| ChannelDescription {
64            name: channel.name.clone(),
65            sample_type: channel.sample_data.sample_type(),
66            quantize_linearly: channel.quantize_linearly,
67            sampling: channel.sampling
68        }).collect())
69    }
70
71    fn infer_level_modes(&self) -> (LevelMode, RoundingMode) {
72        let mode = self.list.iter().next().expect("zero channels in list").sample_data.infer_level_modes();
73
74        debug_assert!(
75            std::iter::repeat(mode).zip(self.list.iter().skip(1))
76                .all(|(first, other)| other.sample_data.infer_level_modes() == first),
77
78            "level mode must be the same across all levels (do not nest resolution levels!)"
79        );
80
81        mode
82    }
83
84    type Writer = AnyChannelsWriter<Samples::Writer>;
85    fn create_writer(&'samples self, header: &Header) -> Self::Writer {
86        let channels = self.list.iter()
87            .map(|chan| chan.sample_data.create_samples_writer(header))
88            .collect();
89
90        AnyChannelsWriter { channels }
91    }
92}
93
94/// A temporary writer for an arbitrary list of channels
95#[derive(Debug, Clone, Eq, PartialEq)]
96pub struct AnyChannelsWriter<SamplesWriter> {
97    channels: SmallVec<[SamplesWriter; 4]>
98}
99
100impl<Samples> ChannelsWriter for AnyChannelsWriter<Samples> where Samples: SamplesWriter {
101    fn extract_uncompressed_block(&self, header: &Header, block_index: BlockIndex) -> Vec<u8> {
102        UncompressedBlock::collect_block_data_from_lines(&header.channels, block_index, |line_ref| {
103            self.channels[line_ref.location.channel].extract_line(line_ref)
104        })
105    }
106}
107
108
109
110
111
112
113impl<'c, Channels, Storage>
114WritableChannels<'c> for SpecificChannels<Storage, Channels>
115where
116    Storage: 'c + GetPixel,
117    Storage::Pixel: IntoRecursive,
118    Channels: 'c + Sync + Clone + IntoRecursive,
119    <Channels as IntoRecursive>::Recursive: WritableChannelsDescription<<Storage::Pixel as IntoRecursive>::Recursive>,
120{
121    fn infer_channel_list(&self) -> ChannelList {
122        let mut vec = self.channels.clone().into_recursive().channel_descriptions_list();
123        vec.sort_unstable_by_key(|channel:&ChannelDescription| channel.name.clone()); // TODO no clone?
124
125        debug_assert!(
126            // check for equal neighbors in sorted vec
127            vec.iter().zip(vec.iter().skip(1)).all(|(prev, next)| prev.name != next.name),
128            "specific channels contain duplicate channel names"
129        );
130
131        ChannelList::new(vec)
132    }
133
134    fn infer_level_modes(&self) -> (LevelMode, RoundingMode) {
135        (LevelMode::Singular, RoundingMode::Down) // TODO
136    }
137
138    type Writer = SpecificChannelsWriter<
139        'c,
140        <<Channels as IntoRecursive>::Recursive as WritableChannelsDescription<<Storage::Pixel as IntoRecursive>::Recursive>>::RecursiveWriter,
141        Storage,
142        Channels
143    >;
144
145    fn create_writer(&'c self, header: &Header) -> Self::Writer {
146        SpecificChannelsWriter {
147            channels: self,
148            recursive_channel_writer: self.channels.clone().into_recursive().create_recursive_writer(&header.channels),
149        }
150    }
151}
152
153
154
155/// A temporary writer for a layer of channels, alpha being optional
156#[derive(Debug, Clone, Eq, PartialEq)]
157pub struct SpecificChannelsWriter<'channels, PixelWriter, Storage, Channels> {
158    channels: &'channels SpecificChannels<Storage, Channels>, // TODO this need not be a reference?? impl writer for specific_channels directly?
159    recursive_channel_writer: PixelWriter,
160}
161
162
163impl<'channels, PxWriter, Storage, Channels> ChannelsWriter
164for SpecificChannelsWriter<'channels, PxWriter, Storage, Channels>
165    where
166        Channels: Sync,
167        Storage: GetPixel,
168        Storage::Pixel: IntoRecursive,
169        PxWriter: Sync + RecursivePixelWriter<<Storage::Pixel as IntoRecursive>::Recursive>,
170{
171    fn extract_uncompressed_block(&self, header: &Header, block_index: BlockIndex) -> Vec<u8> {
172        let block_bytes = block_index.pixel_size.area() * header.channels.bytes_per_pixel;
173        let mut block_bytes = vec![0_u8; block_bytes];
174
175        let width = block_index.pixel_size.0;
176        let line_bytes = width * header.channels.bytes_per_pixel;
177        let byte_lines = block_bytes.chunks_exact_mut(line_bytes);
178        assert_eq!(byte_lines.len(), block_index.pixel_size.height(), "invalid block line splits");
179
180        //dbg!(width, line_bytes, header.channels.bytes_per_pixel, byte_lines.len());
181
182        let mut pixel_line = Vec::with_capacity(width);
183
184        for (y, line_bytes) in byte_lines.enumerate() {
185            pixel_line.clear();
186            pixel_line.extend((0 .. width).map(|x|
187                self.channels.pixels.get_pixel(block_index.pixel_position + Vec2(x, y)).into_recursive()
188            ));
189
190            self.recursive_channel_writer.write_pixels(line_bytes, pixel_line.as_slice(), |px| px);
191        }
192
193        block_bytes
194    }
195}
196
197/// A tuple containing either `ChannelsDescription` or `Option<ChannelsDescription>` entries.
198/// Use an `Option` if you want to dynamically omit a single channel (probably only for roundtrip tests).
199/// The number of entries must match the number of channels.
200pub trait WritableChannelsDescription<Pixel>: Sync {
201
202    /// A type that has a recursive entry for each channel in the image,
203    /// which must accept the desired pixel type.
204    type RecursiveWriter: RecursivePixelWriter<Pixel>;
205
206    /// Create the temporary writer, accepting the sorted list of channels from `channel_descriptions_list`.
207    fn create_recursive_writer(&self, channels: &ChannelList) -> Self::RecursiveWriter;
208
209    /// Return all the channels that should actually end up in the image, in any order.
210    fn channel_descriptions_list(&self) -> SmallVec<[ChannelDescription; 5]>;
211}
212
213impl WritableChannelsDescription<NoneMore> for NoneMore {
214    type RecursiveWriter = NoneMore;
215    fn create_recursive_writer(&self, _: &ChannelList) -> Self::RecursiveWriter { NoneMore }
216    fn channel_descriptions_list(&self) -> SmallVec<[ChannelDescription; 5]> { SmallVec::new() }
217}
218
219impl<InnerDescriptions, InnerPixel, Sample: IntoNativeSample>
220    WritableChannelsDescription<Recursive<InnerPixel, Sample>>
221    for Recursive<InnerDescriptions, ChannelDescription>
222    where InnerDescriptions: WritableChannelsDescription<InnerPixel>
223{
224    type RecursiveWriter = RecursiveWriter<InnerDescriptions::RecursiveWriter, Sample>;
225
226    fn create_recursive_writer(&self, channels: &ChannelList) -> Self::RecursiveWriter {
227        // this linear lookup is required because the order of the channels changed, due to alphabetical sorting
228        let (start_byte_offset, target_sample_type) = channels.channels_with_byte_offset()
229            .find(|(_offset, channel)| channel.name == self.value.name)
230            .map(|(offset, channel)| (offset, channel.sample_type))
231            .expect("a channel has not been put into channel list");
232
233        Recursive::new(self.inner.create_recursive_writer(channels), SampleWriter {
234            start_byte_offset, target_sample_type,
235            px: PhantomData::default()
236        })
237    }
238
239    fn channel_descriptions_list(&self) -> SmallVec<[ChannelDescription; 5]> {
240        let mut inner_list = self.inner.channel_descriptions_list();
241        inner_list.push(self.value.clone());
242        inner_list
243    }
244}
245
246impl<InnerDescriptions, InnerPixel, Sample: IntoNativeSample>
247WritableChannelsDescription<Recursive<InnerPixel, Sample>>
248for Recursive<InnerDescriptions, Option<ChannelDescription>>
249    where InnerDescriptions: WritableChannelsDescription<InnerPixel>
250{
251    type RecursiveWriter = OptionalRecursiveWriter<InnerDescriptions::RecursiveWriter, Sample>;
252
253    fn create_recursive_writer(&self, channels: &ChannelList) -> Self::RecursiveWriter {
254        // this linear lookup is required because the order of the channels changed, due to alphabetical sorting
255
256        let channel = self.value.as_ref().map(|required_channel|
257            channels.channels_with_byte_offset()
258                .find(|(_offset, channel)| channel == &required_channel)
259                .map(|(offset, channel)| (offset, channel.sample_type))
260                .expect("a channel has not been put into channel list")
261        );
262
263        Recursive::new(
264            self.inner.create_recursive_writer(channels),
265            channel.map(|(start_byte_offset, target_sample_type)| SampleWriter {
266                start_byte_offset, target_sample_type,
267                px: PhantomData::default(),
268            })
269        )
270    }
271
272    fn channel_descriptions_list(&self) -> SmallVec<[ChannelDescription; 5]> {
273        let mut inner_list = self.inner.channel_descriptions_list();
274        if let Some(value) = &self.value { inner_list.push(value.clone()); }
275        inner_list
276    }
277}
278
279/// Write pixels to a slice of bytes. The top level writer contains all the other channels,
280/// the most inner channel is `NoneMore`.
281pub trait RecursivePixelWriter<Pixel>: Sync {
282
283    /// Write pixels to a slice of bytes. Recursively do this for all channels.
284    fn write_pixels<FullPixel>(&self, bytes: &mut [u8], pixels: &[FullPixel], get_pixel: impl Fn(&FullPixel) -> &Pixel);
285}
286
287type RecursiveWriter<Inner, Sample> = Recursive<Inner, SampleWriter<Sample>>;
288type OptionalRecursiveWriter<Inner, Sample> = Recursive<Inner, Option<SampleWriter<Sample>>>;
289
290/// Write the pixels of a single channel, unconditionally. Generic over the concrete sample type (f16, f32, u32).
291#[derive(Debug, Clone)]
292pub struct SampleWriter<Sample> {
293    target_sample_type: SampleType,
294    start_byte_offset: usize,
295    px: PhantomData<Sample>,
296}
297
298impl<Sample> SampleWriter<Sample> where Sample: IntoNativeSample {
299    fn write_own_samples(&self, bytes: &mut [u8], samples: impl ExactSizeIterator<Item=Sample>) {
300        let byte_start_index = samples.len() * self.start_byte_offset;
301        let byte_count = samples.len() * self.target_sample_type.bytes_per_sample();
302        let ref mut byte_writer = &mut bytes[byte_start_index..byte_start_index + byte_count];
303
304        let write_error_msg = "invalid memory buffer length when writing";
305
306        // match outside the loop to avoid matching on every single sample
307        match self.target_sample_type {
308            // TODO does this boil down to a `memcpy` where the sample type equals the type parameter?
309            SampleType::F16 => for sample in samples { sample.to_f16().write(byte_writer).expect(write_error_msg); },
310            SampleType::F32 => for sample in samples { sample.to_f32().write(byte_writer).expect(write_error_msg); },
311            SampleType::U32 => for sample in samples { sample.to_u32().write(byte_writer).expect(write_error_msg); },
312        };
313
314        debug_assert!(byte_writer.is_empty(), "all samples are written, but more were expected");
315    }
316}
317
318impl RecursivePixelWriter<NoneMore> for NoneMore {
319    fn write_pixels<FullPixel>(&self, _: &mut [u8], _: &[FullPixel], _: impl Fn(&FullPixel) -> &NoneMore) {}
320}
321
322impl<Inner, InnerPixel, Sample: IntoNativeSample>
323    RecursivePixelWriter<Recursive<InnerPixel, Sample>>
324    for RecursiveWriter<Inner, Sample>
325    where Inner: RecursivePixelWriter<InnerPixel>
326{
327    // TODO impl exact size iterator <item = Self::Pixel>
328    fn write_pixels<FullPixel>(&self, bytes: &mut [u8], pixels: &[FullPixel], get_pixel: impl Fn(&FullPixel) -> &Recursive<InnerPixel, Sample>){
329        self.value.write_own_samples(bytes, pixels.iter().map(|px| get_pixel(px).value));
330        self.inner.write_pixels(bytes, pixels, |px| &get_pixel(px).inner);
331    }
332}
333
334impl<Inner, InnerPixel, Sample> RecursivePixelWriter<Recursive<InnerPixel, Sample>>
335    for OptionalRecursiveWriter<Inner, Sample>
336    where Inner: RecursivePixelWriter<InnerPixel>,
337        Sample: IntoNativeSample
338{
339    fn write_pixels<FullPixel>(&self, bytes: &mut [u8], pixels: &[FullPixel], get_pixel: impl Fn(&FullPixel) -> &Recursive<InnerPixel, Sample>) {
340        if let Some(writer) = &self.value {
341            writer.write_own_samples(bytes, pixels.iter().map(|px| get_pixel(px).value));
342        }
343
344        self.inner.write_pixels(bytes, pixels, |px| &get_pixel(px).inner);
345    }
346}
347
348
349
350
351
352
353
354#[cfg(test)]
355pub mod test {
356    use crate::image::write::channels::WritableChannels;
357    use crate::image::SpecificChannels;
358    use crate::prelude::{f16};
359    use crate::meta::attribute::{ChannelDescription, SampleType};
360    use crate::image::pixel_vec::PixelVec;
361
362    #[test]
363    fn compiles(){
364        let x = 3_f32;
365        let y = f16::from_f32(4.0);
366        let z = 2_u32;
367        let s = 1.3_f32;
368        let px = (x,y,z,s);
369
370        assert_is_writable_channels(
371            SpecificChannels::rgba(|_pos| px)
372        );
373
374        assert_is_writable_channels(SpecificChannels::rgba(
375            PixelVec::new((3, 2), vec![px, px, px, px, px, px])
376        ));
377
378        let px = (2333_u32, 4_f32);
379        assert_is_writable_channels(
380            SpecificChannels::build()
381                .with_channel("A")
382                .with_channel("C")
383                .with_pixels(PixelVec::new((3, 2), vec![px, px, px, px, px, px]))
384        );
385
386        let px = (3_f32, f16::ONE, 2333_u32, 4_f32);
387        assert_is_writable_channels(SpecificChannels::new(
388            (
389                ChannelDescription::named("x", SampleType::F32),
390                ChannelDescription::named("y", SampleType::F16),
391                Some(ChannelDescription::named("z", SampleType::U32)),
392                Some(ChannelDescription::named("p", SampleType::F32)),
393            ),
394
395            PixelVec::new((3, 2), vec![px, px, px, px, px, px])
396        ));
397
398
399
400        fn assert_is_writable_channels<'s>(_channels: impl WritableChannels<'s>){}
401
402    }
403}
404
405
406
407
exr/image/write/channels.rs

exr/image/write/
channels.rs
