Feature Extractor: Allow output shapes without batch

Author: saddam213

OnnxStack.Core/Extensions/OrtValueExtensions.cs

@@ -105,7 +105,19 @@ public static OrtValue CreateOutputBuffer(this OnnxNamedMetadata metadata, ReadO
         public static DenseTensor<float> ToDenseTensor(this OrtValue ortValue)
         {
             var typeInfo = ortValue.GetTensorTypeAndShape();
-            var dimensions = typeInfo.Shape.ToInt();
+            return ortValue.ToDenseTensor(typeInfo.Shape.ToInt());
+        }
+
+
+        /// <summary>
+        /// Converts to DenseTensor<float>.
+        /// TODO: Optimization
+        /// </summary>
+        /// <param name="ortValue">The ort value.</param>
+        /// <returns></returns>
+        public static DenseTensor<float> ToDenseTensor(this OrtValue ortValue, ReadOnlySpan<int> dimensions)
+        {
+            var typeInfo = ortValue.GetTensorTypeAndShape();
             return typeInfo.ElementDataType switch
             {
                 TensorElementType.Float16 => new DenseTensor<float>(ortValue.GetTensorDataAsSpan<Float16>().ToFloat(), dimensions),

OnnxStack.FeatureExtractor/Pipelines/FeatureExtractorPipeline.cs

@@ -121,17 +121,19 @@ private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, Cancellatio
             var controlImage = await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
             var metadata = await _featureExtractorModel.GetMetadataAsync();
             cancellationToken.ThrowIfCancellationRequested();
+            var outputShape = new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize };
+            var outputBuffer = metadata.Outputs[0].Value.Dimensions.Length == 4 ? outputShape : outputShape[1..];
             using (var inferenceParameters = new OnnxInferenceParameters(metadata))
             {
                 inferenceParameters.AddInputTensor(controlImage);
-                inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
+                inferenceParameters.AddOutputBuffer(outputBuffer);
 
                 var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
                 using (var result = results.First())
                 {
                     cancellationToken.ThrowIfCancellationRequested();
 
-                    var resultTensor = result.ToDenseTensor();
+                    var resultTensor = result.ToDenseTensor(outputShape);
                     if (_featureExtractorModel.Normalize)
                         resultTensor.NormalizeMinMax();