Skip to content

Migrate legacy tensorflow image png decoder to memory-safe pixel_bridge #46

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Migrate legacy tensorflow image png decoder to memory-safe pixel_bridge #46

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions pixel_bridge/rust/image.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,14 @@ use crate::{
reader::{Format, ReadSeek},
vec_u8::VecU8,
};
use image::codecs::bmp::BmpDecoder as RustBmpDecoder;
use image::codecs::gif::GifDecoder as RustGifDecoder;
use image::codecs::ico::IcoDecoder as RustIcoDecoder;
use image::codecs::jpeg::JpegDecoder as RustJpegDecoder;
use image::codecs::tiff::TiffDecoder as RustTiffDecoder;
use image::codecs::webp::WebPDecoder as RustWebPDecoder;
use image::{
codecs::bmp::BmpDecoder as RustBmpDecoder, codecs::gif::GifDecoder as RustGifDecoder,
codecs::ico::IcoDecoder as RustIcoDecoder, codecs::jpeg::JpegDecoder as RustJpegDecoder,
codecs::png::PngDecoder as RustPngDecoder, codecs::tiff::TiffDecoder as RustTiffDecoder,
codecs::webp::WebPDecoder as RustWebPDecoder, flat::SampleLayout as RustSampleLayout,
codecs::png::PngDecoder as RustPngDecoder, flat::SampleLayout as RustSampleLayout,
AnimationDecoder, ColorType as RustColorType, ExtendedColorType as RustExtendedColorType,
Frame as RustFrame, ImageBuffer as RustImageBuffer, ImageDecoder as RustImageDecoder,
Rgba as RustRgba,
Expand Down Expand Up @@ -162,6 +165,7 @@ pub enum PixelType {
F32,
}

#[doc(hidden)]
impl From<RustColorType> for PixelType {
fn from(value: RustColorType) -> Self {
match value {
Expand Down Expand Up @@ -202,6 +206,7 @@ pub struct Strides {
pub channels: usize,
}

#[doc(hidden)]
impl From<RustColorType> for ColorType {
fn from(value: RustColorType) -> Self {
match value {
Expand Down
29 changes: 13 additions & 16 deletions pixel_bridge/rust/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,14 @@ use crate::{
image::{GenericImageDecoder, ImageDecoder},
vec_u8::VecU8,
};
use image::{
codecs::bmp::BmpDecoder as RustBmpDecoder, codecs::gif::GifDecoder as RustGifDecoder,
codecs::ico::IcoDecoder as RustIcoDecoder, codecs::jpeg::JpegDecoder as RustJpegDecoder,
codecs::png::PngDecoder as RustPngDecoder, codecs::tiff::TiffDecoder as RustTiffDecoder,
codecs::webp::WebPDecoder as RustWebPDecoder, ImageFormat, ImageReader as RustImageReader,
};
use image::codecs::bmp::BmpDecoder as RustBmpDecoder;
use image::codecs::gif::GifDecoder as RustGifDecoder;
use image::codecs::ico::IcoDecoder as RustIcoDecoder;
use image::codecs::jpeg::JpegDecoder as RustJpegDecoder;
use image::codecs::png::PngDecoder as RustPngDecoder;
use image::codecs::tiff::TiffDecoder as RustTiffDecoder;
use image::codecs::webp::WebPDecoder as RustWebPDecoder;
use image::{ImageFormat, ImageReader as RustImageReader};
use std::fs::File;
use std::io::{BufRead, BufReader, Cursor, Read, Seek};

Expand Down Expand Up @@ -46,6 +48,7 @@ pub enum Format {
Ico,
}

#[doc(hidden)]
impl From<Format> for ImageFormat {
fn from(value: Format) -> ImageFormat {
match value {
Expand Down Expand Up @@ -117,7 +120,7 @@ impl ImageReader {
// Delegate to non-inlined helper functions to reduce the stack frame size by 4k.
match inner.format() {
Some(ImageFormat::Png) => png(inner),
Some(ImageFormat::Jpeg) => jpeg(inner),
Some(ImageFormat::Jpeg) => jpeg(inner, self.jpeg_strict_mode),
Some(ImageFormat::WebP) => webp(inner),
Some(ImageFormat::Gif) => gif(inner),
Some(ImageFormat::Tiff) => tiff(inner),
Expand All @@ -138,10 +141,9 @@ fn png(
}

#[inline(never)]
fn jpeg(
d: RustImageReader<Box<dyn ReadSeek>>,
) -> Result<ImageDecoder, String> {
let decoder = RustJpegDecoder::new(d.into_inner()).map_err(|e| e.to_string())?;
fn jpeg(d: RustImageReader<Box<dyn ReadSeek>>, strict_mode: bool) -> Result<ImageDecoder, String> {
let mut decoder = RustJpegDecoder::new(d.into_inner()).map_err(|e| e.to_string())?;
decoder.set_strict_mode(strict_mode);
Ok(ImageDecoder::new(GenericImageDecoder::Jpeg(Box::new(decoder))))
}

Expand All @@ -163,8 +165,6 @@ fn gif(d: RustImageReader<Box<dyn ReadSeek>>) -> Result<ImageDecoder, String> {
fn tiff(d: RustImageReader<Box<dyn ReadSeek>>) -> Result<ImageDecoder, String> {
let mut reader = d.into_inner();
let mut should_premultiply = false;

// Check if we should premultiply based on ExtraSamples tag
{
let tiff_decoder_result = tiff::decoder::Decoder::new(&mut *reader);
if let Ok(mut decoder) = tiff_decoder_result
Expand All @@ -176,10 +176,7 @@ fn tiff(d: RustImageReader<Box<dyn ReadSeek>>) -> Result<ImageDecoder, String> {
should_premultiply = true;
}
}

// Seek back to start before passing to image crate decoder
reader.seek(std::io::SeekFrom::Start(0)).map_err(|e| e.to_string())?;

let mut image_decoder = ImageDecoder::new(GenericImageDecoder::Tiff(Box::new(
RustTiffDecoder::new(reader).map_err(|e| e.to_string())?,
)));
Expand Down
Loading