martinzink commented on code in PR #2258:
URL: https://github.com/apache/nifi-minifi-cpp/pull/2258#discussion_r3989473421
##########
minifi_rust/extensions/minifi_tensor/src/utils/tensor_helpers.rs:
##########
@@ -0,0 +1,156 @@
+use image::{DynamicImage, ImageResult};
+use minifi_native::{GetAttribute, InputStream, MinifiError};
+use strum_macros::{Display, EnumString};
+use tract::__ndarray_interop::TensorInterface;
+use tract::Tensor;
+use tract::prelude::DatumType;
+
+tract::impl_ndarray_interop!();
+
+fn parse_tensor_shape<Context: GetAttribute>(
+ context: &Context,
+ id: usize,
+) -> Result<Vec<usize>, MinifiError> {
+ let shape_str = context.get_required_attribute(&format!("tensor.{}.shape",
id))?;
+
+ if shape_str.trim().is_empty() {
+ return Ok(Vec::new());
+ }
+
+ let shape = shape_str
+ .split(',')
+ .map(|s| s.trim().parse::<usize>())
+ .collect::<Result<Vec<usize>, _>>()?;
+
+ Ok(shape)
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Display, EnumString)]
+#[strum(serialize_all = "PascalCase", const_into_str)]
+pub(crate) enum MinifiDatumType {
+ F32,
+}
+
+impl From<MinifiDatumType> for DatumType {
+ fn from(value: MinifiDatumType) -> Self {
+ match value {
+ MinifiDatumType::F32 => DatumType::F32,
+ }
+ }
+}
+
+fn numeric_datum_type_from_str(s: &str) -> Option<DatumType> {
+ Some(match s {
+ "U8" => DatumType::U8,
+ "U16" => DatumType::U16,
+ "U32" => DatumType::U32,
+ "U64" => DatumType::U64,
+ "I8" => DatumType::I8,
+ "I16" => DatumType::I16,
+ "I32" => DatumType::I32,
+ "I64" => DatumType::I64,
+ "F16" => DatumType::F16,
+ "F32" => DatumType::F32,
+ "F64" => DatumType::F64,
+ _ => return None,
+ })
+}
+
+fn parse_tensor_dtype<Context: GetAttribute>(
+ context: &Context,
+ id: usize,
+) -> Result<DatumType, MinifiError> {
+ let dtype_str = context.get_required_attribute(&format!("tensor.{}.dtype",
id))?;
+ numeric_datum_type_from_str(&dtype_str).ok_or_else(|| {
+ MinifiError::custom(format!(
+ "Unsupported tensor.{}.dtype '{}': only numeric tensors can be
read",
+ id, dtype_str
+ ))
+ })
+}
+
+pub(crate) fn deserialize_tensors<Context: GetAttribute>(
+ context: &Context,
+ input_stream: &mut dyn InputStream,
+) -> Result<Vec<Tensor>, MinifiError> {
+ let mut result = vec![];
+
+ let mut flow_file_contents = Vec::new();
+ input_stream.read_to_end(&mut flow_file_contents)?;
+ let number_of_tensors = context
+ .get_required_attribute("tensors.len")?
+ .parse::<usize>()?;
+
+ let mut cursor = 0usize;
+ for i in 0..number_of_tensors {
+ let tensor_len = context
+ .get_required_attribute(&format!("tensor.{}.bytes", i))?
+ .parse::<usize>()?;
+ let tensor_shape = parse_tensor_shape(context, i)?;
+ let tensor_dtype = parse_tensor_dtype(context, i)?;
+ let tensor_data = &flow_file_contents[cursor..cursor + tensor_len];
Review Comment:
valid point, [fix panic on malformed tensor payload in
tensor_helpers.rs](https://github.com/apache/nifi-minifi-cpp/pull/2258/commits/76aad697eeaf6d05f085fa862f8b0efe6d7ead76)
--
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]
