Upscale Pipeline Improvments

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.Console/Examples/UpscaleExample.cs

@@ -25,7 +25,7 @@ public async Task RunAsync()
             var inputImage = await OnnxImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
 
             // Create Pipeline
-            var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
+            var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4, 512);
 
             // Run pipeline
             var result = await pipeline.RunAsync(inputImage);

OnnxStack.Console/Examples/UpscaleStreamExample.cs

@@ -26,7 +26,7 @@ public async Task RunAsync()
             var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
 
             // Create pipeline
-            var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
+            var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4, 512);
 
             // Load pipeline
             await pipeline.LoadAsync();

OnnxStack.Core/Extensions/TensorExtension.cs

@@ -5,6 +5,7 @@
 using System.Numerics.Tensors;
 using System.Numerics;
 using OnnxStack.Core.Model;
+using System.Threading.Tasks;
 
 namespace OnnxStack.Core
 {
@@ -415,75 +416,22 @@ private static DenseTensor<float> ConcatenateAxis2(DenseTensor<float> tensor1, D
 
 
         /// <summary>
-        /// Splits the Tensor into 4 equal tiles.
+        /// Normalizes the tensor values from range -1 to 1 to 0 to 1.
         /// </summary>
-        /// <param name="sourceTensor">The source tensor.</param>
-        /// <returns>TODO: Optimize</returns>
-        public static ImageTiles SplitTiles(this DenseTensor<float> sourceTensor)
+        /// <param name="imageTensor">The image tensor.</param>
+        public static void NormalizeOneOneToZeroOne(this DenseTensor<float> imageTensor)
         {
-            int tileWidth = sourceTensor.Dimensions[3] / 2;
-            int tileHeight = sourceTensor.Dimensions[2] / 2;
-
-            return new ImageTiles(
-                SplitTile(sourceTensor, 0, 0, tileHeight, tileWidth),
-                SplitTile(sourceTensor, 0, tileWidth, tileHeight, tileWidth * 2),
-                SplitTile(sourceTensor, tileHeight, 0, tileHeight * 2, tileWidth),
-                SplitTile(sourceTensor, tileHeight, tileWidth, tileHeight * 2, tileWidth * 2));
-        }
-
-        private static DenseTensor<float> SplitTile(DenseTensor<float> tensor, int startRow, int startCol, int endRow, int endCol)
-        {
-            int height = endRow - startRow;
-            int width = endCol - startCol;
-            int channels = tensor.Dimensions[1];
-            var slicedData = new DenseTensor<float>(new[] { 1, channels, height, width });
-            for (int c = 0; c < channels; c++)
-            {
-                for (int i = 0; i < height; i++)
-                {
-                    for (int j = 0; j < width; j++)
-                    {
-                        slicedData[0, c, i, j] = tensor[0, c, startRow + i, startCol + j];
-                    }
-                }
-            }
-            return slicedData;
+            Parallel.For(0, (int)imageTensor.Length, (i) => imageTensor.SetValue(i, imageTensor.GetValue(i) / 2f + 0.5f));
         }
 
 
         /// <summary>
-        /// Rejoins the tiles into a single Tensor.
+        /// Normalizes the tensor values from range 0 to 1 to -1 to 1.
         /// </summary>
-        /// <param name="tiles">The tiles.</param>
-        /// <returns>TODO: Optimize</returns>
-        public static DenseTensor<float> RejoinTiles(this ImageTiles tiles)
+        /// <param name="imageTensor">The image tensor.</param>
+        public static void NormalizeZeroOneToOneOne(this DenseTensor<float> imageTensor)
         {
-            int totalHeight = tiles.Tile1.Dimensions[2] + tiles.Tile3.Dimensions[2];
-            int totalWidth = tiles.Tile1.Dimensions[3] + tiles.Tile2.Dimensions[3];
-            int channels = tiles.Tile1.Dimensions[1];
-            var destination = new DenseTensor<float>(new[] { 1, channels, totalHeight, totalWidth });
-            RejoinTile(destination, tiles.Tile1, 0, 0);
-            RejoinTile(destination, tiles.Tile2, 0, tiles.Tile1.Dimensions[3]);
-            RejoinTile(destination, tiles.Tile3, tiles.Tile1.Dimensions[2], 0);
-            RejoinTile(destination, tiles.Tile4, tiles.Tile1.Dimensions[2], tiles.Tile1.Dimensions[3]);
-            return destination;
-        }
-
-        private static void RejoinTile(DenseTensor<float> destination, DenseTensor<float> tile, int startRow, int startCol)
-        {
-            int channels = tile.Dimensions[1];
-            int height = tile.Dimensions[2];
-            int width = tile.Dimensions[3];
-            for (int c = 0; c < channels; c++)
-            {
-                for (int i = 0; i < height; i++)
-                {
-                    for (int j = 0; j < width; j++)
-                    {
-                        destination[0, c, startRow + i, startCol + j] = tile[0, c, i, j];
-                    }
-                }
-            }
+            Parallel.For(0, (int)imageTensor.Length, (i) => imageTensor.SetValue(i, 2f * imageTensor.GetValue(i) - 1f));
         }
     }
 }

OnnxStack.Core/Image/Extensions.cs

@@ -1,7 +1,9 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core.Model;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
 using SixLabors.ImageSharp.Processing;
+using System.Threading.Tasks;
 
 namespace OnnxStack.Core.Image
 {
@@ -14,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];
@@ -26,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>();
             }
         }
 
@@ -40,6 +53,104 @@ public static ResizeMode ToResizeMode(this ImageResizeMode resizeMode)
             };
         }
 
+
+        /// <summary>
+        /// Splits the Tensor into 4 equal tiles.
+        /// </summary>
+        /// <param name="sourceTensor">The source tensor.</param>
+        /// <returns>TODO: Optimize</returns>
+        public static ImageTiles SplitImageTiles(this DenseTensor<float> sourceTensor, int overlap = 20)
+        {
+            var tileWidth = sourceTensor.Dimensions[3] / 2;
+            var tileHeight = sourceTensor.Dimensions[2] / 2;
+            return new ImageTiles(tileWidth, tileHeight, overlap,
+                SplitImageTile(sourceTensor, 0, 0, tileHeight + overlap, tileWidth + overlap),
+                SplitImageTile(sourceTensor, 0, tileWidth - overlap, tileHeight + overlap, tileWidth * 2),
+                SplitImageTile(sourceTensor, tileHeight - overlap, 0, tileHeight * 2, tileWidth + overlap),
+                SplitImageTile(sourceTensor, tileHeight - overlap, tileWidth - overlap, tileHeight * 2, tileWidth * 2));
+        }
+
+
+        /// <summary>
+        /// Splits a tile from the source.
+        /// </summary>
+        /// <param name="source">The tensor.</param>
+        /// <param name="startRow">The start row.</param>
+        /// <param name="startCol">The start col.</param>
+        /// <param name="endRow">The end row.</param>
+        /// <param name="endCol">The end col.</param>
+        /// <returns></returns>
+        private static DenseTensor<float> SplitImageTile(DenseTensor<float> source, int startRow, int startCol, int endRow, int endCol)
+        {
+            int height = endRow - startRow;
+            int width = endCol - startCol;
+            int channels = source.Dimensions[1];
+            var splitTensor = new DenseTensor<float>(new[] { 1, channels, height, width });
+            Parallel.For(0, channels, (c) =>
+            {
+                Parallel.For(0, height, (i) =>
+                {
+                    Parallel.For(0, width, (j) =>
+                    {
+                        splitTensor[0, c, i, j] = source[0, c, startRow + i, startCol + j];
+                    });
+                });
+            });
+            return splitTensor;
+        }
+
+
+        /// <summary>
+        /// Joins the tiles into a single Tensor.
+        /// </summary>
+        /// <param name="tiles">The tiles.</param>
+        /// <returns>TODO: Optimize</returns>
+        public static DenseTensor<float> JoinImageTiles(this ImageTiles tiles)
+        {
+            var totalWidth = tiles.Width * 2;
+            var totalHeight = tiles.Height * 2;
+            var channels = tiles.Tile1.Dimensions[1];
+            var destination = new DenseTensor<float>(new[] { 1, channels, totalHeight, totalWidth });
+            JoinImageTile(destination, tiles.Tile1, 0, 0, tiles.Height + tiles.Overlap, tiles.Width + tiles.Overlap);
+            JoinImageTile(destination, tiles.Tile2, 0, tiles.Width - tiles.Overlap, tiles.Height + tiles.Overlap, totalWidth);
+            JoinImageTile(destination, tiles.Tile3, tiles.Height - tiles.Overlap, 0, totalHeight, tiles.Width + tiles.Overlap);
+            JoinImageTile(destination, tiles.Tile4, tiles.Height - tiles.Overlap, tiles.Width - tiles.Overlap, totalHeight, totalWidth);
+            return destination;
+        }
+
+
+        /// <summary>
+        /// Joins the tile to the destination tensor.
+        /// </summary>
+        /// <param name="destination">The destination.</param>
+        /// <param name="tile">The tile.</param>
+        /// <param name="startRow">The start row.</param>
+        /// <param name="startCol">The start col.</param>
+        /// <param name="endRow">The end row.</param>
+        /// <param name="endCol">The end col.</param>
+        private static void JoinImageTile(DenseTensor<float> destination, DenseTensor<float> tile, int startRow, int startCol, int endRow, int endCol)
+        {
+            int height = endRow - startRow;
+            int width = endCol - startCol;
+            int channels = tile.Dimensions[1];
+            Parallel.For(0, channels, (c) =>
+            {
+                Parallel.For(0, height, (i) =>
+                {
+                    Parallel.For(0, width, (j) =>
+                    {
+                        var value = tile[0, c, i, j];
+                        var existing = destination[0, c, startRow + i, startCol + j];
+                        if (existing > 0)
+                        {
+                            // Blend ovelap
+                            value = (existing + value) / 2f;
+                        }
+                        destination[0, c, startRow + i, startCol + j] = value;
+                    });
+                });
+            });
+        }
     }
 
     public enum ImageNormalizeType

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.Core/Model/ImageTiles.cs

@@ -2,5 +2,5 @@
 
 namespace OnnxStack.Core.Model
 {
-    public record ImageTiles(DenseTensor<float> Tile1, DenseTensor<float> Tile2, DenseTensor<float> Tile3, DenseTensor<float> Tile4);
+    public record ImageTiles(int Width, int Height, int Overlap, DenseTensor<float> Tile1, DenseTensor<float> Tile2, DenseTensor<float> Tile3, DenseTensor<float> Tile4);
 }

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; }
     }
 }