firefox-main/third_party/rust/jxl/src/frame/lf_preview.rs

Source code

Revision control

Copy as Markdown

Other Tools

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use crate::{
api::{JxlColorProfile, JxlColorType, JxlDataFormat, JxlOutputBuffer, JxlPixelFormat},
error::Result,
frame::Frame,
headers::{Orientation, frame_header::FrameType},
image::{DataTypeTag, Rect},
render::{
Channels, ChannelsMut, RenderPipelineInOutStage, RenderPipelineInPlaceStage,
buffer_splitter::{BufferSplitter, SaveStageBufferInfo},
low_memory_pipeline::row_buffers::RowBuffer,
save::SaveStage,
stages::{
ConvertF32ToF16Stage, ConvertF32ToU8Stage, ConvertF32ToU16Stage, FromLinearStage,
OutputColorInfo, TransferFunction, Upsample8x, XybStage,
},
},
util::{f16, mirror},
};
impl Frame {
#[allow(clippy::too_many_arguments)]
fn render_lf_frame_rect(
&mut self,
color_type: JxlColorType,
data_format: JxlDataFormat,
rect: Rect,
upsampled_rect: Rect,
orientation: Orientation,
output_buffers: &mut [Option<JxlOutputBuffer<'_>>],
full_size: (usize, usize),
output_color_info: &OutputColorInfo,
output_tf: &TransferFunction,
) -> Result<()> {
let save_stage = SaveStage::new(
if color_type.has_alpha() {
&[0, 1, 2, 3]
} else {
&[0, 1, 2]
},
orientation,
0,
color_type,
data_format,
color_type.has_alpha(),
);
let len = rect.size.0;
let ulen = len * 8;
enum DataFormatConverter {
U8(ConvertF32ToU8Stage),
U16(ConvertF32ToU16Stage),
F16(ConvertF32ToF16Stage),
None,
}
let (converter, constant_alpha) = match data_format {
JxlDataFormat::U8 { bit_depth } => (
DataFormatConverter::U8(ConvertF32ToU8Stage::new(0, bit_depth)),
RowBuffer::new_filled(DataTypeTag::U8, ulen, &(1u8 << bit_depth).to_ne_bytes())?,
),
JxlDataFormat::U16 { bit_depth, .. } => (
DataFormatConverter::U16(ConvertF32ToU16Stage::new(0, bit_depth)),
RowBuffer::new_filled(DataTypeTag::U16, ulen, &(1u16 << bit_depth).to_ne_bytes())?,
),
JxlDataFormat::F16 { .. } => (
DataFormatConverter::F16(ConvertF32ToF16Stage::new(0)),
RowBuffer::new_filled(
DataTypeTag::F16,
ulen,
&(f16::from_f32(1.0).to_bits().to_ne_bytes()),
)?,
),
JxlDataFormat::F32 { .. } => (
DataFormatConverter::None,
RowBuffer::new_filled(DataTypeTag::F32, ulen, &1.0f32.to_ne_bytes())?,
),
};
let upsample_stage = Upsample8x::new(&self.decoder_state.file_header.transform_data, 0);
let mut upsample_state = upsample_stage.init_local_state(0)?.unwrap();
let xyb_stage = XybStage::new(0, output_color_info.clone());
let from_linear_stage = FromLinearStage::new(0, output_tf.clone());
let mut lf_rows = [
RowBuffer::new(DataTypeTag::F32, 2, 0, 0, len)?,
RowBuffer::new(DataTypeTag::F32, 2, 0, 0, len)?,
RowBuffer::new(DataTypeTag::F32, 2, 0, 0, len)?,
];
// Converted to RGB in place.
let mut upsampled_rows = [
RowBuffer::new(DataTypeTag::F32, 0, 3, 3, ulen)?,
RowBuffer::new(DataTypeTag::F32, 0, 3, 3, ulen)?,
RowBuffer::new(DataTypeTag::F32, 0, 3, 3, ulen)?,
];
let mut output_rows = [
RowBuffer::new(data_format.data_type(), 0, 0, 0, ulen)?,
RowBuffer::new(data_format.data_type(), 0, 0, 0, ulen)?,
RowBuffer::new(data_format.data_type(), 0, 0, 0, ulen)?,
];
let src = if self.header.frame_type == FrameType::RegularFrame {
self.decoder_state.lf_frames[0].as_ref().unwrap()
} else {
self.lf_frame_data.as_ref().unwrap()
};
const LF_ROW_OFFSET: usize = 8;
let x0 = rect.origin.0;
let x1 = rect.end().0;
let y0 = rect.origin.1 as isize - 2;
let y1 = rect.end().1 as isize + 2;
let lf_size = src[0].size();
for yy in y0..y1 {
let sy = mirror(yy, lf_size.1);
// Fill in input.
for c in 0..3 {
let bufy = (yy + LF_ROW_OFFSET as isize) as usize;
let row = lf_rows[c].get_row_mut::<f32>(bufy);
let srow = src[c].row(sy);
let off = RowBuffer::x0_offset::<f32>();
row[off..off + len].copy_from_slice(&srow[x0..x1]);
row[off - 1] = srow[mirror(x0 as isize - 1, lf_size.0)];
row[off - 2] = srow[mirror(x0 as isize - 2, lf_size.0)];
row[off + len] = srow[mirror(x1 as isize, lf_size.0)];
row[off + len + 1] = srow[mirror(x1 as isize + 1, lf_size.0)];
}
if yy < y0 + 4 {
continue;
}
let y = yy as usize - 2;
// Upsample.
for c in 0..3 {
let off = RowBuffer::x0_offset::<f32>() - 2;
let input_rows_refs = [
&lf_rows[c].get_row::<f32>(y + LF_ROW_OFFSET - 2)[off..],
&lf_rows[c].get_row::<f32>(y + LF_ROW_OFFSET - 1)[off..],
&lf_rows[c].get_row::<f32>(y + LF_ROW_OFFSET)[off..],
&lf_rows[c].get_row::<f32>(y + LF_ROW_OFFSET + 1)[off..],
&lf_rows[c].get_row::<f32>(y + LF_ROW_OFFSET + 2)[off..],
]
.into_iter()
.collect();
let input_channels = Channels::new(input_rows_refs, 1, 5);
let output_rows_refs =
upsampled_rows[c].get_rows_mut(y * 8..y * 8 + 8, RowBuffer::x0_offset::<f32>());
let mut output_channels = ChannelsMut::new(output_rows_refs, 1, 8);
upsample_stage.process_row_chunk(
(0, 0),
len,
&input_channels,
&mut output_channels,
Some(upsample_state.as_mut()),
);
}
// un-XYB, convert and save.
for uy in y * 8..y * 8 + 8 {
// XYB
let [x, y, b] = &mut upsampled_rows;
let off = RowBuffer::x0_offset::<f32>();
let mut rows = [
&mut x.get_row_mut(uy)[off..],
&mut y.get_row_mut(uy)[off..],
&mut b.get_row_mut(uy)[off..],
];
xyb_stage.process_row_chunk((0, 0), ulen, &mut rows, None);
from_linear_stage.process_row_chunk((0, 0), ulen, &mut rows, None);
macro_rules! convert {
($s: expr, $t: ty) => {
for c in 0..3 {
let input_rows_refs = std::iter::once(
&upsampled_rows[c].get_row(uy)[RowBuffer::x0_offset::<f32>()..],
)
.collect();
let input_channels = Channels::new(input_rows_refs, 1, 1);
let output_rows_refs = output_rows[c]
.get_rows_mut(uy..uy + 1, RowBuffer::x0_offset::<$t>());
let mut output_channels = ChannelsMut::new(output_rows_refs, 1, 1);
$s.process_row_chunk(
(0, 0),
ulen,
&input_channels,
&mut output_channels,
None,
);
}
};
}
// Convert
let save_input = match &converter {
DataFormatConverter::U8(s) => {
convert!(s, u8);
&output_rows
}
DataFormatConverter::U16(s) => {
convert!(s, u16);
&output_rows
}
DataFormatConverter::F16(s) => {
convert!(s, f16);
&output_rows
}
DataFormatConverter::None => &upsampled_rows,
};
let input_no_alpha = [&save_input[0], &save_input[1], &save_input[2]];
let input_alpha = [
&save_input[0],
&save_input[1],
&save_input[2],
&constant_alpha,
];
save_stage.save_lowmem(
if color_type.has_alpha() {
&input_alpha
} else {
&input_no_alpha
},
output_buffers,
upsampled_rect.size,
uy,
upsampled_rect.origin,
full_size,
(0, 0),
)?;
}
}
Ok(())
}
pub fn maybe_preview_lf_frame(
&mut self,
pixel_format: &JxlPixelFormat,
output_buffers: &mut [JxlOutputBuffer<'_>],
changed_regions: Option<&[Rect]>,
output_profile: &JxlColorProfile,
) -> Result<()> {
if self.header.needs_blending() {
return Ok(());
}
if !((self.header.has_lf_frame() && self.header.frame_type == FrameType::RegularFrame)
|| (self.header.frame_type == FrameType::LFFrame && self.header.lf_level == 1))
{
return Ok(());
}
let output_color_info = OutputColorInfo::from_header(&self.decoder_state.file_header)?;
let Some(output_tf) = output_profile.transfer_function().map(|tf| {
TransferFunction::from_api_tf(
tf,
output_color_info.intensity_target,
output_color_info.luminances,
)
}) else {
return Ok(());
};
if output_tf.is_linear() {
return Ok(());
}
let image_metadata = &self.decoder_state.file_header.image_metadata;
if !image_metadata.xyb_encoded || !image_metadata.extra_channel_info.is_empty() {
// We only render LF frames for XYB VarDCT images with no extra channels.
// TODO(veluca): we might want to relax this to "no alpha".
return Ok(());
}
let color_type = pixel_format.color_type;
let data_format = pixel_format.color_data_format.unwrap();
if pixel_format.color_data_format.is_none()
|| output_buffers.is_empty()
|| !matches!(
color_type,
JxlColorType::Rgb | JxlColorType::Rgba | JxlColorType::Bgr | JxlColorType::Bgra,
)
{
// We only render color data, and only to 3- or 4- channel output buffers.
return Ok(());
}
// We already have a fully-rendered frame and we are not requesting to re-render
// specific regions.
if self.decoder_state.lf_frame_was_rendered && changed_regions.is_none() {
return Ok(());
}
if changed_regions.is_none() {
self.decoder_state.lf_frame_was_rendered = true;
}
let sz = &self.decoder_state.file_header.size;
let xsize = sz.xsize() as usize;
let ysize = sz.ysize() as usize;
let mut regions_storage;
let regions = if let Some(regions) = changed_regions {
regions
} else {
regions_storage = vec![];
for i in (0..xsize.div_ceil(8)).step_by(256) {
let x0 = i;
let x1 = (i + 256).min(xsize.div_ceil(8));
regions_storage.push(Rect {
origin: (x0, 0),
size: (x1 - x0, ysize.div_ceil(8)),
});
}
&regions_storage[..]
};
let orientation = image_metadata.orientation;
let info = SaveStageBufferInfo {
downsample: (0, 0),
orientation,
byte_size: data_format.bytes_per_sample() * color_type.samples_per_pixel(),
after_extend: false,
};
let info = [Some(info)];
let mut bufs = [Some(JxlOutputBuffer::reborrow(&mut output_buffers[0]))];
let mut bufs = BufferSplitter::new(&mut bufs);
for r in regions {
let upsampled_rect = Rect {
size: (r.size.0 * 8, r.size.1 * 8),
origin: (r.origin.0 * 8, r.origin.1 * 8),
};
let upsampled_rect = upsampled_rect.clip((xsize, ysize));
let mut bufs = bufs.get_local_buffers(
&info,
upsampled_rect,
false,
(xsize, ysize),
(xsize, ysize),
(0, 0),
);
self.render_lf_frame_rect(
color_type,
data_format,
*r,
upsampled_rect,
orientation,
&mut bufs,
(xsize, ysize),
&output_color_info,
&output_tf,
)?;
}
Ok(())
}
}