use crate::error::{CodecError, CodecResult};
use crate::frame::{Plane, VideoFrame};
use bytes::Bytes;
use oximedia_core::PixelFormat;
use std::io::Cursor;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ImageFormat {
Png,
Jpeg,
WebP,
}
impl ImageFormat {
pub fn from_bytes(data: &[u8]) -> CodecResult<Self> {
if data.len() < 12 {
return Err(CodecError::InvalidData("Data too short".into()));
}
if data.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) {
return Ok(Self::Png);
}
if data.starts_with(&[0xFF, 0xD8, 0xFF]) {
return Ok(Self::Jpeg);
}
if data.starts_with(b"RIFF") && data.len() >= 12 && &data[8..12] == b"WEBP" {
return Ok(Self::WebP);
}
Err(CodecError::UnsupportedFeature(
"Unknown image format".into(),
))
}
#[must_use]
pub const fn extension(&self) -> &'static str {
match self {
Self::Png => "png",
Self::Jpeg => "jpg",
Self::WebP => "webp",
}
}
#[must_use]
pub const fn supports_alpha(&self) -> bool {
match self {
Self::Png | Self::WebP => true,
Self::Jpeg => false,
}
}
}
#[derive(Clone, Debug)]
pub struct EncoderConfig {
pub format: ImageFormat,
pub quality: u8,
pub lossless: bool,
}
impl EncoderConfig {
#[must_use]
pub const fn png() -> Self {
Self {
format: ImageFormat::Png,
quality: 100,
lossless: true,
}
}
#[must_use]
pub const fn webp_lossy(quality: u8) -> Self {
Self {
format: ImageFormat::WebP,
quality,
lossless: false,
}
}
#[must_use]
pub const fn webp_lossless() -> Self {
Self {
format: ImageFormat::WebP,
quality: 100,
lossless: true,
}
}
}
impl Default for EncoderConfig {
fn default() -> Self {
Self::png()
}
}
pub struct ImageDecoder {
format: ImageFormat,
data: Bytes,
}
impl ImageDecoder {
pub fn new(data: &[u8]) -> CodecResult<Self> {
let format = ImageFormat::from_bytes(data)?;
Ok(Self {
format,
data: Bytes::copy_from_slice(data),
})
}
#[must_use]
pub const fn format(&self) -> ImageFormat {
self.format
}
#[allow(clippy::too_many_lines)]
pub fn decode(&self) -> CodecResult<VideoFrame> {
match self.format {
ImageFormat::Png => self.decode_png(),
ImageFormat::Jpeg => self.decode_jpeg(),
ImageFormat::WebP => self.decode_webp(),
}
}
#[cfg(feature = "image-io")]
fn decode_png(&self) -> CodecResult<VideoFrame> {
let decoder = png::Decoder::new(Cursor::new(&self.data));
let mut reader = decoder
.read_info()
.map_err(|e| CodecError::DecoderError(format!("PNG decode error: {e}")))?;
let info = reader.info();
let width = info.width;
let height = info.height;
let color_type = info.color_type;
let buffer_size = reader.output_buffer_size().ok_or_else(|| {
CodecError::DecoderError("Cannot determine PNG output buffer size".into())
})?;
let mut buf = vec![0u8; buffer_size];
let output_info = reader
.next_frame(&mut buf)
.map_err(|e| CodecError::DecoderError(format!("PNG decode error: {e}")))?;
let (format, data) = match color_type {
png::ColorType::Rgb => {
(
PixelFormat::Rgb24,
buf[..output_info.buffer_size()].to_vec(),
)
}
png::ColorType::Rgba => {
(
PixelFormat::Rgba32,
buf[..output_info.buffer_size()].to_vec(),
)
}
png::ColorType::Grayscale => {
(
PixelFormat::Gray8,
buf[..output_info.buffer_size()].to_vec(),
)
}
png::ColorType::GrayscaleAlpha => {
let size = (width * height) as usize;
let mut rgba = Vec::with_capacity(size * 4);
for chunk in buf[..output_info.buffer_size()].chunks_exact(2) {
let gray = chunk[0];
let alpha = chunk[1];
rgba.extend_from_slice(&[gray, gray, gray, alpha]);
}
(PixelFormat::Rgba32, rgba)
}
png::ColorType::Indexed => {
return Err(CodecError::UnsupportedFeature(
"Indexed PNG not supported".into(),
))
}
};
let stride = data.len() / height as usize;
let plane = Plane {
data,
stride,
width,
height,
};
let mut frame = VideoFrame::new(format, width, height);
frame.planes = vec![plane];
Ok(frame)
}
#[cfg(not(feature = "image-io"))]
fn decode_png(&self) -> CodecResult<VideoFrame> {
Err(CodecError::UnsupportedFeature(
"PNG support not enabled".into(),
))
}
#[cfg(feature = "image-io")]
fn decode_jpeg(&self) -> CodecResult<VideoFrame> {
let mut decoder = jpeg_decoder::Decoder::new(Cursor::new(&self.data));
let pixels = decoder
.decode()
.map_err(|e| CodecError::DecoderError(format!("JPEG decode error: {e}")))?;
let info = decoder
.info()
.ok_or_else(|| CodecError::DecoderError("No JPEG info available".into()))?;
let width = u32::from(info.width);
let height = u32::from(info.height);
let (format, data) = match info.pixel_format {
jpeg_decoder::PixelFormat::RGB24 => (PixelFormat::Rgb24, pixels),
jpeg_decoder::PixelFormat::L8 => (PixelFormat::Gray8, pixels),
jpeg_decoder::PixelFormat::CMYK32 => {
let mut rgb = Vec::with_capacity((width * height * 3) as usize);
for chunk in pixels.chunks_exact(4) {
let c = f32::from(chunk[0]) / 255.0;
let m = f32::from(chunk[1]) / 255.0;
let y = f32::from(chunk[2]) / 255.0;
let k = f32::from(chunk[3]) / 255.0;
let r = ((1.0 - c) * (1.0 - k) * 255.0) as u8;
let g = ((1.0 - m) * (1.0 - k) * 255.0) as u8;
let b = ((1.0 - y) * (1.0 - k) * 255.0) as u8;
rgb.extend_from_slice(&[r, g, b]);
}
(PixelFormat::Rgb24, rgb)
}
_ => {
return Err(CodecError::UnsupportedFeature(format!(
"JPEG pixel format {:?} not supported",
info.pixel_format
)))
}
};
let stride = data.len() / height as usize;
let plane = Plane {
data,
stride,
width,
height,
};
let mut frame = VideoFrame::new(format, width, height);
frame.planes = vec![plane];
Ok(frame)
}
#[cfg(not(feature = "image-io"))]
fn decode_jpeg(&self) -> CodecResult<VideoFrame> {
Err(CodecError::UnsupportedFeature(
"JPEG support not enabled".into(),
))
}
#[cfg(feature = "image-io")]
fn decode_webp(&self) -> CodecResult<VideoFrame> {
use crate::webp::alpha::decode_alpha;
use crate::webp::riff::{WebPContainer, WebPEncoding};
use crate::webp::vp8l_decoder::Vp8lDecoder;
let container = WebPContainer::parse(&self.data)?;
let (width, height) = container.dimensions()?;
match container.encoding {
WebPEncoding::Lossless => {
let chunk = container.bitstream_chunk().ok_or_else(|| {
CodecError::DecoderError("No VP8L bitstream chunk found".into())
})?;
let mut decoder = Vp8lDecoder::new();
let decoded = decoder.decode(&chunk.data)?;
Self::decoded_image_to_frame(&decoded)
}
WebPEncoding::Lossy => {
let chunk = container.bitstream_chunk().ok_or_else(|| {
CodecError::DecoderError("No VP8 bitstream chunk found".into())
})?;
Self::decode_vp8_to_frame(&chunk.data, width, height)
}
WebPEncoding::Extended => {
let vp8l_chunk = container
.chunks
.iter()
.find(|c| c.chunk_type == crate::webp::riff::ChunkType::Vp8L);
let vp8_chunk = container
.chunks
.iter()
.find(|c| c.chunk_type == crate::webp::riff::ChunkType::Vp8);
let alpha_chunk = container.alpha_chunk();
if let Some(vp8l) = vp8l_chunk {
let mut decoder = Vp8lDecoder::new();
let decoded = decoder.decode(&vp8l.data)?;
Self::decoded_image_to_frame(&decoded)
} else if let Some(vp8) = vp8_chunk {
let mut frame = Self::decode_vp8_to_frame(&vp8.data, width, height)?;
if let Some(alph) = alpha_chunk {
let alpha_plane = decode_alpha(&alph.data, width, height)?;
if frame.format == PixelFormat::Rgb24 && !frame.planes.is_empty() {
let rgb = &frame.planes[0].data;
let mut rgba =
Vec::with_capacity((width as usize) * (height as usize) * 4);
for (i, rgb_chunk) in rgb.chunks_exact(3).enumerate() {
rgba.extend_from_slice(rgb_chunk);
let a = alpha_plane.get(i).copied().unwrap_or(255);
rgba.push(a);
}
let stride = (width as usize) * 4;
frame = VideoFrame::new(PixelFormat::Rgba32, width, height);
frame.planes = vec![Plane {
data: rgba,
stride,
width,
height,
}];
}
}
Ok(frame)
} else {
Err(CodecError::DecoderError(
"Extended WebP has no VP8 or VP8L bitstream".into(),
))
}
}
}
}
fn decoded_image_to_frame(
decoded: &crate::webp::vp8l_decoder::DecodedImage,
) -> CodecResult<VideoFrame> {
let width = decoded.width;
let height = decoded.height;
let has_alpha = decoded.has_alpha;
let mut rgba = Vec::with_capacity((width as usize) * (height as usize) * 4);
for &pixel in &decoded.pixels {
let a = ((pixel >> 24) & 0xFF) as u8;
let r = ((pixel >> 16) & 0xFF) as u8;
let g = ((pixel >> 8) & 0xFF) as u8;
let b = (pixel & 0xFF) as u8;
rgba.extend_from_slice(&[r, g, b, a]);
}
if !has_alpha {
let mut rgb = Vec::with_capacity((width as usize) * (height as usize) * 3);
for chunk in rgba.chunks_exact(4) {
rgb.extend_from_slice(&chunk[..3]);
}
let stride = rgb.len() / height as usize;
let plane = Plane {
data: rgb,
stride,
width,
height,
};
let mut frame = VideoFrame::new(PixelFormat::Rgb24, width, height);
frame.planes = vec![plane];
Ok(frame)
} else {
let stride = rgba.len() / height as usize;
let plane = Plane {
data: rgba,
stride,
width,
height,
};
let mut frame = VideoFrame::new(PixelFormat::Rgba32, width, height);
frame.planes = vec![plane];
Ok(frame)
}
}
#[cfg(feature = "vp8")]
fn decode_vp8_to_frame(data: &[u8], _width: u32, _height: u32) -> CodecResult<VideoFrame> {
use crate::traits::{DecoderConfig, VideoDecoder};
use crate::vp8::Vp8Decoder;
let config = DecoderConfig::default();
let mut decoder = Vp8Decoder::new(config)?;
decoder.send_packet(data, 0)?;
let yuv_frame = decoder
.receive_frame()?
.ok_or_else(|| CodecError::DecoderError("VP8 decoder produced no frame".into()))?;
if yuv_frame.format == PixelFormat::Yuv420p {
convert_yuv420p_to_rgb(&yuv_frame)
} else if yuv_frame.format == PixelFormat::Rgb24 {
Ok(yuv_frame)
} else {
Err(CodecError::UnsupportedFeature(format!(
"VP8 decoder produced unexpected format: {}",
yuv_frame.format
)))
}
}
#[cfg(not(feature = "vp8"))]
fn decode_vp8_to_frame(_data: &[u8], _width: u32, _height: u32) -> CodecResult<VideoFrame> {
Err(CodecError::UnsupportedFeature(
"VP8 lossy decoding requires the 'vp8' feature".into(),
))
}
#[cfg(not(feature = "image-io"))]
fn decode_webp(&self) -> CodecResult<VideoFrame> {
Err(CodecError::UnsupportedFeature(
"WebP support not enabled".into(),
))
}
}
pub struct ImageEncoder {
config: EncoderConfig,
}
impl ImageEncoder {
#[must_use]
pub const fn new(config: EncoderConfig) -> Self {
Self { config }
}
pub fn encode(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
match self.config.format {
ImageFormat::Png => self.encode_png(frame),
ImageFormat::Jpeg => Err(CodecError::UnsupportedFeature(
"JPEG encoding not supported (patent concerns)".into(),
)),
ImageFormat::WebP => self.encode_webp(frame),
}
}
#[cfg(feature = "image-io")]
#[allow(clippy::too_many_lines)]
fn encode_png(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
let mut output = Vec::new();
let mut encoder = png::Encoder::new(Cursor::new(&mut output), frame.width, frame.height);
let (color_type, bit_depth) = match frame.format {
PixelFormat::Rgb24 => (png::ColorType::Rgb, png::BitDepth::Eight),
PixelFormat::Rgba32 => (png::ColorType::Rgba, png::BitDepth::Eight),
PixelFormat::Gray8 => (png::ColorType::Grayscale, png::BitDepth::Eight),
PixelFormat::Gray16 => (png::ColorType::Grayscale, png::BitDepth::Sixteen),
_ => {
return Err(CodecError::UnsupportedFeature(format!(
"Pixel format {} not supported for PNG encoding",
frame.format
)))
}
};
encoder.set_color(color_type);
encoder.set_depth(bit_depth);
encoder.set_compression(png::Compression::default());
let mut writer = encoder
.write_header()
.map_err(|e| CodecError::Internal(format!("PNG encode error: {e}")))?;
if frame.planes.is_empty() {
return Err(CodecError::InvalidData("Frame has no planes".into()));
}
writer
.write_image_data(&frame.planes[0].data)
.map_err(|e| CodecError::Internal(format!("PNG encode error: {e}")))?;
writer
.finish()
.map_err(|e| CodecError::Internal(format!("PNG encode error: {e}")))?;
Ok(output)
}
#[cfg(not(feature = "image-io"))]
fn encode_png(&self, _frame: &VideoFrame) -> CodecResult<Vec<u8>> {
Err(CodecError::UnsupportedFeature(
"PNG support not enabled".into(),
))
}
#[cfg(feature = "image-io")]
fn encode_webp(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
let (width, height, data) = match frame.format {
PixelFormat::Rgb24 | PixelFormat::Rgba32 => {
if frame.planes.is_empty() {
return Err(CodecError::InvalidData("Frame has no planes".into()));
}
(frame.width, frame.height, &frame.planes[0].data)
}
PixelFormat::Gray8 => {
if frame.planes.is_empty() {
return Err(CodecError::InvalidData("Frame has no planes".into()));
}
let gray_data = &frame.planes[0].data;
let mut rgb = Vec::with_capacity(gray_data.len() * 3);
for &gray in gray_data.iter() {
rgb.extend_from_slice(&[gray, gray, gray]);
}
return self.encode_webp_rgb(frame.width, frame.height, &rgb, false);
}
_ => {
return Err(CodecError::UnsupportedFeature(format!(
"Pixel format {} not supported for WebP encoding",
frame.format
)))
}
};
let has_alpha = frame.format == PixelFormat::Rgba32;
self.encode_webp_rgb(width, height, data, has_alpha)
}
fn encode_webp_rgb(
&self,
width: u32,
height: u32,
data: &[u8],
has_alpha: bool,
) -> CodecResult<Vec<u8>> {
use crate::webp::alpha::encode_alpha;
use crate::webp::encoder::WebPLossyEncoder;
use crate::webp::riff::WebPWriter;
use crate::webp::vp8l_encoder::Vp8lEncoder;
if self.config.lossless {
let pixel_count = (width as usize) * (height as usize);
let mut pixels = Vec::with_capacity(pixel_count);
if has_alpha {
for chunk in data.chunks_exact(4) {
let r = u32::from(chunk[0]);
let g = u32::from(chunk[1]);
let b = u32::from(chunk[2]);
let a = u32::from(chunk[3]);
pixels.push((a << 24) | (r << 16) | (g << 8) | b);
}
} else {
for chunk in data.chunks_exact(3) {
let r = u32::from(chunk[0]);
let g = u32::from(chunk[1]);
let b = u32::from(chunk[2]);
pixels.push((0xFF << 24) | (r << 16) | (g << 8) | b);
}
}
let encoder = Vp8lEncoder::new(100);
let vp8l_data = encoder.encode(&pixels, width, height, has_alpha)?;
Ok(WebPWriter::write_lossless(&vp8l_data))
} else {
let quality = self.config.quality.clamp(0, 100);
let lossy_encoder = WebPLossyEncoder::new(quality);
if has_alpha {
let (vp8_data, alpha_data) = lossy_encoder.encode_rgba(data, width, height)?;
let alpha_chunk = encode_alpha(&alpha_data, width, height)?;
Ok(WebPWriter::write_extended(
&vp8_data,
Some(&alpha_chunk),
width,
height,
))
} else {
let vp8_data = lossy_encoder.encode_rgb(data, width, height)?;
Ok(WebPWriter::write_lossy(&vp8_data))
}
}
}
#[cfg(not(feature = "image-io"))]
fn encode_webp(&self, _frame: &VideoFrame) -> CodecResult<Vec<u8>> {
Err(CodecError::UnsupportedFeature(
"WebP support not enabled".into(),
))
}
}
#[must_use]
#[allow(clippy::cast_possible_truncation)]
#[allow(clippy::cast_sign_loss)]
pub fn rgb_to_yuv(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
let r = f32::from(r);
let g = f32::from(g);
let b = f32::from(b);
let y = 0.2126 * r + 0.7152 * g + 0.0722 * b;
let u = (b - y) / 1.8556 + 128.0;
let v = (r - y) / 1.5748 + 128.0;
(
y.clamp(0.0, 255.0) as u8,
u.clamp(0.0, 255.0) as u8,
v.clamp(0.0, 255.0) as u8,
)
}
#[must_use]
#[allow(clippy::cast_possible_truncation)]
#[allow(clippy::cast_sign_loss)]
pub fn yuv_to_rgb(y: u8, u: u8, v: u8) -> (u8, u8, u8) {
let y = f32::from(y);
let u = f32::from(u) - 128.0;
let v = f32::from(v) - 128.0;
let r = y + 1.5748 * v;
let g = y - 0.1873 * u - 0.4681 * v;
let b = y + 1.8556 * u;
(
r.clamp(0.0, 255.0) as u8,
g.clamp(0.0, 255.0) as u8,
b.clamp(0.0, 255.0) as u8,
)
}
pub fn convert_rgb_to_yuv420p(frame: &VideoFrame) -> CodecResult<VideoFrame> {
if !matches!(frame.format, PixelFormat::Rgb24 | PixelFormat::Rgba32) {
return Err(CodecError::InvalidParameter(
"Frame must be RGB24 or Rgba32".into(),
));
}
if frame.planes.is_empty() {
return Err(CodecError::InvalidData("Frame has no planes".into()));
}
let width = frame.width as usize;
let height = frame.height as usize;
let rgb_data = &frame.planes[0].data;
let bytes_per_pixel = if frame.format == PixelFormat::Rgb24 {
3
} else {
4
};
let y_size = width * height;
let uv_width = width / 2;
let uv_height = height / 2;
let uv_size = uv_width * uv_height;
let mut y_plane = vec![0u8; y_size];
let mut u_plane = vec![0u8; uv_size];
let mut v_plane = vec![0u8; uv_size];
for y in 0..height {
for x in 0..width {
let rgb_idx = (y * width + x) * bytes_per_pixel;
let r = rgb_data[rgb_idx];
let g = rgb_data[rgb_idx + 1];
let b = rgb_data[rgb_idx + 2];
let (y_val, u_val, v_val) = rgb_to_yuv(r, g, b);
y_plane[y * width + x] = y_val;
if x % 2 == 0 && y % 2 == 0 {
let uv_idx = (y / 2) * uv_width + (x / 2);
u_plane[uv_idx] = u_val;
v_plane[uv_idx] = v_val;
}
}
}
let mut yuv_frame = VideoFrame::new(PixelFormat::Yuv420p, frame.width, frame.height);
yuv_frame.planes = vec![
Plane {
data: y_plane,
stride: width,
width: frame.width,
height: frame.height,
},
Plane {
data: u_plane,
stride: uv_width,
width: frame.width / 2,
height: frame.height / 2,
},
Plane {
data: v_plane,
stride: uv_width,
width: frame.width / 2,
height: frame.height / 2,
},
];
yuv_frame.timestamp = frame.timestamp;
yuv_frame.frame_type = frame.frame_type;
yuv_frame.color_info = frame.color_info;
Ok(yuv_frame)
}
pub fn convert_yuv420p_to_rgb(frame: &VideoFrame) -> CodecResult<VideoFrame> {
if frame.format != PixelFormat::Yuv420p {
return Err(CodecError::InvalidParameter("Frame must be YUV420p".into()));
}
if frame.planes.len() != 3 {
return Err(CodecError::InvalidData("YUV420p requires 3 planes".into()));
}
let width = frame.width as usize;
let height = frame.height as usize;
let y_data = &frame.planes[0].data;
let u_data = &frame.planes[1].data;
let v_data = &frame.planes[2].data;
let rgb_size = width * height * 3;
let mut rgb_data = vec![0u8; rgb_size];
let uv_width = width / 2;
for y in 0..height {
for x in 0..width {
let y_val = y_data[y * width + x];
let uv_idx = (y / 2) * uv_width + (x / 2);
let u_val = u_data[uv_idx];
let v_val = v_data[uv_idx];
let (r, g, b) = yuv_to_rgb(y_val, u_val, v_val);
let rgb_idx = (y * width + x) * 3;
rgb_data[rgb_idx] = r;
rgb_data[rgb_idx + 1] = g;
rgb_data[rgb_idx + 2] = b;
}
}
let mut rgb_frame = VideoFrame::new(PixelFormat::Rgb24, frame.width, frame.height);
rgb_frame.planes = vec![Plane {
data: rgb_data,
stride: width * 3,
width: frame.width,
height: frame.height,
}];
rgb_frame.timestamp = frame.timestamp;
rgb_frame.frame_type = frame.frame_type;
rgb_frame.color_info = frame.color_info;
Ok(rgb_frame)
}