Skip to content
This repository was archived by the owner on Nov 27, 2024. It is now read-only.

Commit c80f7bd

Browse files
saddam213saddam213
committed
Upscale service prototype
1 parent 5ba9c1e commit c80f7bd

File tree

13 files changed

+380
-0
lines changed

13 files changed

+380
-0
lines changed

‎OnnxStack.Console/Examples/UpscaleExample.cs‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
using OnnxStack.Core;
2+
using OnnxStack.Core.Config;
3+
using OnnxStack.Core.Services;
4+
using OnnxStack.ImageUpscaler.Services;
5+
using SixLabors.ImageSharp;
6+
using SixLabors.ImageSharp.PixelFormats;
7+
8+
namespace OnnxStack.Console.Runner
9+
{
10+
public sealed class UpscaleExample : IExampleRunner
11+
{
12+
private readonly string _outputDirectory;
13+
private readonly IOnnxModelService _modelService;
14+
private readonly IUpscaleService _imageUpscaleService;
15+
16+
17+
public UpscaleExample(IOnnxModelService modelService, IUpscaleService imageUpscaleService)
18+
{
19+
_modelService = modelService;
20+
_imageUpscaleService = imageUpscaleService;
21+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(UpscaleExample));
22+
Directory.CreateDirectory(_outputDirectory);
23+
}
24+
25+
public string Name => "Image Upscale Demo";
26+
27+
public string Description => "Upscales images";
28+
29+
public async Task RunAsync()
30+
{
31+
32+
var modelset = new OnnxModelSetConfig
33+
{
34+
Name = "Upscaler",
35+
IsEnabled = true,
36+
ExecutionProvider = ExecutionProvider.DirectML,
37+
ModelConfigurations = new List<OnnxModelSessionConfig>
38+
{
39+
new OnnxModelSessionConfig
40+
{
41+
Type = OnnxModelType.Upscaler,
42+
OnnxModelPath = "D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx"
43+
}
44+
}
45+
};
46+
modelset.ApplyConfigurationOverrides();
47+
_modelService.UpdateModelSet(modelset);
48+
49+
OutputHelpers.WriteConsole("Enter Image Path", ConsoleColor.Yellow);
50+
var imageFile = OutputHelpers.ReadConsole(ConsoleColor.Gray);
51+
if (!File.Exists(imageFile))
52+
{
53+
OutputHelpers.WriteConsole("File not found!", ConsoleColor.Red);
54+
return;
55+
}
56+
57+
OutputHelpers.WriteConsole("Loading Model...", ConsoleColor.Cyan);
58+
await _modelService.LoadModelAsync(modelset);
59+
OutputHelpers.WriteConsole("Model Loaded.", ConsoleColor.Cyan);
60+
61+
var inputImage = await Image.LoadAsync<Rgba32>(imageFile);
62+
63+
OutputHelpers.WriteConsole("Upscaling Image...", ConsoleColor.Cyan);
64+
var result = await _imageUpscaleService.GenerateAsync(modelset, inputImage);
65+
await result.SaveAsPngAsync(Path.Combine(_outputDirectory, "Result.png"));
66+
OutputHelpers.WriteConsole("Upscaling Complete.", ConsoleColor.Cyan);
67+
}
68+
69+
}
70+
}

‎OnnxStack.Console/OnnxStack.Console.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
</ItemGroup>
1616

1717
<ItemGroup>
18+
<ProjectReference Include="..\OnnxStack.ImageUpscaler\OnnxStack.ImageUpscaler.csproj" />
1819
<ProjectReference Include="..\OnnxStack.StableDiffusion\OnnxStack.StableDiffusion.csproj" />
1920
</ItemGroup>
2021

‎OnnxStack.Console/Program.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
using Microsoft.Extensions.Hosting;
33
using Microsoft.Extensions.Logging;
44
using OnnxStack.Core;
5+
using OnnxStack.ImageUpscaler;
6+
using OnnxStack.ImageUpscaler.Config;
57
using System.Reflection;
68

79
namespace OnnxStack.Console
@@ -17,6 +19,7 @@ static async Task Main(string[] _)
1719
// Add OnnxStack
1820
builder.Services.AddOnnxStack();
1921
builder.Services.AddOnnxStackStableDiffusion();
22+
builder.Services.AddOnnxStackImageUpscaler(new ImageUpscalerConfig());
2023

2124
// Add AppService
2225
builder.Services.AddHostedService<AppService>();

‎OnnxStack.Core/Config/OnnxModelType.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,5 +9,7 @@ public enum OnnxModelType
99
TextEncoder2 = 21,
1010
VaeEncoder = 30,
1111
VaeDecoder = 40,
12+
13+
Upscaler = 100
1214
}
1315
}

‎OnnxStack.ImageUpscaler/Extensions/ImageExtensions.cs‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
using Microsoft.ML.OnnxRuntime;
2+
using Microsoft.ML.OnnxRuntime.Tensors;
3+
using OnnxStack.Core;
4+
using SixLabors.ImageSharp;
5+
using SixLabors.ImageSharp.PixelFormats;
6+
using System;
7+
8+
namespace OnnxStack.ImageUpscaler.Extensions
9+
{
10+
internal static class ImageExtensions
11+
{
12+
13+
/// <summary>
14+
/// Converts to DenseTensor.
15+
/// </summary>
16+
/// <param name="image">The image.</param>
17+
/// <param name="dimensions">The dimensions.</param>
18+
/// <returns></returns>
19+
public static DenseTensor<float> ToDenseTensor(this Image<Rgba32> image, ReadOnlySpan<int> dimensions)
20+
{
21+
using (image)
22+
{
23+
return ProcessPixels(image, dimensions);
24+
}
25+
}
26+
27+
28+
/// <summary>
29+
/// Converts to Image.
30+
/// </summary>
31+
/// <param name="ortValue">The ort value.</param>
32+
/// <returns></returns>
33+
public static Image<Rgba32> ToImage(this OrtValue ortValue)
34+
{
35+
return ortValue.ToDenseTensor().ToImage();
36+
}
37+
38+
39+
/// <summary>
40+
/// Converts to image.
41+
/// </summary>
42+
/// <param name="imageTensor">The image tensor.</param>
43+
/// <returns></returns>
44+
public static Image<Rgba32> ToImage(this DenseTensor<float> imageTensor)
45+
{
46+
var height = imageTensor.Dimensions[2];
47+
var width = imageTensor.Dimensions[3];
48+
var result = new Image<Rgba32>(width, height);
49+
for (var y = 0; y < height; y++)
50+
{
51+
for (var x = 0; x < width; x++)
52+
{
53+
result[x, y] = new Rgba32(
54+
CalculateByte(imageTensor, 0, y, x),
55+
CalculateByte(imageTensor, 1, y, x),
56+
CalculateByte(imageTensor, 2, y, x)
57+
);
58+
}
59+
}
60+
return result;
61+
}
62+
63+
64+
/// <summary>
65+
/// Processes the pixels.
66+
/// </summary>
67+
/// <param name="image">The image.</param>
68+
/// <param name="dimensions">The dimensions.</param>
69+
/// <returns></returns>
70+
private static DenseTensor<float> ProcessPixels(Image<Rgba32> image, ReadOnlySpan<int> dimensions)
71+
{
72+
var width = dimensions[3];
73+
var height = dimensions[2];
74+
var channels = dimensions[1];
75+
var imageArray = new DenseTensor<float>(new[] { 1, channels, height, width });
76+
image.ProcessPixelRows(img =>
77+
{
78+
for (int x = 0; x < width; x++)
79+
{
80+
for (int y = 0; y < height; y++)
81+
{
82+
var pixelSpan = img.GetRowSpan(y);
83+
imageArray[0, 0, y, x] = (pixelSpan[x].R / 255.0f);
84+
imageArray[0, 1, y, x] = (pixelSpan[x].G / 255.0f);
85+
imageArray[0, 2, y, x] = (pixelSpan[x].B / 255.0f);
86+
}
87+
}
88+
});
89+
return imageArray;
90+
}
91+
92+
93+
/// <summary>
94+
/// Calculates the byte.
95+
/// </summary>
96+
/// <param name="imageTensor">The image tensor.</param>
97+
/// <param name="index">The index.</param>
98+
/// <param name="y">The y.</param>
99+
/// <param name="x">The x.</param>
100+
/// <returns></returns>
101+
private static byte CalculateByte(Tensor<float> imageTensor, int index, int y, int x)
102+
{
103+
return (byte)Math.Round(Math.Clamp(imageTensor[0, index, y, x], 0, 1) * 255);
104+
}
105+
}
106+
}

‎OnnxStack.ImageUpscaler/Models/ImageTile.cs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
using SixLabors.ImageSharp;
2+
using SixLabors.ImageSharp.PixelFormats;
3+
4+
namespace OnnxStack.ImageUpscaler.Models
5+
{
6+
public record ImageTile(Image<Rgba32> Image, Rectangle Destination);
7+
}
8+
9+
10+

‎OnnxStack.ImageUpscaler/Models/UpscaleInput.cs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
using System.Collections.Generic;
2+
3+
namespace OnnxStack.ImageUpscaler.Models
4+
{
5+
public record UpscaleInput(List<ImageTile> ImageTiles, int OutputWidth, int OutputHeight);
6+
}
7+
8+
9+

‎OnnxStack.ImageUpscaler/OnnxStack.ImageUpscaler.csproj‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@
3131
</None>
3232
</ItemGroup>
3333

34+
<ItemGroup>
35+
<PackageReference Include="SixLabors.ImageSharp" Version="3.0.2" />
36+
</ItemGroup>
37+
3438
<ItemGroup>
3539
<ProjectReference Include="..\OnnxStack.Core\OnnxStack.Core.csproj" />
3640
</ItemGroup>

‎OnnxStack.ImageUpscaler/Registration.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using OnnxStack.Core;
33
using OnnxStack.Core.Config;
44
using OnnxStack.ImageUpscaler.Config;
5+
using OnnxStack.ImageUpscaler.Services;
56

67
namespace OnnxStack.ImageUpscaler
78
{
@@ -33,6 +34,8 @@ public static void AddOnnxStackImageUpscaler(this IServiceCollection serviceColl
3334

3435
private static void RegisterServices(this IServiceCollection serviceCollection)
3536
{
37+
serviceCollection.AddSingleton<IImageService, ImageService>();
38+
serviceCollection.AddSingleton<IUpscaleService, UpscaleService>();
3639
}
3740
}
3841
}

‎OnnxStack.ImageUpscaler/Services/IImageService.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
using OnnxStack.ImageUpscaler.Models;
2+
using SixLabors.ImageSharp;
3+
using SixLabors.ImageSharp.PixelFormats;
4+
using System.Collections.Generic;
5+
6+
namespace OnnxStack.ImageUpscaler.Services
7+
{
8+
public interface IImageService
9+
{
10+
List<ImageTile> GenerateTiles(Image<Rgba32> imageSource, int maxTileSize, int scaleFactor);
11+
}
12+
}

0 commit comments

Comments
 (0)