1use std::io::Cursor;
19use std::sync::Arc;
20
21use alpha2_model::{LinearPrimaries, ProviderSource};
22use alpha2_time::MediaTime;
23use exr::block::reader::ChunksReader as _;
24use exr::meta::attribute::{ChannelDescription, ChannelList, SampleType};
25use half::f16;
26
27use crate::api::{
28 ColorOrigin, Diagnostic, ExtensionError, Format, FramePixels, LinearColor, ProbeReport,
29 ResourceResolver, Timing, VisualFrame, VisualProvider, VisualSource,
30};
31use crate::budget::{MAX_CANVAS_PIXELS, check_canvas};
32
33const MAGIC: [u8; 4] = [0x76, 0x2f, 0x31, 0x01];
35
36pub(crate) fn is_exr(bytes: &[u8]) -> bool {
38 bytes.starts_with(&MAGIC)
39}
40
41pub struct ExrProvider;
43
44impl VisualProvider for ExrProvider {
45 fn name(&self) -> &'static str {
46 "exr"
47 }
48
49 fn probe(&self, bytes: &[u8]) -> Option<Result<ProbeReport, ExtensionError>> {
50 is_exr(bytes).then(|| decode_exr(bytes, None).map(|image| still_report(&image)))
51 }
52
53 fn open(
54 &self,
55 bytes: Arc<[u8]>,
56 source: &ProviderSource,
57 _resolver: &dyn ResourceResolver,
58 ) -> Result<Box<dyn VisualSource>, ExtensionError> {
59 let ProviderSource::Exr { primaries } = source else {
60 return Err(ExtensionError::Mismatch("not an OpenEXR source".into()));
61 };
62 if !is_exr(&bytes) {
63 return Err(ExtensionError::Mismatch("not an OpenEXR file".into()));
64 }
65 let image = decode_exr(&bytes, primaries.map(project_color))?;
66 let report = still_report(&image);
67 Ok(Box::new(ExrSource {
68 report,
69 frame: VisualFrame {
70 width: image.width,
71 height: image.height,
72 pixels: FramePixels::LinearF16(image.pixels.into()),
73 index: 0,
74 start: MediaTime::from_ticks(i64::MIN),
75 end: None,
76 diagnostics: Vec::new(),
77 },
78 }))
79 }
80}
81
82fn still_report(image: &ExrImage) -> ProbeReport {
83 ProbeReport {
84 format: Format::Exr,
85 width: image.width,
86 height: image.height,
87 timing: Timing::Static,
88 loop_count: None,
89 items: Vec::new(),
90 item: None,
91 color: Some(image.color.clone()),
92 diagnostics: image.diagnostics.clone(),
93 }
94}
95
96struct ExrSource {
98 report: ProbeReport,
99 frame: VisualFrame,
100}
101
102impl VisualSource for ExrSource {
103 fn report(&self) -> &ProbeReport {
104 &self.report
105 }
106
107 fn render(
108 &mut self,
109 _time: MediaTime,
110 _fit: (u32, u32),
111 ) -> Result<VisualFrame, ExtensionError> {
112 Ok(self.frame.clone())
113 }
114
115 fn resident_bytes(&self) -> usize {
116 self.frame.pixels.byte_len()
117 }
118}
119
120pub(crate) struct ExrImage {
122 pub(crate) width: u32,
123 pub(crate) height: u32,
124 pub(crate) pixels: Vec<u16>,
126 pub(crate) color: LinearColor,
128 pub(crate) diagnostics: Vec<Diagnostic>,
129}
130
131pub(crate) fn project_color(primaries: LinearPrimaries) -> LinearColor {
133 LinearColor {
134 primaries: Some(primaries),
135 chromaticities: named_chromaticities(primaries),
136 origin: ColorOrigin::Project,
137 }
138}
139
140pub(crate) fn decode_exr(
147 bytes: &[u8],
148 color: Option<LinearColor>,
149) -> Result<ExrImage, ExtensionError> {
150 check_channel_lists(bytes)?;
151 let reader = exr::block::read(Cursor::new(bytes), false).map_err(exr_error)?;
152 let headers = reader.headers();
153 if headers.iter().all(|header| header.deep) {
154 return Err(ExtensionError::Unsupported(
155 "deep OpenEXR data is not supported".into(),
156 ));
157 }
158 let Some(part) = headers.iter().position(|header| {
159 !header.deep
160 && channels_shown(header.channels.list.iter().map(|c| c.name.to_string())).is_some()
161 }) else {
162 let names = headers
163 .iter()
164 .flat_map(|header| {
165 header
166 .channels
167 .list
168 .iter()
169 .map(|channel| channel.name.to_string())
170 })
171 .collect::<Vec<_>>()
172 .join(", ");
173 return Err(ExtensionError::Unsupported(format!(
174 "no part has R, G, and B channels or a Y channel (it has {names})"
175 )));
176 };
177 let header = headers[part].clone();
178 let display = header.shared_attributes.display_window;
179 let (width, height) = (
180 u32::try_from(display.size.x()).unwrap_or(u32::MAX),
181 u32::try_from(display.size.y()).unwrap_or(u32::MAX),
182 );
183 let mut diagnostics = check_part(&header, width, height, part, headers.len())?;
184 let color = color.unwrap_or_else(|| {
185 let declared = file_color(header.shared_attributes.chromaticities);
186 if defines_color_space(declared.chromaticities) {
187 return declared;
188 }
189 diagnostics.push(Diagnostic::UnsupportedFeature {
192 feature: "chromaticities that define no color space".into(),
193 path: "/chromaticities".into(),
194 count: 1,
195 });
196 file_color(None)
197 });
198
199 let channels = &header.channels.list;
200 let shown = channels_shown(channels.iter().map(|channel| channel.name.to_string()))
201 .ok_or_else(|| ExtensionError::Malformed("the part's channels could not be read".into()))?;
202 let others = (0..channels.len())
203 .filter(|index| !shown.contains(&Some(*index)))
204 .collect::<Vec<_>>();
205 if let Some(first) = others.first() {
206 diagnostics.push(Diagnostic::UnsupportedFeature {
207 feature: "channels besides R, G, B, A, and Y".into(),
208 path: format!("/channels/{}", channels[*first].name),
209 count: others.len(),
210 });
211 }
212 for channel in shown.iter().flatten().map(|index| &channels[*index]) {
213 check_channel(channel)?;
214 }
215 if shown
216 .iter()
217 .flatten()
218 .any(|index| channels[*index].sample_type == SampleType::F32)
219 {
220 diagnostics.push(Diagnostic::ReducedPrecision);
221 }
222
223 let matrix = to_bt709(color.chromaticities);
224 let mut canvas = Canvas::new(&header, display, (width, height), shown, matrix);
225 reader
226 .filter_chunks(false, |_, tile, block| {
227 block.layer == part && tile.is_largest_resolution_level()
228 })
229 .map_err(exr_error)?
230 .decompress_sequential(false, |_, block| {
231 canvas.place_block(&block, &header.channels);
232 Ok(())
233 })
234 .map_err(exr_error)?;
235 let Canvas { pixels, counts, .. } = canvas;
236 if counts.non_finite > 0 {
237 diagnostics.push(Diagnostic::NonFiniteSamples {
238 count: counts.non_finite,
239 });
240 }
241 if counts.out_of_range > 0 {
242 diagnostics.push(Diagnostic::OutOfHalfRange {
243 count: counts.out_of_range,
244 });
245 }
246 Ok(ExrImage {
247 width,
248 height,
249 pixels,
250 color,
251 diagnostics,
252 })
253}
254
255fn check_channel_lists(bytes: &[u8]) -> Result<(), ExtensionError> {
268 fn text(bytes: &[u8], at: &mut usize) -> Option<Vec<u8>> {
271 let rest = bytes.get(*at..)?;
272 let end = 1 + rest
273 .get(1..)?
274 .iter()
275 .take(256)
276 .position(|byte| *byte == 0)?;
277 *at += end + 1;
278 Some(rest[..end].to_vec())
279 }
280 let Some(flags) = bytes.get(4..8) else {
281 return Ok(());
282 };
283 let multipart = flags[1] & 0x10 != 0;
284 let mut at = 8;
285 loop {
286 let mut attributes = 0;
287 loop {
288 match bytes.get(at) {
289 None => return Ok(()),
290 Some(0) => {
291 at += 1;
292 break;
293 }
294 Some(_) => {}
295 }
296 let Some(name) = text(bytes, &mut at) else {
297 return Ok(());
298 };
299 attributes += 1;
300 let Some(kind) = text(bytes, &mut at) else {
301 return Ok(());
302 };
303 let Some(size) = bytes
304 .get(at..at + 4)
305 .and_then(|size| usize::try_from(i32::from_le_bytes(size.try_into().ok()?)).ok())
306 else {
307 return Ok(());
308 };
309 if kind == b"chlist" && name != b"channels" {
310 return Err(ExtensionError::Malformed(format!(
311 "a header holds a channel list besides `channels`: `{}`",
312 String::from_utf8_lossy(&name)
313 )));
314 }
315 at = at.saturating_add(4).saturating_add(size);
316 }
317 if !multipart || attributes == 0 {
319 return Ok(());
320 }
321 }
322}
323
324fn check_channel(channel: &ChannelDescription) -> Result<(), ExtensionError> {
326 if channel.sampling.x() != 1 || channel.sampling.y() != 1 {
327 return Err(ExtensionError::Unsupported(format!(
328 "channel {} is subsampled",
329 channel.name
330 )));
331 }
332 if channel.sample_type == SampleType::U32 {
333 return Err(ExtensionError::Unsupported(format!(
334 "channel {} holds integers, not colors",
335 channel.name
336 )));
337 }
338 Ok(())
339}
340
341fn check_part(
343 header: &exr::meta::header::Header,
344 width: u32,
345 height: u32,
346 part: usize,
347 parts: usize,
348) -> Result<Vec<Diagnostic>, ExtensionError> {
349 check_canvas(width, height)?;
350 let data_pixels = u64::try_from(header.layer_size.x())
351 .unwrap_or(u64::MAX)
352 .saturating_mul(u64::try_from(header.layer_size.y()).unwrap_or(u64::MAX));
353 if data_pixels > MAX_CANVAS_PIXELS {
354 return Err(ExtensionError::Budget(format!(
355 "a {}x{} data window is larger than {MAX_CANVAS_PIXELS} pixels",
356 header.layer_size.x(),
357 header.layer_size.y()
358 )));
359 }
360 let mut diagnostics = Vec::new();
361 if parts > 1 {
362 diagnostics.push(Diagnostic::UnsupportedFeature {
363 feature: "more than one part".into(),
364 path: format!("/parts/{}", usize::from(part == 0)),
365 count: parts - 1,
366 });
367 }
368 let aspect = header.shared_attributes.pixel_aspect;
369 if (aspect - 1.0).abs() > 1e-6 {
370 diagnostics.push(Diagnostic::UnsupportedFeature {
371 feature: format!("pixel aspect ratio {aspect}"),
372 path: "/pixelAspectRatio".into(),
373 count: 1,
374 });
375 }
376 Ok(diagnostics)
377}
378
379fn place(offset: i64, at: usize, size: usize) -> Option<usize> {
382 usize::try_from(offset + i64::try_from(at).ok()?)
383 .ok()
384 .filter(|position| *position < size)
385}
386
387struct Canvas {
391 width: usize,
392 height: usize,
393 offset: (i64, i64),
396 shown: [Option<usize>; 4],
398 matrix: Option<[[f32; 3]; 3]>,
399 row: Vec<Vec<f32>>,
401 pixels: Vec<u16>,
402 counts: Counts,
403}
404
405impl Canvas {
406 fn new(
407 header: &exr::meta::header::Header,
408 display: exr::meta::attribute::IntegerBounds,
409 (width, height): (u32, u32),
410 shown: [Option<usize>; 4],
411 matrix: Option<[[f32; 3]; 3]>,
412 ) -> Self {
413 let position = header.own_attributes.layer_position;
414 let mut canvas = Self {
415 width: width as usize,
416 height: height as usize,
417 offset: (
418 i64::from(position.x()) - i64::from(display.position.x()),
419 i64::from(position.y()) - i64::from(display.position.y()),
420 ),
421 shown,
422 matrix,
423 row: vec![Vec::new(); header.channels.list.len()],
424 pixels: vec![0; width as usize * height as usize * 4],
425 counts: Counts::default(),
426 };
427 let blank = canvas.pixel(|_| Some(0.0));
430 let size = header.layer_size;
431 for data_y in 0..size.y() {
432 let Some(y) = place(canvas.offset.1, data_y, canvas.height) else {
433 continue;
434 };
435 for data_x in 0..size.x() {
436 if let Some(x) = place(canvas.offset.0, data_x, canvas.width) {
437 canvas.pixels[(y * canvas.width + x) * 4..][..4].copy_from_slice(&blank);
438 }
439 }
440 }
441 canvas
442 }
443
444 fn pixel(&mut self, mut sample: impl FnMut(usize) -> Option<f32>) -> [u16; 4] {
447 let shown = self.shown;
448 let counts = &mut self.counts;
449 let mut value = |slot: usize, missing: f32| {
450 shown[slot].and_then(&mut sample).map_or(missing, |value| {
451 if value.is_finite() {
452 value
453 } else {
454 counts.non_finite += 1;
455 0.0
456 }
457 })
458 };
459 let alpha = value(3, 1.0).clamp(0.0, 1.0);
460 let mut rgb = [value(0, 0.0), value(1, 0.0), value(2, 0.0)];
461 if let Some(matrix) = &self.matrix {
462 rgb = multiply(matrix, rgb);
463 }
464 [rgb[0], rgb[1], rgb[2], alpha].map(|value| to_half(value, counts))
465 }
466
467 fn place_block(&mut self, block: &exr::block::UncompressedBlock, channels: &ChannelList) {
470 let last = channels.list.len() - 1;
471 for line in block.lines(channels) {
472 let channel = line.location.channel;
473 if self.shown.contains(&Some(channel)) {
474 let row = &mut self.row[channel];
475 row.clear();
476 match channels.list[channel].sample_type {
477 SampleType::F16 => row.extend(
478 line.read_samples::<f16>()
479 .map(|sample| sample.map_or(f32::NAN, f16::to_f32)),
480 ),
481 _ => row.extend(
482 line.read_samples::<f32>()
483 .map(|sample| sample.unwrap_or(f32::NAN)),
484 ),
485 }
486 }
487 if channel != last {
488 continue;
489 }
490 let Some(y) = place(self.offset.1, line.location.position.y(), self.height) else {
491 continue;
492 };
493 for at in 0..line.location.sample_count {
494 let data_x = line.location.position.x() + at;
495 let Some(x) = place(self.offset.0, data_x, self.width) else {
496 continue;
497 };
498 let row = std::mem::take(&mut self.row);
499 let pixel = self.pixel(|channel| row[channel].get(at).copied());
500 self.row = row;
501 self.pixels[(y * self.width + x) * 4..][..4].copy_from_slice(&pixel);
502 }
503 }
504 }
505}
506
507#[derive(Default)]
509struct Counts {
510 non_finite: usize,
511 out_of_range: usize,
512}
513
514fn to_half(value: f32, counts: &mut Counts) -> u16 {
518 if value.is_nan() {
519 counts.non_finite += 1;
520 return 0;
521 }
522 let half = f16::from_f32(value);
523 if half.is_infinite() {
524 counts.out_of_range += 1;
525 return if value > 0.0 { f16::MAX } else { f16::MIN }.to_bits();
526 }
527 half.to_bits()
528}
529
530fn channels_shown(names: impl Iterator<Item = String>) -> Option<[Option<usize>; 4]> {
534 let names = names.collect::<Vec<_>>();
535 let find = |name: &str| names.iter().position(|candidate| candidate == name);
536 let alpha = find("A");
537 match (find("R"), find("G"), find("B"), find("Y")) {
538 (Some(red), Some(green), Some(blue), _) => {
539 Some([Some(red), Some(green), Some(blue), alpha])
540 }
541 (_, _, _, Some(luma)) => Some([Some(luma), Some(luma), Some(luma), alpha]),
542 _ => None,
543 }
544}
545
546fn exr_error(error: exr::error::Error) -> ExtensionError {
547 match error {
548 exr::error::Error::NotSupported(message) => ExtensionError::Unsupported(message.into()),
549 other => ExtensionError::Malformed(other.to_string()),
550 }
551}
552
553fn file_color(declared: Option<exr::meta::attribute::Chromaticities>) -> LinearColor {
556 let Some(declared) = declared else {
557 return LinearColor {
558 primaries: Some(LinearPrimaries::Bt709),
559 chromaticities: named_chromaticities(LinearPrimaries::Bt709),
560 origin: ColorOrigin::FormatDefault,
561 };
562 };
563 let chromaticities = [declared.red, declared.green, declared.blue, declared.white]
564 .map(|point| [point.x(), point.y()]);
565 LinearColor {
566 primaries: PRIMARIES.into_iter().find(|named| {
567 named_chromaticities(*named)
568 .iter()
569 .flatten()
570 .zip(chromaticities.iter().flatten())
571 .all(|(a, b)| (a - b).abs() < 5e-4)
572 }),
573 chromaticities,
574 origin: ColorOrigin::File,
575 }
576}
577
578const PRIMARIES: [LinearPrimaries; 5] = [
579 LinearPrimaries::Bt709,
580 LinearPrimaries::Bt2020,
581 LinearPrimaries::P3D65,
582 LinearPrimaries::AcesAp0,
583 LinearPrimaries::AcesAp1,
584];
585
586const D65: [f32; 2] = [0.3127, 0.3290];
587const ACES_WHITE: [f32; 2] = [0.32168, 0.33767];
588
589pub(crate) const fn named_chromaticities(primaries: LinearPrimaries) -> [[f32; 2]; 4] {
591 match primaries {
592 LinearPrimaries::Bt709 => [[0.64, 0.33], [0.30, 0.60], [0.15, 0.06], D65],
593 LinearPrimaries::Bt2020 => [[0.708, 0.292], [0.170, 0.797], [0.131, 0.046], D65],
594 LinearPrimaries::P3D65 => [[0.680, 0.320], [0.265, 0.690], [0.150, 0.060], D65],
595 LinearPrimaries::AcesAp0 => [[0.7347, 0.2653], [0.0, 1.0], [0.0001, -0.0770], ACES_WHITE],
596 LinearPrimaries::AcesAp1 => [[0.713, 0.293], [0.165, 0.830], [0.128, 0.044], ACES_WHITE],
597 }
598}
599
600type Matrix = [[f64; 3]; 3];
601
602fn to_bt709(chromaticities: [[f32; 2]; 4]) -> Option<[[f32; 3]; 3]> {
605 let bt709 = named_chromaticities(LinearPrimaries::Bt709);
606 if chromaticities == bt709 {
607 return None;
608 }
609 let to_xyz = rgb_to_xyz(chromaticities);
610 let adapt = bradford(chromaticities[3], D65);
611 let from_xyz = invert(&rgb_to_xyz(bt709));
612 let matrix = product(&from_xyz, &product(&adapt, &to_xyz));
613 #[allow(clippy::cast_possible_truncation)]
615 Some(matrix.map(|row| row.map(|value| value as f32)))
616}
617
618fn defines_color_space(chromaticities: [[f32; 2]; 4]) -> bool {
621 to_bt709(chromaticities)
622 .is_none_or(|matrix| matrix.iter().flatten().all(|value| value.is_finite()))
623}
624
625fn multiply(matrix: &[[f32; 3]; 3], rgb: [f32; 3]) -> [f32; 3] {
626 matrix.map(|row| row[0] * rgb[0] + row[1] * rgb[1] + row[2] * rgb[2])
627}
628
629fn xyz_of(point: [f32; 2]) -> [f64; 3] {
630 let (x, y) = (f64::from(point[0]), f64::from(point[1]));
631 [x / y, 1.0, (1.0 - x - y) / y]
632}
633
634fn rgb_to_xyz(chromaticities: [[f32; 2]; 4]) -> Matrix {
636 let columns = [0, 1, 2].map(|index| xyz_of(chromaticities[index]));
637 let primaries = [0, 1, 2].map(|row| columns.map(|column| column[row]));
638 let white = xyz_of(chromaticities[3]);
639 let scale = apply(&invert(&primaries), white);
640 primaries.map(|row| [row[0] * scale[0], row[1] * scale[1], row[2] * scale[2]])
641}
642
643fn bradford(from: [f32; 2], to: [f32; 2]) -> Matrix {
645 const CONE: Matrix = [
646 [0.8951, 0.2664, -0.1614],
647 [-0.7502, 1.7135, 0.0367],
648 [0.0389, -0.0685, 1.0296],
649 ];
650 let source = apply(&CONE, xyz_of(from));
651 let target = apply(&CONE, xyz_of(to));
652 let scale = [0, 1, 2].map(|index| {
653 let mut row = [0.0; 3];
654 row[index] = target[index] / source[index];
655 row
656 });
657 product(&invert(&CONE), &product(&scale, &CONE))
658}
659
660fn apply(matrix: &Matrix, vector: [f64; 3]) -> [f64; 3] {
661 matrix.map(|row| row[0] * vector[0] + row[1] * vector[1] + row[2] * vector[2])
662}
663
664fn product(a: &Matrix, b: &Matrix) -> Matrix {
665 let mut out = [[0.0; 3]; 3];
666 for (row, out_row) in out.iter_mut().enumerate() {
667 for (column, value) in out_row.iter_mut().enumerate() {
668 *value = (0..3).map(|k| a[row][k] * b[k][column]).sum();
669 }
670 }
671 out
672}
673
674fn invert(m: &Matrix) -> Matrix {
675 let cofactor =
676 |r0: usize, r1: usize, c0: usize, c1: usize| m[r0][c0] * m[r1][c1] - m[r0][c1] * m[r1][c0];
677 let adjugate = [
678 [
679 cofactor(1, 2, 1, 2),
680 -cofactor(0, 2, 1, 2),
681 cofactor(0, 1, 1, 2),
682 ],
683 [
684 -cofactor(1, 2, 0, 2),
685 cofactor(0, 2, 0, 2),
686 -cofactor(0, 1, 0, 2),
687 ],
688 [
689 cofactor(1, 2, 0, 1),
690 -cofactor(0, 2, 0, 1),
691 cofactor(0, 1, 0, 1),
692 ],
693 ];
694 let determinant =
695 m[0][0] * adjugate[0][0] + m[0][1] * adjugate[1][0] + m[0][2] * adjugate[2][0];
696 adjugate.map(|row| row.map(|value| value / determinant))
697}
698
699#[cfg(test)]
700pub(crate) fn matrix_to_bt709(primaries: LinearPrimaries) -> Option<[[f32; 3]; 3]> {
701 to_bt709(named_chromaticities(primaries))
702}