Files
craft-libs/crates/exr/src/decode.rs
T
Dylan M. Taylor 99b5141700 craft-exr: deep data and Cryptomatte samples from PhotoCraft (feature decode, stacked on #992) (#1001)
* craft-exr: bounded part listing, flat decoding and writing (feature decode)

PhotoCraft's exr-crate wrapper (preview check, per-part and total size limits,
file-size plausibility, part ranking, one-part decode) and EffectCraft's
streaming channel filter, as one reader that returns a neutral model: part
headers with text attributes, and the kept channels in their stored type.

* craft-exr: deep data and Cryptomatte samples (feature decode)

PhotoCraft's deep OpenEXR reader (NONE, RLE and ZIPS deep scan lines and
tiles, every allocation admitted against the byte budget) and its flattening,
and its Cryptomatte layer listing and rank decoding, on the decode module.
2026-11-11 20:23:00 -05:00

573 lines
24 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Bounded OpenEXR reading and writing on top of the `exr` crate (feature `decode`).
//!
//! - [`parts`] lists every part from the headers alone (size, windows, channels, compression,
//! text attributes), after checking the declared sizes against [`Limits`] and the file's size.
//! - [`pick_part`] chooses the part a plain picture decode shows: the best colour channel group
//! ([`crate::channels::color_channels`]), the earliest part on a tie.
//! - [`decode_part`] decodes the channels of one flat part that a filter keeps, in their stored
//! sample type, over the part's data window. Every compression the `exr` crate reads is read,
//! scan lines and tiles alike (the largest resolution level). Other parts' chunks are skipped
//! unread, and channels the filter drops are never allocated.
//! - [`write`] writes one scan-line part.
//!
//! The `exr` crate sizes buffers from headers before it reads pixels, so everything that reaches
//! it is checked first: preview attributes against their own size, each part's samples against
//! [`Limits`], all parts together against the budget, and the samples against what a file of
//! this size can decompress to.
use std::io::Cursor;
use exr::block::UncompressedBlock;
use exr::block::chunk::TileCoordinates;
use exr::image::read::image::ReadImage;
use exr::image::read::layers::{ChannelsReader, ReadAllLayers, ReadChannels};
use exr::meta::MetaData;
use exr::meta::attribute::{AttributeValue, SampleType as ExrSample};
use exr::meta::header::Header;
use exr::prelude::{
AnyChannel, AnyChannels, Blocks, Compression as ExrCompression, Encoding, FlatSamples, Layer, LayerAttributes, LineOrder, WritableImage,
};
pub use half::f16;
use crate::channels::color_channels;
/// Why a file could not be read or written.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum Error {
/// The file is damaged or not OpenEXR.
#[error("OpenEXR: {0}")]
Malformed(String),
/// A valid file this reader does not cover (deep data in a flat read, HTJ2K, subsampled
/// channels kept by the filter).
#[error("unsupported OpenEXR: {0}")]
Unsupported(String),
/// The file declares more than [`Limits`] allow.
#[error("OpenEXR exceeds a limit: {0}")]
LimitExceeded(String),
/// The picture handed to [`write`] cannot be stored.
#[error("OpenEXR: cannot write: {0}")]
Encode(String),
}
/// Result of this module's functions.
pub type Result<T> = std::result::Result<T, Error>;
fn malformed(e: impl std::fmt::Display) -> Error {
Error::Malformed(e.to_string())
}
/// An `exr`-crate error, its `NotSupported` kind kept apart from damaged data.
fn map_exr(e: exr::error::Error) -> Error {
match e {
exr::error::Error::NotSupported(message) => Error::Unsupported(message.to_string()),
e => malformed(e),
}
}
/// Caps on what a file may ask the reader to allocate. `#[non_exhaustive]`: start from
/// [`Limits::default`] and set the fields you change.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct Limits {
/// Largest data window width and height of a part.
pub max_width: u32,
pub max_height: u32,
/// Largest data window of a part, in pixels.
pub max_pixels: u64,
/// Most bytes of samples all parts together may declare, counting 4 bytes a sample. A decode
/// allocates at most this for the channels it keeps.
pub max_bytes: u64,
/// The most a compressed byte is taken to expand to (DWA parts get 4 times as much). A file
/// declaring more samples than its size times this, plus 1 MiB, is refused before reading.
pub max_ratio: u64,
}
impl Default for Limits {
fn default() -> Self {
let max_bytes = if cfg!(target_pointer_width = "64") { 8 << 30 } else { 1 << 30 };
Limits { max_width: 1 << 18, max_height: 1 << 18, max_pixels: 1 << 30, max_bytes, max_ratio: 1 << 16 }
}
}
impl Limits {
/// Effectively unlimited (still bounded by overflow checks).
pub fn none() -> Self {
Limits { max_width: u32::MAX, max_height: u32::MAX, max_pixels: u64::MAX, max_bytes: u64::MAX, max_ratio: u64::MAX }
}
/// Checks a `width × height` buffer of `bytes_per_pixel` against the size limits.
pub(crate) fn check_bytes(&self, width: u32, height: u32, bytes_per_pixel: u64) -> Result<()> {
if width == 0 || height == 0 {
return Err(Error::Malformed(format!("zero-sized picture {width}x{height}")));
}
if width > self.max_width || height > self.max_height {
return Err(Error::LimitExceeded(format!("{width}x{height} exceeds {}x{}", self.max_width, self.max_height)));
}
let pixels = u64::from(width) * u64::from(height);
if pixels > self.max_pixels {
return Err(Error::LimitExceeded(format!("{pixels} pixels exceed {}", self.max_pixels)));
}
let bytes = pixels.saturating_mul(bytes_per_pixel);
if bytes > self.max_bytes || usize::try_from(bytes).is_err() {
return Err(Error::LimitExceeded(format!("{bytes} bytes exceed {}", self.max_bytes)));
}
Ok(())
}
}
/// How a channel's samples are stored.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SampleType {
F16,
F32,
U32,
}
impl SampleType {
pub(crate) fn of(t: ExrSample) -> Self {
match t {
ExrSample::F16 => SampleType::F16,
ExrSample::F32 => SampleType::F32,
ExrSample::U32 => SampleType::U32,
}
}
}
/// A channel as the header declares it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChannelInfo {
/// The name as stored (`diffuse.R`); the part name is not prefixed.
pub name: String,
pub sample: SampleType,
/// Horizontal and vertical sampling (1, 1 for a full-resolution channel).
pub sampling: (usize, usize),
}
/// A window: position and size.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Window {
pub x: i32,
pub y: i32,
pub width: u32,
pub height: u32,
}
/// Header facts of one part.
#[derive(Debug, Clone, PartialEq)]
pub struct PartInfo {
/// Position in the file.
pub index: usize,
/// The part's name (`name` attribute), the prefix renderers give its channels.
pub name: Option<String>,
/// The `view` attribute of a stereo part.
pub view: Option<String>,
pub data_window: Window,
pub display_window: Window,
pub deep: bool,
pub tiled: bool,
/// The compression's name as the `exr` crate spells it (`ZIP16`, `PIZ`, `DWAA(None)`, …).
pub compression: String,
pub channels: Vec<ChannelInfo>,
/// Every text attribute of the part and of the shared header (`cryptomatte/…/manifest`,
/// `owner`, …), name and value, sorted by name.
pub text_attributes: Vec<(String, String)>,
}
impl PartInfo {
/// The value of the text attribute called `name`.
pub fn text(&self, name: &str) -> Option<&str> {
self.text_attributes.iter().find(|(k, _)| k == name).map(|(_, v)| v.as_str())
}
/// How much the part looks like a colour picture: 4 RGBA, 3 RGB, 2 YA, 1 Y (or a lone
/// channel), 0 anything else.
pub fn color_rank(&self) -> usize {
color_channels(self.channels.iter().map(|c| c.name.as_str())).map_or(0, |(layout, _)| layout.channels())
}
}
/// Refuses headers whose `preview` attributes declare more pixels than their value holds.
///
/// The `exr` crate sizes a preview's buffer from its declared dimensions (up to 4 MiB, zero
/// filled) before it finds the attribute too short, and in its non-pedantic mode then skips it
/// and goes on: thousands of tiny malformed previews would each cost that allocation. This walks
/// the raw headers and checks each preview against its own size. Anything else odd is left to
/// the crate to report.
pub(crate) fn check_previews(b: &[u8]) -> Result<()> {
if b.get(..4) != Some(&[0x76, 0x2f, 0x31, 0x01][..]) {
return Ok(());
}
let Some(flags) = b.get(4..8).and_then(|v| <[u8; 4]>::try_from(v).ok()).map(u32::from_le_bytes) else { return Ok(()) };
let multipart = flags & 0x1000 != 0;
let cstr = |at: usize| -> Option<(&[u8], usize)> {
let rest = b.get(at..)?;
let n = rest.iter().position(|&c| c == 0)?;
Some((rest.get(..n)?, at + n + 1))
};
let u32_at = |at: usize| b.get(at..at.checked_add(4)?).and_then(|v| <[u8; 4]>::try_from(v).ok()).map(u32::from_le_bytes);
let mut at = 8;
loop {
// One header's attributes: name, type, size, value; an empty name ends the header.
loop {
let Some((name, next)) = cstr(at) else { return Ok(()) };
at = next;
if name.is_empty() {
break;
}
let Some((kind, next)) = cstr(at) else { return Ok(()) };
let Some(size) = u32_at(next).and_then(|n| i32::try_from(n).ok()).and_then(|n| usize::try_from(n).ok()) else { return Ok(()) };
let value = next + 4;
let end = match value.checked_add(size) {
Some(end) if end <= b.len() => end,
_ => return Ok(()),
};
if kind == b"preview" {
let declared = u32_at(value)
.zip(u32_at(value + 4))
.and_then(|(w, h)| u64::from(w).checked_mul(u64::from(h))?.checked_mul(4))
.filter(|_| size >= 8);
if declared.is_none_or(|n| n > (size - 8) as u64) {
return Err(malformed("a preview attribute declares more pixels than it holds"));
}
}
at = end;
}
// A multi-part header list ends with an empty header (a lone NUL).
if !multipart || b.get(at) == Some(&0) {
return Ok(());
}
}
}
fn window_size(size: exr::math::Vec2<usize>) -> (u32, u32) {
(u32::try_from(size.0).unwrap_or(u32::MAX), u32::try_from(size.1).unwrap_or(u32::MAX))
}
/// Reads the headers and checks every part's declared size, alone and together, before any pixel
/// allocation. Headers are read without the crate's validation, which refuses deep parts.
fn read_meta(bytes: &[u8], limits: &Limits) -> Result<MetaData> {
check_previews(bytes)?;
let meta = MetaData::read_from_buffered(Cursor::new(bytes), false).map_err(map_exr)?;
let (mut total, mut dwa) = (0u64, 0u64);
for header in meta.headers.iter() {
let (w, h) = window_size(header.layer_size);
if w > limits.max_width || h > limits.max_height {
return Err(Error::LimitExceeded(format!("a {w}x{h} part exceeds {}x{}", limits.max_width, limits.max_height)));
}
let pixels = u64::from(w) * u64::from(h);
if pixels > limits.max_pixels {
return Err(Error::LimitExceeded(format!("a part of {pixels} pixels exceeds {}", limits.max_pixels)));
}
let bytes = pixels.saturating_mul(header.channels.list.len().max(1) as u64 * 4);
total = total.saturating_add(bytes);
if matches!(header.compression, ExrCompression::DWAA(_) | ExrCompression::DWAB(_)) {
dwa = dwa.saturating_add(bytes);
}
}
if total > limits.max_bytes {
return Err(Error::LimitExceeded(format!("{total} bytes across {} parts exceed {}", meta.headers.len(), limits.max_bytes)));
}
// The crate allocates a flat part's samples before reading its chunks, so the file must be
// able to hold them. DWA's run-length channels are zlib-packed on top and reach about
// 132000:1, so DWA parts are counted at a quarter. Deep parts size their samples from the
// file's own counts.
if !meta.requirements.has_deep_data {
let declared = total.saturating_sub(dwa).saturating_add(dwa / 4);
let capacity = (bytes.len() as u64).saturating_mul(limits.max_ratio).saturating_add(1 << 20);
if declared > capacity {
return Err(Error::LimitExceeded(format!("a {} byte file cannot hold {declared} bytes of samples", bytes.len())));
}
}
Ok(meta)
}
fn text_attributes(header: &Header) -> Vec<(String, String)> {
let mut out: Vec<(String, String)> = header
.shared_attributes
.other
.iter()
.chain(header.own_attributes.other.iter())
.filter_map(|(k, v)| match v {
AttributeValue::Text(t) => Some((k.to_string(), t.to_string())),
_ => None,
})
.collect();
out.sort();
out.dedup_by(|a, b| a.0 == b.0);
out
}
fn part_info(index: usize, header: &Header) -> PartInfo {
let (w, h) = window_size(header.layer_size);
let display = &header.shared_attributes.display_window;
let (dw, dh) = window_size(display.size);
PartInfo {
index,
name: header.own_attributes.layer_name.as_ref().map(|n| n.to_string()),
view: header.own_attributes.view_name.as_ref().map(|v| v.to_string()),
data_window: Window { x: header.own_attributes.layer_position.0, y: header.own_attributes.layer_position.1, width: w, height: h },
display_window: Window { x: display.position.0, y: display.position.1, width: dw, height: dh },
deep: header.deep,
tiled: matches!(header.blocks, exr::meta::BlockDescription::Tiles(_)),
compression: format!("{:?}", header.compression),
channels: header
.channels
.list
.iter()
.map(|c| ChannelInfo { name: c.name.to_string(), sample: SampleType::of(c.sample_type), sampling: (c.sampling.0, c.sampling.1) })
.collect(),
text_attributes: text_attributes(header),
}
}
/// Lists every part of the file from its headers, after checking the declared sizes.
pub fn parts(bytes: &[u8], limits: &Limits) -> Result<Vec<PartInfo>> {
let meta = read_meta(bytes, limits)?;
Ok(meta.headers.iter().enumerate().map(|(i, h)| part_info(i, h)).collect())
}
/// The part a plain picture decode shows: the highest [`PartInfo::color_rank`], the earliest on a
/// tie. `None` when no part has colour channels.
pub fn pick_part(parts: &[PartInfo]) -> Option<usize> {
let mut best: Option<(usize, usize)> = None;
for p in parts {
let rank = p.color_rank();
if rank > 0 && best.is_none_or(|(_, r)| rank > r) {
best = Some((p.index, rank));
}
}
best.map(|(i, _)| i)
}
/// The samples of a decoded channel, in their stored type.
#[derive(Debug, Clone, PartialEq)]
pub enum Samples {
F16(Vec<f16>),
F32(Vec<f32>),
U32(Vec<u32>),
}
impl Samples {
pub fn len(&self) -> usize {
match self {
Samples::F16(v) => v.len(),
Samples::F32(v) => v.len(),
Samples::U32(v) => v.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// Sample `i` as `f32` (U32 IDs exact up to 2^24). 0 past the end.
pub fn get_f32(&self, i: usize) -> f32 {
match self {
Samples::F16(v) => v.get(i).map_or(0.0, |s| s.to_f32()),
Samples::F32(v) => v.get(i).copied().unwrap_or(0.0),
Samples::U32(v) => v.get(i).map_or(0.0, |&s| s as f32),
}
}
/// All samples as `f32`.
pub fn to_f32(&self) -> Vec<f32> {
match self {
Samples::F16(v) => v.iter().map(|s| s.to_f32()).collect(),
Samples::F32(v) => v.clone(),
Samples::U32(v) => v.iter().map(|&s| s as f32).collect(),
}
}
}
/// A decoded channel.
#[derive(Debug, Clone, PartialEq)]
pub struct Channel {
/// The name as stored.
pub name: String,
/// `width × height` samples of the part's data window, top row first.
pub samples: Samples,
}
/// A decoded flat part: its header facts and the channels the filter kept, in file order.
#[derive(Debug, Clone, PartialEq)]
pub struct Part {
pub info: PartInfo,
pub channels: Vec<Channel>,
}
impl Part {
/// The kept channel called `name`.
pub fn channel(&self, name: &str) -> Option<&Channel> {
self.channels.iter().find(|c| c.name == name)
}
}
/// Places a data-window plane in a display window: samples outside the display window are cut,
/// the display window's pixels the data window does not cover are `fill`. `None` when the plane
/// is not `data.width × data.height` samples or the display window has more than `max_pixels`
/// pixels.
pub fn to_display_window(plane: &[f32], data: Window, display: Window, fill: f32, max_pixels: u64) -> Option<Vec<f32>> {
let (w, h) = (data.width as usize, data.height as usize);
if plane.len() != w.checked_mul(h)? || u64::from(display.width) * u64::from(display.height) > max_pixels {
return None;
}
let (dw, dh) = (display.width as usize, display.height as usize);
let mut out = vec![fill; dw.checked_mul(dh)?];
let (ox, oy) = (i64::from(data.x) - i64::from(display.x), i64::from(data.y) - i64::from(display.y));
for (y, row) in plane.chunks_exact(w.max(1)).enumerate() {
let Some(ty) = usize::try_from(y as i64 + oy).ok().filter(|ty| *ty < dh) else { continue };
let skip = usize::try_from(-ox).unwrap_or(0).min(w);
let tx = usize::try_from(ox.max(0)).unwrap_or(usize::MAX);
let take = (w - skip).min(dw.saturating_sub(tx));
if take == 0 {
continue;
}
let at = ty * dw + tx;
if let (Some(dst), Some(src)) = (out.get_mut(at..at + take), row.get(skip..skip + take)) {
dst.copy_from_slice(src);
}
}
Some(out)
}
/// Which channels of which part a decode keeps.
struct Select<'k> {
index: usize,
keep: &'k dyn Fn(&str) -> bool,
seen: std::cell::Cell<usize>,
}
/// One part's kept channels while it is decoded: each line goes straight into its channel's plane.
struct SelectReader {
/// Per channel of the part: its plane, `None` when not kept.
planes: Vec<Option<Samples>>,
width: usize,
height: usize,
}
impl<'s> ReadChannels<'s> for Select<'_> {
type Reader = SelectReader;
fn create_channels_reader(&'s self, header: &Header) -> exr::error::Result<SelectReader> {
let position = self.seen.get();
self.seen.set(position.saturating_add(1));
let (width, height) = (header.layer_size.0, header.layer_size.1);
if position != self.index {
// Another part: keep nothing, so its chunks are skipped unread.
return Ok(SelectReader { planes: header.channels.list.iter().map(|_| None).collect(), width, height });
}
let n = width.checked_mul(height).ok_or(exr::error::Error::Invalid("part size".into()))?;
let mut planes = Vec::with_capacity(header.channels.list.len());
for c in header.channels.list.iter() {
if !(self.keep)(&c.name.to_string()) {
planes.push(None);
continue;
}
if c.sampling.0 != 1 || c.sampling.1 != 1 {
return Err(exr::error::Error::NotSupported("subsampled channels".into()));
}
planes.push(Some(match c.sample_type {
ExrSample::F16 => Samples::F16(vec![f16::ZERO; n]),
ExrSample::F32 => Samples::F32(vec![0.0; n]),
ExrSample::U32 => Samples::U32(vec![0; n]),
}));
}
Ok(SelectReader { planes, width, height })
}
}
impl ChannelsReader for SelectReader {
type Channels = Vec<Option<Samples>>;
fn filter_block(&self, tile: TileCoordinates) -> bool {
tile.is_largest_resolution_level() && self.planes.iter().any(Option::is_some)
}
fn read_block(&mut self, header: &Header, block: UncompressedBlock) -> exr::error::UnitResult {
for line in block.lines(&header.channels) {
let Some(Some(plane)) = self.planes.get_mut(line.location.channel) else { continue };
let (x, y, count) = (line.location.position.0, line.location.position.1, line.location.sample_count);
if y >= self.height || x.checked_add(count).is_none_or(|end| end > self.width) {
continue;
}
let at = y * self.width + x;
match plane {
Samples::F16(v) => line.read_samples_into_slice(v.get_mut(at..at + count).unwrap_or_default())?,
Samples::F32(v) => line.read_samples_into_slice(v.get_mut(at..at + count).unwrap_or_default())?,
Samples::U32(v) => line.read_samples_into_slice(v.get_mut(at..at + count).unwrap_or_default())?,
}
}
Ok(())
}
fn into_channels(self) -> Vec<Option<Samples>> {
self.planes
}
}
/// Decodes the channels of flat part `index` that `keep` accepts (by the name as stored), in
/// their stored type. A file with deep parts, or a kept channel that is subsampled, is
/// [`Error::Unsupported`].
pub fn decode_part(bytes: &[u8], index: usize, keep: &dyn Fn(&str) -> bool, limits: &Limits) -> Result<Part> {
let meta = read_meta(bytes, limits)?;
if meta.requirements.has_deep_data {
return Err(Error::Unsupported("deep data (use the deep reader)".into()));
}
let header = meta.headers.get(index).ok_or_else(|| malformed(format!("part {index} of {}", meta.headers.len())))?;
if matches!(header.compression, ExrCompression::HTJ2K32 | ExrCompression::HTJ2K256) {
return Err(Error::Unsupported("HTJ2K compression".into()));
}
let info = part_info(index, header);
let select = Select { index, keep, seen: std::cell::Cell::new(0) };
let image = ReadImage::new(ReadAllLayers { read_channels: select }, |_| {}).from_buffered(Cursor::new(bytes)).map_err(map_exr)?;
let planes = image.layer_data.into_iter().nth(index).map(|l| l.channel_data).ok_or_else(|| malformed("the part was not read"))?;
let channels = info.channels.iter().zip(planes).filter_map(|(c, p)| Some(Channel { name: c.name.clone(), samples: p? })).collect();
Ok(Part { info, channels })
}
/// Compression for [`write`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Compression {
None,
Rle,
Zip1,
#[default]
Zip16,
Piz,
}
/// Writes one scan-line part of `width × height` with these channels (each `width × height`
/// samples, top row first; the order given does not matter, the file stores them sorted).
pub fn write(width: u32, height: u32, channels: &[Channel], compression: Compression) -> Result<Vec<u8>> {
let enc = |m: &str| Error::Encode(m.into());
if width == 0 || height == 0 || channels.is_empty() {
return Err(enc("empty picture"));
}
let n = (width as usize).checked_mul(height as usize).ok_or_else(|| enc("picture too large"))?;
if channels.iter().any(|c| c.samples.len() != n) {
return Err(enc("a channel has the wrong number of samples"));
}
let mut list = Vec::with_capacity(channels.len());
for c in channels {
let name = exr::meta::attribute::Text::new_or_none(c.name.as_str()).ok_or_else(|| enc("a channel name is not valid OpenEXR text"))?;
let samples = match &c.samples {
Samples::F16(v) => FlatSamples::F16(v.clone()),
Samples::F32(v) => FlatSamples::F32(v.clone()),
Samples::U32(v) => FlatSamples::U32(v.clone()),
};
list.push(AnyChannel::new(name, samples));
}
let compression = match compression {
Compression::None => ExrCompression::Uncompressed,
Compression::Rle => ExrCompression::RLE,
Compression::Zip1 => ExrCompression::ZIP1,
Compression::Zip16 => ExrCompression::ZIP16,
Compression::Piz => ExrCompression::PIZ,
};
let encoding = Encoding { compression, blocks: Blocks::ScanLines, line_order: LineOrder::Increasing };
let layer = Layer::new((width as usize, height as usize), LayerAttributes::default(), encoding, AnyChannels::sort(list.into_iter().collect()));
let mut out = Cursor::new(Vec::new());
exr::image::Image::from_layer(layer).write().to_buffered(&mut out).map_err(|e| Error::Encode(e.to_string()))?;
Ok(out.into_inner())
}