martinzink commented on code in PR #2258: URL: https://github.com/apache/nifi-minifi-cpp/pull/2258#discussion_r4105363546
########## minifi_rust/extensions/minifi_tensor/src/low_level_processors/image_to_tensor.rs: ########## @@ -0,0 +1,625 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +pub(crate) mod image_to_tensor_def; + +use crate::low_level_processors::image_to_tensor::image_to_tensor_def::TENSOR_BYTES_ATTR; +use crate::utils::dimensions::Dimensions; +use crate::utils::per_channel_f32::PerChannelF32; +use crate::utils::tensor_helpers::{MinifiDatumType, load_as_image}; +pub(crate) use image_to_tensor_def::{ + COLOR_FORMAT, LETTERBOX_PAD_VALUE, MEAN, PIXEL_DIVISOR, RESIZE_FILTER, RESIZE_MODE, STD_DEV, + TARGET_HEIGHT, TARGET_WIDTH, TENSOR_SHAPE_FORMAT, +}; +use image_to_tensor_def::{ + IMG_ORG_HEIGHT_ATTR, IMG_ORG_WIDTH_ATTR, IMG_RESIZE_MODE_ATTR, SUCCESS, TENSOR_DTYPE_ATTR, + TENSOR_SHAPE_ATTR, +}; +use image_to_tensor_def::{IMG_TRG_HEIGHT_ATTR, IMG_TRG_WIDTH_ATTR, TENSORS_LEN_ATTR}; +use minifi_native::macros::{ComponentIdentifier, PropertyType}; +use minifi_native::{ + FlowFileTransform, GetAttribute, GetControllerService, GetId, GetProperty, InputStream, Logger, + MinifiError, ProcessError, RouteErrorExt, Schedule, TransformedFlowFile, +}; +use strum_macros::{Display, EnumString, IntoStaticStr, VariantNames}; +use tract::Tensor; + +tract::impl_ndarray_interop!(); + +#[derive( + Debug, Clone, Copy, PartialEq, Display, EnumString, VariantNames, IntoStaticStr, PropertyType, +)] +#[strum(serialize_all = "PascalCase", const_into_str)] +pub(crate) enum ResizeFilter { + Nearest, + Bilinear, + Bicubic, + Lanczos3, +} + +impl From<ResizeFilter> for image::imageops::FilterType { + fn from(filter: ResizeFilter) -> Self { + match filter { + ResizeFilter::Nearest => image::imageops::FilterType::Nearest, + ResizeFilter::Bilinear => image::imageops::FilterType::Triangle, + ResizeFilter::Bicubic => image::imageops::FilterType::CatmullRom, + ResizeFilter::Lanczos3 => image::imageops::FilterType::Lanczos3, + } + } +} + +#[derive( + Debug, Clone, Copy, PartialEq, Display, EnumString, VariantNames, IntoStaticStr, PropertyType, +)] +#[strum(serialize_all = "UPPERCASE", const_into_str)] +pub(crate) enum ColorFormat { + Rgb, + Bgr, + Grayscale, +} + +#[derive( + Debug, Clone, Copy, PartialEq, Display, EnumString, VariantNames, IntoStaticStr, PropertyType, +)] +#[strum(serialize_all = "UPPERCASE", const_into_str)] +pub(crate) enum TensorShapeFormat { + Chw, // channel, height, width + Hwc, // height, width, channel +} + +#[derive( + Debug, Clone, Copy, PartialEq, Display, EnumString, VariantNames, IntoStaticStr, PropertyType, +)] +#[strum(serialize_all = "PascalCase", const_into_str)] +pub(crate) enum ResizeMode { + Stretch, + Letterbox, +} + +#[derive(ComponentIdentifier)] +pub(crate) struct ImageToTensor { + target_width: u32, + target_height: u32, + resize_filter: ResizeFilter, + resize_mode: ResizeMode, + color_format: ColorFormat, + tensor_shape_format: TensorShapeFormat, + mean: PerChannelF32, + std_dev: PerChannelF32, + pixel_divisor: f32, + letterbox_pad_value: f32, +} + +impl Schedule for ImageToTensor { + fn schedule<Ctx: GetProperty, L: Logger>( + context: &Ctx, + _logger: &L, + ) -> Result<Self, MinifiError> + where + Self: Sized, + { + let target_width = context.get_property(&TARGET_WIDTH)?; + let target_height = context.get_property(&TARGET_HEIGHT)?; + let resize_filter = context.get_property(&RESIZE_FILTER)?; + let resize_mode = context.get_property(&RESIZE_MODE)?; + let color_format = context.get_property(&COLOR_FORMAT)?; + let tensor_shape_format = context.get_property(&TENSOR_SHAPE_FORMAT)?; + let mean = context.get_property(&MEAN)?; + let std_dev = context.get_property(&STD_DEV)?; + if std_dev.contains_zero() { + return Err(MinifiError::validation( + "Standard Deviation components must be non-zero", + )); + } + let pixel_divisor = context.get_property(&PIXEL_DIVISOR)?; + if pixel_divisor == 0.0 { + return Err(MinifiError::validation("Pixel divisor must be non-zero")); + } + let letterbox_pad_value = context.get_property(&LETTERBOX_PAD_VALUE)?; + + Ok(Self { + target_width, + target_height, + resize_filter, + resize_mode, + color_format, + tensor_shape_format, + mean, + std_dev, + pixel_divisor, + letterbox_pad_value, + }) + } +} + +struct MaskedRgbImage { + img: image::RgbImage, + mask: Vec<bool>, +} + +impl ImageToTensor { + fn stretch_resize(&self, img: image::DynamicImage) -> MaskedRgbImage { + let resized = img + .resize_exact( + self.target_width, + self.target_height, + self.resize_filter.into(), + ) + .to_rgb8(); + let mask = vec![true; (self.target_width * self.target_height) as usize]; + MaskedRgbImage { img: resized, mask } + } + + fn letterbox_resize(&self, img: image::DynamicImage) -> MaskedRgbImage { + let (src_w, src_h) = (img.width() as f32, img.height() as f32); + let scale = (self.target_width as f32 / src_w).min(self.target_height as f32 / src_h); + let new_w = (src_w * scale).round().max(1.0) as u32; + let new_h = (src_h * scale).round().max(1.0) as u32; + let scaled = img + .resize_exact(new_w, new_h, self.resize_filter.into()) + .to_rgb8(); + + let pad_x = (self.target_width - new_w) / 2; + let pad_y = (self.target_height - new_h) / 2; + + let mut canvas = image::RgbImage::from_pixel( + self.target_width, + self.target_height, + image::Rgb([0, 0, 0]), + ); + image::imageops::overlay(&mut canvas, &scaled, pad_x as i64, pad_y as i64); + + // Mask to track which pixel is part of source and which is padding + let mut mask = vec![false; (self.target_width * self.target_height) as usize]; + for y in pad_y..(new_h + pad_y) { + for x in pad_x..(new_w + pad_x) { + mask[(y * self.target_width + x) as usize] = true; + } + } + MaskedRgbImage { img: canvas, mask } + } + + fn resize_rgb(&self, img: image::DynamicImage) -> MaskedRgbImage { + match self.resize_mode { + ResizeMode::Stretch => self.stretch_resize(img), + ResizeMode::Letterbox => self.letterbox_resize(img), + } + } + + pub fn tensor_bytes(&self, img: image::DynamicImage) -> Vec<u8> { + let num_channels: usize = match self.color_format { + ColorFormat::Grayscale => 1, + _ => 3, + }; + let total_pixels = (self.target_width * self.target_height) as usize; + let mut tensor_bytes = Vec::with_capacity(total_pixels * num_channels * 4); + + let masked_img = self.resize_rgb(img); Review Comment: image crate already checks inside, so if the dimensions do check out no transformation will be done ``` // check if the new dimensions are the same as the old. if they are, make a copy instead of resampling if (nwidth, nheight) == image.dimensions() { let mut tmp = image.buffer_like(); tmp.copy_from(image, 0, 0).unwrap(); return tmp; } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
