FeatureExtractor normalization

Author: saddam213

OnnxStack.Console/Examples/ControlNetFeatureExample.cs

@@ -35,7 +35,7 @@ public async Task RunAsync()
             var inputImage = await OnnxImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
 
             // Create Annotation pipeline
-            var annotationPipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutputTensor: true);
+            var annotationPipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutput: true);
 
             // Create Depth Image
             var controlImage = await annotationPipeline.RunAsync(inputImage);

OnnxStack.Console/Examples/FeatureExtractorExample.cs

@@ -37,7 +37,7 @@ public async Task RunAsync()
             {
                 FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\canny.onnx"),
                 FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\hed.onnx"),
-                FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutputTensor: true, inputResizeMode: ImageResizeMode.Stretch),
+                FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutput: true, inputResizeMode: ImageResizeMode.Stretch),
                 FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024, setOutputToInputAlpha: true, inputResizeMode: ImageResizeMode.Stretch)
             };
 

OnnxStack.Core/Image/Extensions.cs

@@ -16,6 +16,17 @@ public static class Extensions
         /// <param name="imageTensor">The image tensor.</param>
         /// <returns></returns>
         public static OnnxImage ToImageMask(this DenseTensor<float> imageTensor)
+        {
+            return new OnnxImage(imageTensor.FromMaskTensor());
+        }
+
+
+        /// <summary>
+        /// Convert from single channle mask tensor to Rgba32 (Greyscale)
+        /// </summary>
+        /// <param name="imageTensor">The image tensor.</param>
+        /// <returns></returns>
+        public static Image<Rgba32> FromMaskTensor(this DenseTensor<float> imageTensor)
         {
             var width = imageTensor.Dimensions[3];
             var height = imageTensor.Dimensions[2];
@@ -28,7 +39,7 @@ public static OnnxImage ToImageMask(this DenseTensor<float> imageTensor)
                         result[x, y] = new L8((byte)(imageTensor[0, 0, y, x] * 255.0f));
                     }
                 }
-                return new OnnxImage(result.CloneAs<Rgba32>());
+                return result.CloneAs<Rgba32>();
             }
         }
 

OnnxStack.Core/Image/OnnxImage.cs

@@ -64,27 +64,35 @@ public OnnxImage(DenseTensor<float> imageTensor, ImageNormalizeType normalizeTyp
         {
             var height = imageTensor.Dimensions[2];
             var width = imageTensor.Dimensions[3];
-            var hasTransparency = imageTensor.Dimensions[1] == 4;
-            _imageData = new Image<Rgba32>(width, height);
-            for (var y = 0; y < height; y++)
+            var channels = imageTensor.Dimensions[1];
+            if (channels == 1)
             {
-                for (var x = 0; x < width; x++)
+                _imageData = imageTensor.FromMaskTensor();
+            }
+            else
+            {
+                var hasTransparency = channels == 4;
+                _imageData = new Image<Rgba32>(width, height);
+                for (var y = 0; y < height; y++)
                 {
-                    if (normalizeType == ImageNormalizeType.ZeroToOne)
-                    {
-                        _imageData[x, y] = new Rgba32(
-                            DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
-                            DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
-                            DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
-                            hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
-                    }
-                    else
+                    for (var x = 0; x < width; x++)
                     {
-                        _imageData[x, y] = new Rgba32(
-                            DenormalizeOneToOneToByte(imageTensor, 0, y, x),
-                            DenormalizeOneToOneToByte(imageTensor, 1, y, x),
-                            DenormalizeOneToOneToByte(imageTensor, 2, y, x),
-                            hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
+                        if (normalizeType == ImageNormalizeType.ZeroToOne)
+                        {
+                            _imageData[x, y] = new Rgba32(
+                                DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
+                                DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
+                                DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
+                                hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
+                        }
+                        else
+                        {
+                            _imageData[x, y] = new Rgba32(
+                                DenormalizeOneToOneToByte(imageTensor, 0, y, x),
+                                DenormalizeOneToOneToByte(imageTensor, 1, y, x),
+                                DenormalizeOneToOneToByte(imageTensor, 2, y, x),
+                                hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
+                        }
                     }
                 }
             }

OnnxStack.FeatureExtractor/Common/FeatureExtractorModel.cs

@@ -17,17 +17,18 @@ public FeatureExtractorModel(FeatureExtractorModelConfig configuration)
 
         public int OutputChannels => _configuration.OutputChannels;
         public int SampleSize => _configuration.SampleSize;
-        public bool NormalizeOutputTensor => _configuration.NormalizeOutputTensor;
+        public bool NormalizeOutput => _configuration.NormalizeOutput;
         public bool SetOutputToInputAlpha => _configuration.SetOutputToInputAlpha;
         public ImageResizeMode InputResizeMode => _configuration.InputResizeMode;
-        public ImageNormalizeType InputNormalization => _configuration.NormalizeInputTensor;
+        public ImageNormalizeType NormalizeType => _configuration.NormalizeType;
+        public bool NormalizeInput => _configuration.NormalizeInput;
 
         public static FeatureExtractorModel Create(FeatureExtractorModelConfig configuration)
         {
             return new FeatureExtractorModel(configuration);
         }
 
-        public static FeatureExtractorModel Create(string modelFile, int sampleSize = 0, int outputChannels = 1, bool normalizeOutputTensor = false, ImageNormalizeType normalizeInputTensor = ImageNormalizeType.ZeroToOne, ImageResizeMode inputResizeMode = ImageResizeMode.Crop, bool setOutputToInputAlpha = false, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
+        public static FeatureExtractorModel Create(string modelFile, int sampleSize = 0, int outputChannels = 1, ImageNormalizeType normalizeType = ImageNormalizeType.ZeroToOne, bool normalizeInput = true, bool normalizeOutput = false, ImageResizeMode inputResizeMode = ImageResizeMode.Crop, bool setOutputToInputAlpha = false, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
         {
             var configuration = new FeatureExtractorModelConfig
             {
@@ -38,12 +39,12 @@ public static FeatureExtractorModel Create(string modelFile, int sampleSize = 0,
                 IntraOpNumThreads = 0,
                 OnnxModelPath = modelFile,
 
-            
                 SampleSize = sampleSize,
                 OutputChannels = outputChannels,
-                NormalizeOutputTensor = normalizeOutputTensor,
+                NormalizeType = normalizeType,
+                NormalizeInput = normalizeInput,
+                NormalizeOutput = normalizeOutput,
                 SetOutputToInputAlpha = setOutputToInputAlpha,
-                NormalizeInputTensor = normalizeInputTensor,
                 InputResizeMode = inputResizeMode
             };
             return new FeatureExtractorModel(configuration);

OnnxStack.FeatureExtractor/Common/FeatureExtractorModelConfig.cs

@@ -7,9 +7,10 @@ public record FeatureExtractorModelConfig : OnnxModelConfig
     {
         public int SampleSize { get; set; }
         public int OutputChannels { get; set; }
-        public bool NormalizeOutputTensor { get; set; }
+        public bool NormalizeOutput { get; set; }
         public bool SetOutputToInputAlpha { get; set; }
         public ImageResizeMode InputResizeMode { get; set; }
-        public ImageNormalizeType NormalizeInputTensor { get; set; }
+        public ImageNormalizeType NormalizeType { get; set; }
+        public bool NormalizeInput { get; set; }
     }
 }

OnnxStack.FeatureExtractor/Pipelines/FeatureExtractorPipeline.cs

@@ -62,16 +62,30 @@ public async Task UnloadAsync()
         }
 
 
+        /// <summary>
+        /// Generates the feature extractor image
+        /// </summary>
+        /// <param name="inputImage">The input image.</param>
+        /// <returns></returns>
+        public async Task<DenseTensor<float>> RunAsync(DenseTensor<float> inputTensor, CancellationToken cancellationToken = default)
+        {
+            var timestamp = _logger?.LogBegin("Extracting DenseTensor feature...");
+            var result = await ExtractTensorAsync(inputTensor, cancellationToken);
+            _logger?.LogEnd("Extracting DenseTensor feature complete.", timestamp);
+            return result;
+        }
+
+
         /// <summary>
         /// Generates the feature extractor image
         /// </summary>
         /// <param name="inputImage">The input image.</param>
         /// <returns></returns>
         public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
         {
-            var timestamp = _logger?.LogBegin("Extracting image feature...");
-            var result = await RunInternalAsync(inputImage, cancellationToken);
-            _logger?.LogEnd("Extracting image feature complete.", timestamp);
+            var timestamp = _logger?.LogBegin("Extracting OnnxImage feature...");
+            var result = await ExtractImageAsync(inputImage, cancellationToken);
+            _logger?.LogEnd("Extracting OnnxImage feature complete.", timestamp);
             return result;
         }
 
@@ -83,13 +97,13 @@ public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken ca
         /// <returns></returns>
         public async Task<OnnxVideo> RunAsync(OnnxVideo video, CancellationToken cancellationToken = default)
         {
-            var timestamp = _logger?.LogBegin("Extracting video features...");
+            var timestamp = _logger?.LogBegin("Extracting OnnxVideo features...");
             var featureFrames = new List<OnnxImage>();
             foreach (var videoFrame in video.Frames)
             {
                 featureFrames.Add(await RunAsync(videoFrame, cancellationToken));
             }
-            _logger?.LogEnd("Extracting video features complete.", timestamp);
+            _logger?.LogEnd("Extracting OnnxVideo features complete.", timestamp);
             return new OnnxVideo(video.Info, featureFrames);
         }
 
@@ -102,28 +116,62 @@ public async Task<OnnxVideo> RunAsync(OnnxVideo video, CancellationToken cancell
         /// <returns></returns>
         public async IAsyncEnumerable<OnnxImage> RunAsync(IAsyncEnumerable<OnnxImage> imageFrames, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
-            var timestamp = _logger?.LogBegin("Extracting video stream features...");
+            var timestamp = _logger?.LogBegin("Extracting OnnxImage stream features...");
             await foreach (var imageFrame in imageFrames)
             {
-                yield return await RunInternalAsync(imageFrame, cancellationToken);
+                yield return await ExtractImageAsync(imageFrame, cancellationToken);
             }
-            _logger?.LogEnd("Extracting video stream features complete.", timestamp);
+            _logger?.LogEnd("Extracting OnnxImage stream features complete.", timestamp);
         }
 
 
         /// <summary>
-        /// Runs the pipeline
+        /// Extracts the feature to OnnxImage.
         /// </summary>
         /// <param name="inputImage">The input image.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
+        private async Task<OnnxImage> ExtractImageAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
         {
             var originalWidth = inputImage.Width;
             var originalHeight = inputImage.Height;
             var inputTensor = _featureExtractorModel.SampleSize <= 0
-                ? await inputImage.GetImageTensorAsync(_featureExtractorModel.InputNormalization)
-                : await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, _featureExtractorModel.InputNormalization, resizeMode: _featureExtractorModel.InputResizeMode);
+                ? await inputImage.GetImageTensorAsync(_featureExtractorModel.NormalizeType)
+                : await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, _featureExtractorModel.NormalizeType, resizeMode: _featureExtractorModel.InputResizeMode);
+
+            var outputTensor = await RunInternalAsync(inputTensor, cancellationToken);
+            var imageResult = new OnnxImage(outputTensor, _featureExtractorModel.NormalizeType);
+
+            if (_featureExtractorModel.InputResizeMode == ImageResizeMode.Stretch && (imageResult.Width != originalWidth || imageResult.Height != originalHeight))
+                imageResult.Resize(originalHeight, originalWidth, _featureExtractorModel.InputResizeMode);
+
+            return imageResult;
+        }
+
+
+        /// <summary>
+        /// Extracts the feature to DenseTensor.
+        /// </summary>
+        /// <param name="inputTensor">The input tensor.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public async Task<DenseTensor<float>> ExtractTensorAsync(DenseTensor<float> inputTensor, CancellationToken cancellationToken = default)
+        {
+            if (_featureExtractorModel.NormalizeInput && _featureExtractorModel.NormalizeType == ImageNormalizeType.ZeroToOne)
+                inputTensor.NormalizeOneOneToZeroOne();
+
+            return await RunInternalAsync(inputTensor, cancellationToken); 
+        }
+
+
+        /// <summary>
+        /// Runs the pipeline
+        /// </summary>
+        /// <param name="inputTensor">The input tensor.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        private async Task<DenseTensor<float>> RunInternalAsync(DenseTensor<float> inputTensor, CancellationToken cancellationToken = default)
+        {
             var metadata = await _featureExtractorModel.GetMetadataAsync();
             cancellationToken.ThrowIfCancellationRequested();
             var outputShape = new[] { 1, _featureExtractorModel.OutputChannels, inputTensor.Dimensions[2], inputTensor.Dimensions[3] };
@@ -139,21 +187,13 @@ private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, Cancellatio
                     cancellationToken.ThrowIfCancellationRequested();
 
                     var outputTensor = inferenceResult.ToDenseTensor(outputShape);
-                    if (_featureExtractorModel.NormalizeOutputTensor)
+                    if (_featureExtractorModel.NormalizeOutput)
                         outputTensor.NormalizeMinMax();
 
-                    var imageResult = default(OnnxImage);
                     if (_featureExtractorModel.SetOutputToInputAlpha)
-                        imageResult = new OnnxImage(AddAlphaChannel(inputTensor, outputTensor), _featureExtractorModel.InputNormalization);
-                    else if (_featureExtractorModel.OutputChannels >= 3)
-                        imageResult = new OnnxImage(outputTensor, _featureExtractorModel.InputNormalization);
-                    else
-                        imageResult = outputTensor.ToImageMask();
-
-                    if (_featureExtractorModel.InputResizeMode == ImageResizeMode.Stretch && (imageResult.Width != originalWidth || imageResult.Height != originalHeight))
-                        imageResult.Resize(originalHeight, originalWidth, _featureExtractorModel.InputResizeMode);
+                        return AddAlphaChannel(inputTensor, outputTensor);
 
-                    return imageResult;
+                    return outputTensor;
                 }
             }
         }
@@ -200,7 +240,7 @@ public static FeatureExtractorPipeline CreatePipeline(FeatureExtractorModelSet m
         /// <param name="executionProvider">The execution provider.</param>
         /// <param name="logger">The logger.</param>
         /// <returns></returns>
-        public static FeatureExtractorPipeline CreatePipeline(string modelFile, int sampleSize = 0, int outputChannels = 1, bool normalizeOutputTensor = false, ImageNormalizeType normalizeInputTensor = ImageNormalizeType.ZeroToOne, ImageResizeMode inputResizeMode = ImageResizeMode.Crop, bool setOutputToInputAlpha = false, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML, ILogger logger = default)
+        public static FeatureExtractorPipeline CreatePipeline(string modelFile, int sampleSize = 0, int outputChannels = 1, ImageNormalizeType normalizeType = ImageNormalizeType.ZeroToOne, bool normalizeInput = true, bool normalizeOutput = false, ImageResizeMode inputResizeMode = ImageResizeMode.Crop, bool setOutputToInputAlpha = false, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML, ILogger logger = default)
         {
             var name = Path.GetFileNameWithoutExtension(modelFile);
             var configuration = new FeatureExtractorModelSet
@@ -214,9 +254,10 @@ public static FeatureExtractorPipeline CreatePipeline(string modelFile, int samp
                     OnnxModelPath = modelFile,
                     SampleSize = sampleSize,
                     OutputChannels = outputChannels,
-                    NormalizeOutputTensor = normalizeOutputTensor,
+                    NormalizeOutput = normalizeOutput,
+                    NormalizeInput = normalizeInput,
+                    NormalizeType = normalizeType,
                     SetOutputToInputAlpha = setOutputToInputAlpha,
-                    NormalizeInputTensor = normalizeInputTensor,
                     InputResizeMode = inputResizeMode
                 }
             };