martinzink commented on code in PR #2258:
URL: https://github.com/apache/nifi-minifi-cpp/pull/2258#discussion_r4084324126
##########
minifi_rust/extensions/minifi_tensor/src/utils/tensor_helpers.rs:
##########
@@ -0,0 +1,210 @@
+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)?;
+ if cursor + tensor_len > flow_file_contents.len() {
+ return Err(MinifiError::custom(
+ "FlowFile contents are not in sync with tensor attributes",
+ ));
+ }
+ let tensor_data = &flow_file_contents[cursor..cursor + tensor_len];
+ result.push(Tensor::from_bytes(
+ tensor_dtype,
+ &tensor_shape,
+ tensor_data,
+ )?);
+ cursor += tensor_len;
+ }
+
+ if cursor != flow_file_contents.len() {
+ Err(MinifiError::custom(
+ "FlowFile contents are not in sync with tensor attributes",
+ ))
+ } else {
+ Ok(result)
+ }
+}
+
+pub(crate) fn tensor_as_f32(tensors: &[Tensor], index: usize) ->
Result<Vec<f32>, MinifiError> {
+ let tensor = tensors
+ .get(index)
+ .ok_or(MinifiError::custom("Invalid shape of tensors"))?;
+ let casted = tensor.convert_to(DatumType::F32)?;
+ Ok(casted.as_slice::<f32>()?.to_vec())
+}
Review Comment:
yeah thats a low hanging fruit and we should revisit this first thing after
this PR merges 👍
--
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]
