Move common InputImage to OnnxStack.Core, Tidy Upscale implementation

Author: saddam213

OnnxStack.Console/Examples/StableDiffusionGif.cs

@@ -1,4 +1,5 @@
-using OnnxStack.StableDiffusion;
+using OnnxStack.Core.Image;
+using OnnxStack.StableDiffusion;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.StableDiffusion.Enums;

OnnxStack.Console/Examples/UpscaleExample.cs

@@ -1,4 +1,5 @@
-using OnnxStack.ImageUpscaler.Services;
+using OnnxStack.Core.Image;
+using OnnxStack.ImageUpscaler.Services;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
 
@@ -42,7 +43,7 @@ public async Task RunAsync()
             var inputImage = await Image.LoadAsync<Rgba32>(imageFile);
 
             OutputHelpers.WriteConsole("Upscaling Image...", ConsoleColor.Cyan);
-            var result = await _imageUpscaleService.GenerateAsync(modelSet, inputImage);
+            var result = await _imageUpscaleService.GenerateAsImageAsync(modelSet, new InputImage(inputImage));
             await result.SaveAsPngAsync(Path.Combine(_outputDirectory, "Result.png"));
             OutputHelpers.WriteConsole("Upscaling Complete.", ConsoleColor.Cyan);
         }

OnnxStack.StableDiffusion/Models/InputImage.cs renamed to OnnxStack.Core/Image/InputImage.cs

@@ -4,7 +4,7 @@
 using System.IO;
 using System.Text.Json.Serialization;
 
-namespace OnnxStack.StableDiffusion.Models
+namespace OnnxStack.Core.Image
 {
     public class InputImage
     {

OnnxStack.Core/OnnxStack.Core.csproj

@@ -41,6 +41,7 @@
     <PackageReference Include="Microsoft.ML" Version="2.0.1" />
     <PackageReference Include="Microsoft.ML.OnnxRuntime.Extensions" Version="0.9.0" />
     <PackageReference Include="Microsoft.ML.OnnxRuntime.Managed" Version="1.16.2" />
+    <PackageReference Include="SixLabors.ImageSharp" Version="3.1.0" />
   </ItemGroup>
 
   <ItemGroup>

OnnxStack.ImageUpscaler/Extensions/ImageExtensions.cs

@@ -1,14 +1,89 @@
 using Microsoft.ML.OnnxRuntime;
 using Microsoft.ML.OnnxRuntime.Tensors;
 using OnnxStack.Core;
+using OnnxStack.Core.Image;
+using OnnxStack.ImageUpscaler.Models;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
 using System;
+using System.Collections.Generic;
 
 namespace OnnxStack.ImageUpscaler.Extensions
 {
     internal static class ImageExtensions
     {
+        /// <summary>
+        /// Generates the image tiles.
+        /// </summary>
+        /// <param name="imageSource">The image source.</param>
+        /// <param name="sampleSize">Maximum size of the tile.</param>
+        /// <param name="scaleFactor">The scale factor.</param>
+        /// <returns></returns>
+        public static List<ImageTile> GenerateTiles(this Image<Rgba32> imageSource, int sampleSize, int scaleFactor)
+        {
+            var tiles = new List<ImageTile>();
+            var tileSizeX = Math.Min(sampleSize, imageSource.Width);
+            var tileSizeY = Math.Min(sampleSize, imageSource.Height);
+            var columns = (int)Math.Ceiling((double)imageSource.Width / tileSizeX);
+            var rows = (int)Math.Ceiling((double)imageSource.Height / tileSizeY);
+            var tileWidth = imageSource.Width / columns;
+            var tileHeight = imageSource.Height / rows;
+
+            for (int y = 0; y < rows; y++)
+            {
+                for (int x = 0; x < columns; x++)
+                {
+                    var tileRect = new Rectangle(x * tileWidth, y * tileHeight, tileWidth, tileHeight);
+                    var tileDest = new Rectangle(tileRect.X * scaleFactor, tileRect.Y * scaleFactor, tileWidth * scaleFactor, tileHeight * scaleFactor);
+                    var tileImage = ExtractTile(imageSource, tileRect);
+                    tiles.Add(new ImageTile(tileImage, tileDest));
+                }
+            }
+            return tiles;
+        }
+
+
+        /// <summary>
+        /// Extracts an image tile from a source image.
+        /// </summary>
+        /// <param name="sourceImage">The source image.</param>
+        /// <param name="sourceArea">The source area.</param>
+        /// <returns></returns>
+        public static Image<Rgba32> ExtractTile(this Image<Rgba32> sourceImage, Rectangle sourceArea)
+        {
+            var height = sourceArea.Height;
+            var targetImage = new Image<Rgba32>(sourceArea.Width, sourceArea.Height);
+            sourceImage.ProcessPixelRows(targetImage, (sourceAccessor, targetAccessor) =>
+            {
+                for (int i = 0; i < height; i++)
+                {
+                    var sourceRow = sourceAccessor.GetRowSpan(sourceArea.Y + i);
+                    var targetRow = targetAccessor.GetRowSpan(i);
+                    sourceRow.Slice(sourceArea.X, sourceArea.Width).CopyTo(targetRow);
+                }
+            });
+            return targetImage;
+        }
+
+
+        /// <summary>
+        /// Converts to InputImage to an Image<Rgba32>
+        /// </summary>
+        /// <param name="inputImage">The input image.</param>
+        /// <returns></returns>
+        public static Image<Rgba32> ToImage(this InputImage inputImage)
+        {
+            if (!string.IsNullOrEmpty(inputImage.ImageBase64))
+                return Image.Load<Rgba32>(Convert.FromBase64String(inputImage.ImageBase64.Split(',')[1]));
+            if (inputImage.ImageBytes != null)
+                return Image.Load<Rgba32>(inputImage.ImageBytes);
+            if (inputImage.ImageStream != null)
+                return Image.Load<Rgba32>(inputImage.ImageStream);
+            if (inputImage.ImageTensor != null)
+                return inputImage.ImageTensor.ToImage();
+
+            return inputImage.Image;
+        }
 
         /// <summary>
         /// Converts to DenseTensor.
@@ -20,7 +95,7 @@ public static DenseTensor<float> ToDenseTensor(this Image<Rgba32> image, ReadOnl
         {
             using (image)
             {
-                return ProcessPixels(image, dimensions);
+                return NormalizeToZeroToOne(image, dimensions);
             }
         }
 
@@ -51,23 +126,24 @@ public static Image<Rgba32> ToImage(this DenseTensor<float> imageTensor)
                 for (var x = 0; x < width; x++)
                 {
                     result[x, y] = new Rgba32(
-                        CalculateByte(imageTensor, 0, y, x),
-                        CalculateByte(imageTensor, 1, y, x),
-                        CalculateByte(imageTensor, 2, y, x)
+                        DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
+                        DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
+                        DenormalizeZeroToOneToByte(imageTensor, 2, y, x)
                     );
                 }
             }
             return result;
         }
 
 
+
         /// <summary>
-        /// Processes the pixels.
+        /// Normalizes the pixels from 0-255 to 0-1
         /// </summary>
         /// <param name="image">The image.</param>
         /// <param name="dimensions">The dimensions.</param>
         /// <returns></returns>
-        private static DenseTensor<float> ProcessPixels(Image<Rgba32> image, ReadOnlySpan<int> dimensions)
+        private static DenseTensor<float> NormalizeToZeroToOne(Image<Rgba32> image, ReadOnlySpan<int> dimensions)
         {
             var width = dimensions[3];
             var height = dimensions[2];
@@ -90,15 +166,16 @@ private static DenseTensor<float> ProcessPixels(Image<Rgba32> image, ReadOnlySpa
         }
 
 
+
         /// <summary>
-        /// Calculates the byte.
+        /// Denormalizes the pixels from 0-1 to 0-255
         /// </summary>
         /// <param name="imageTensor">The image tensor.</param>
         /// <param name="index">The index.</param>
         /// <param name="y">The y.</param>
         /// <param name="x">The x.</param>
         /// <returns></returns>
-        private static byte CalculateByte(Tensor<float> imageTensor, int index, int y, int x)
+        private static byte DenormalizeZeroToOneToByte(DenseTensor<float> imageTensor, int index, int y, int x)
         {
             return (byte)Math.Round(Math.Clamp(imageTensor[0, index, y, x], 0, 1) * 255);
         }

OnnxStack.ImageUpscaler/OnnxStack.ImageUpscaler.csproj

@@ -31,10 +31,6 @@
     </None>
   </ItemGroup>
 
-  <ItemGroup>
-    <PackageReference Include="SixLabors.ImageSharp" Version="3.0.2" />
-  </ItemGroup>
-
   <ItemGroup>
     <ProjectReference Include="..\OnnxStack.Core\OnnxStack.Core.csproj" />
   </ItemGroup>

OnnxStack.ImageUpscaler/Registration.cs

@@ -38,7 +38,6 @@ public static void AddOnnxStackImageUpscaler(this IServiceCollection serviceColl
         /// <param name="serviceCollection">The service collection.</param>
         private static void RegisterServices(this IServiceCollection serviceCollection)
         {
-            serviceCollection.AddSingleton<IImageService, ImageService>();
             serviceCollection.AddSingleton<IUpscaleService, UpscaleService>();
         }
 

OnnxStack.ImageUpscaler/Services/IImageService.cs

@@ -1,8 +1,11 @@
-using OnnxStack.ImageUpscaler.Config;
+using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core.Image;
+using OnnxStack.ImageUpscaler.Config;
 using OnnxStack.StableDiffusion.Config;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
 using System.Collections.Generic;
+using System.IO;
 using System.Threading.Tasks;
 
 namespace OnnxStack.ImageUpscaler.Services
@@ -44,6 +47,32 @@ public interface IUpscaleService
         /// <param name="modelOptions">The model options.</param>
         /// <param name="inputImage">The input image.</param>
         /// <returns></returns>
-        Task<Image<Rgba32>> GenerateAsync(UpscaleModelSet modelOptions, Image<Rgba32> inputImage);
+        Task<DenseTensor<float>> GenerateAsync(UpscaleModelSet modelOptions, InputImage inputImage);
+
+        /// <summary>
+        /// Generates the upscaled image.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <param name="inputImage">The input image.</param>
+        /// <returns></returns>
+        Task<Image<Rgba32>> GenerateAsImageAsync(UpscaleModelSet modelOptions, InputImage inputImage);
+
+
+        /// <summary>
+        /// Generates the upscaled image.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <param name="inputImage">The input image.</param>
+        /// <returns></returns>
+        Task<byte[]> GenerateAsByteAsync(UpscaleModelSet modelOptions, InputImage inputImage);
+
+
+        /// <summary>
+        /// Generates the upscaled image.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <param name="inputImage">The input image.</param>
+        /// <returns></returns>
+        Task<Stream> GenerateAsStreamAsync(UpscaleModelSet modelOptions, InputImage inputImage);
     }
 }