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

Commit d6267e6

Browse files
saddam213saddam213
committed
Update ImageUpscale console example
1 parent 1f6d747 commit d6267e6

File tree

8 files changed

+229
-61
lines changed

8 files changed

+229
-61
lines changed

‎OnnxStack.Console/Examples/UpscaleExample.cs

Lines changed: 15 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,18 @@
11
using OnnxStack.Core.Image;
2+
using OnnxStack.FeatureExtractor.Pipelines;
23
using OnnxStack.ImageUpscaler.Config;
3-
using OnnxStack.ImageUpscaler.Services;
44
using SixLabors.ImageSharp;
5-
using SixLabors.ImageSharp.PixelFormats;
65

76
namespace OnnxStack.Console.Runner
87
{
98
public sealed class UpscaleExample : IExampleRunner
109
{
1110
private readonly string _outputDirectory;
1211
private readonly ImageUpscalerConfig _configuration;
13-
private readonly IUpscaleService _imageUpscaleService;
1412

15-
16-
public UpscaleExample(ImageUpscalerConfig configuration, IUpscaleService imageUpscaleService)
13+
public UpscaleExample(ImageUpscalerConfig configuration)
1714
{
1815
_configuration = configuration;
19-
_imageUpscaleService = imageUpscaleService;
2016
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(UpscaleExample));
2117
Directory.CreateDirectory(_outputDirectory);
2218
}
@@ -29,28 +25,24 @@ public UpscaleExample(ImageUpscalerConfig configuration, IUpscaleService imageUp
2925

3026
public async Task RunAsync()
3127
{
32-
var modelSet = _configuration.ModelSets.FirstOrDefault(x => x.Name == "RealSR BSRGAN x4");
33-
28+
// Load Input Image
29+
var inputImage = await InputImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
3430

31+
// Create Pipeline
32+
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
3533

36-
OutputHelpers.WriteConsole("Enter Image Path", ConsoleColor.Yellow);
37-
var imageFile = OutputHelpers.ReadConsole(ConsoleColor.Gray);
38-
if (!File.Exists(imageFile))
39-
{
40-
OutputHelpers.WriteConsole("File not found!", ConsoleColor.Red);
41-
return;
42-
}
34+
// Run pipeline
35+
var result = await pipeline.RunAsync(inputImage);
4336

44-
OutputHelpers.WriteConsole("Loading Model...", ConsoleColor.Cyan);
45-
await _imageUpscaleService.LoadModelAsync(modelSet);
46-
OutputHelpers.WriteConsole("Model Loaded.", ConsoleColor.Cyan);
37+
// Create Image from Tensor result
38+
var image = result.ToImage(ImageNormalizeType.ZeroToOne);
4739

48-
var inputImage = await Image.LoadAsync<Rgba32>(imageFile);
40+
// Save Image File
41+
var outputFilename = Path.Combine(_outputDirectory, $"Upscaled.png");
42+
await image.SaveAsPngAsync(outputFilename);
4943

50-
OutputHelpers.WriteConsole("Upscaling Image...", ConsoleColor.Cyan);
51-
var result = await _imageUpscaleService.GenerateAsImageAsync(modelSet, new InputImage(inputImage));
52-
await result.SaveAsPngAsync(Path.Combine(_outputDirectory, "Result.png"));
53-
OutputHelpers.WriteConsole("Upscaling Complete.", ConsoleColor.Cyan);
44+
// Unload
45+
await pipeline.UnloadAsync();
5446
}
5547

5648
}

‎OnnxStack.ImageUpscaler/Common/UpscaleModel.cs

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
using Microsoft.ML.OnnxRuntime;
2+
using OnnxStack.Core.Config;
3+
using OnnxStack.Core.Model;
4+
5+
namespace OnnxStack.ImageUpscaler.Common
6+
{
7+
public class UpscaleModel : OnnxModelSession
8+
{
9+
private readonly int _channels;
10+
private readonly int _sampleSize;
11+
private readonly int _scaleFactor;
12+
13+
public UpscaleModel(UpscaleModelConfig configuration) : base(configuration)
14+
{
15+
_channels = configuration.Channels;
16+
_sampleSize = configuration.SampleSize;
17+
_scaleFactor = configuration.ScaleFactor;
18+
}
19+
20+
public int Channels => _channels;
21+
public int SampleSize => _sampleSize;
22+
public int ScaleFactor => _scaleFactor;
23+
24+
25+
public static UpscaleModel Create(UpscaleModelConfig configuration)
26+
{
27+
return new UpscaleModel(configuration);
28+
}
29+
30+
public static UpscaleModel Create(string modelFile, int scaleFactor, int sampleSize = 512, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
31+
{
32+
var configuration = new UpscaleModelConfig
33+
{
34+
Channels = 3,
35+
SampleSize = sampleSize,
36+
ScaleFactor = scaleFactor,
37+
DeviceId = deviceId,
38+
ExecutionProvider = executionProvider,
39+
ExecutionMode = ExecutionMode.ORT_SEQUENTIAL,
40+
InterOpNumThreads = 0,
41+
IntraOpNumThreads = 0,
42+
OnnxModelPath = modelFile
43+
};
44+
return new UpscaleModel(configuration);
45+
}
46+
}
47+
48+
49+
public record UpscaleModelConfig : OnnxModelConfig
50+
{
51+
public int Channels { get; set; }
52+
public int SampleSize { get; set; }
53+
public int ScaleFactor { get; set; }
54+
}
55+
}

‎OnnxStack.ImageUpscaler/Config/UpscaleModelSet.cs renamed to ‎OnnxStack.ImageUpscaler/Common/UpscaleModelSet.cs

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
11
using Microsoft.ML.OnnxRuntime;
22
using OnnxStack.Core.Config;
3-
using OnnxStack.ImageUpscaler.Models;
4-
using System.Collections.Generic;
53
using System.Text.Json.Serialization;
64

7-
namespace OnnxStack.StableDiffusion.Config
5+
namespace OnnxStack.ImageUpscaler.Common
86
{
97
public record UpscaleModelSet : IOnnxModelSetConfig
108
{
119
public string Name { get; set; }
12-
public bool IsEnabled { get; set; }
10+
public bool IsEnabled { get; set; }
1311
public int DeviceId { get; set; }
1412
public int InterOpNumThreads { get; set; }
1513
public int IntraOpNumThreads { get; set; }

‎OnnxStack.ImageUpscaler/Config/ImageUpscalerConfig.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
using OnnxStack.Common.Config;
22
using OnnxStack.Core;
3-
using OnnxStack.StableDiffusion.Config;
3+
using OnnxStack.ImageUpscaler.Common;
44
using System.Collections.Generic;
55

66
namespace OnnxStack.ImageUpscaler.Config

‎OnnxStack.ImageUpscaler/Models/UpscaleModel.cs

Lines changed: 0 additions & 31 deletions
This file was deleted.

‎OnnxStack.ImageUpscaler/Pipelines/ImageUpscalePipeline.cs

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
using Microsoft.Extensions.Logging;
2+
using Microsoft.ML.OnnxRuntime.Tensors;
3+
using OnnxStack.Core;
4+
using OnnxStack.Core.Config;
5+
using OnnxStack.Core.Image;
6+
using OnnxStack.Core.Model;
7+
using OnnxStack.ImageUpscaler.Common;
8+
using OnnxStack.ImageUpscaler.Extensions;
9+
using OnnxStack.ImageUpscaler.Models;
10+
using SixLabors.ImageSharp;
11+
using SixLabors.ImageSharp.PixelFormats;
12+
using System.IO;
13+
using System.Linq;
14+
using System.Threading;
15+
using System.Threading.Tasks;
16+
17+
namespace OnnxStack.FeatureExtractor.Pipelines
18+
{
19+
public class ImageUpscalePipeline
20+
{
21+
private readonly string _name;
22+
private readonly ILogger _logger;
23+
private readonly UpscaleModel _upscaleModel;
24+
25+
26+
/// <summary>
27+
/// Initializes a new instance of the <see cref="ImageUpscalePipeline"/> class.
28+
/// </summary>
29+
/// <param name="name">The name.</param>
30+
/// <param name="upscaleModel">The upscale model.</param>
31+
/// <param name="logger">The logger.</param>
32+
public ImageUpscalePipeline(string name, UpscaleModel upscaleModel, ILogger logger = default)
33+
{
34+
_name = name;
35+
_logger = logger;
36+
_upscaleModel = upscaleModel;
37+
}
38+
39+
40+
/// <summary>
41+
/// Gets the name.
42+
/// </summary>
43+
/// <value>
44+
public string Name => _name;
45+
46+
47+
/// <summary>
48+
/// Loads the model.
49+
/// </summary>
50+
public async Task LoadAsync()
51+
{
52+
await _upscaleModel.LoadAsync();
53+
}
54+
55+
56+
/// <summary>
57+
/// Unloads the models.
58+
/// </summary>
59+
public async Task UnloadAsync()
60+
{
61+
await Task.Yield();
62+
_upscaleModel?.Dispose();
63+
}
64+
65+
66+
/// <summary>
67+
/// Runs the upscale pipeline.
68+
/// </summary>
69+
/// <param name="inputImage">The input image.</param>
70+
/// <param name="cancellationToken">The cancellation token.</param>
71+
/// <returns></returns>
72+
public async Task<DenseTensor<float>> RunAsync(InputImage inputImage, CancellationToken cancellationToken = default)
73+
{
74+
using (var image = await inputImage.ToImageAsync())
75+
{
76+
var upscaleInput = CreateInputParams(image, _upscaleModel.SampleSize, _upscaleModel.ScaleFactor);
77+
var metadata = await _upscaleModel.GetMetadataAsync();
78+
79+
var outputTensor = new DenseTensor<float>(new[] { 1, _upscaleModel.Channels, upscaleInput.OutputHeight, upscaleInput.OutputWidth });
80+
foreach (var imageTile in upscaleInput.ImageTiles)
81+
{
82+
cancellationToken.ThrowIfCancellationRequested();
83+
84+
var outputDimension = new[] { 1, _upscaleModel.Channels, imageTile.Destination.Height, imageTile.Destination.Width };
85+
var inputTensor = imageTile.Image.ToDenseTensor(ImageNormalizeType.ZeroToOne, _upscaleModel.Channels);
86+
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
87+
{
88+
inferenceParameters.AddInputTensor(inputTensor);
89+
inferenceParameters.AddOutputBuffer(outputDimension);
90+
91+
var results = await _upscaleModel.RunInferenceAsync(inferenceParameters);
92+
using (var result = results.First())
93+
{
94+
outputTensor.ApplyImageTile(result.ToDenseTensor(), imageTile.Destination);
95+
}
96+
}
97+
}
98+
99+
return outputTensor;
100+
}
101+
}
102+
103+
104+
private static UpscaleInput CreateInputParams(Image<Rgba32> imageSource, int maxTileSize, int scaleFactor)
105+
{
106+
var tiles = imageSource.GenerateTiles(maxTileSize, scaleFactor);
107+
var width = imageSource.Width * scaleFactor;
108+
var height = imageSource.Height * scaleFactor;
109+
return new UpscaleInput(tiles, width, height);
110+
}
111+
112+
113+
/// <summary>
114+
/// Creates the pipeline from a UpscaleModelSet.
115+
/// </summary>
116+
/// <param name="modelSet">The model set.</param>
117+
/// <param name="logger">The logger.</param>
118+
/// <returns></returns>
119+
public static ImageUpscalePipeline CreatePipeline(UpscaleModelSet modelSet, ILogger logger = default)
120+
{
121+
var upscaleModel = new UpscaleModel(modelSet.UpscaleModelConfig.ApplyDefaults(modelSet));
122+
return new ImageUpscalePipeline(modelSet.Name, upscaleModel, logger);
123+
}
124+
125+
126+
/// <summary>
127+
/// Creates the pipeline from the specified folder.
128+
/// </summary>
129+
/// <param name="modelFolder">The model folder.</param>
130+
/// <param name="deviceId">The device identifier.</param>
131+
/// <param name="executionProvider">The execution provider.</param>
132+
/// <param name="logger">The logger.</param>
133+
/// <returns></returns>
134+
public static ImageUpscalePipeline CreatePipeline(string modelFile, int scaleFactor, int sampleSize = 512, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML, ILogger logger = default)
135+
{
136+
var name = Path.GetFileNameWithoutExtension(modelFile);
137+
var configuration = new UpscaleModelSet
138+
{
139+
Name = name,
140+
IsEnabled = true,
141+
DeviceId = deviceId,
142+
ExecutionProvider = executionProvider,
143+
UpscaleModelConfig = new UpscaleModelConfig
144+
{
145+
Channels = 3,
146+
SampleSize = sampleSize,
147+
ScaleFactor = scaleFactor,
148+
OnnxModelPath = modelFile
149+
}
150+
};
151+
return CreatePipeline(configuration, logger);
152+
}
153+
}
154+
}

‎OnnxStack.ImageUpscaler/Services/IUpscaleService.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
using Microsoft.ML.OnnxRuntime.Tensors;
22
using OnnxStack.Core.Image;
33
using OnnxStack.Core.Video;
4-
using OnnxStack.StableDiffusion.Config;
4+
using OnnxStack.ImageUpscaler.Common;
55
using SixLabors.ImageSharp;
66
using SixLabors.ImageSharp.PixelFormats;
77
using System.IO;

‎OnnxStack.ImageUpscaler/Services/UpscaleService.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@
55
using OnnxStack.Core.Model;
66
using OnnxStack.Core.Services;
77
using OnnxStack.Core.Video;
8+
using OnnxStack.ImageUpscaler.Common;
89
using OnnxStack.ImageUpscaler.Config;
910
using OnnxStack.ImageUpscaler.Extensions;
1011
using OnnxStack.ImageUpscaler.Models;
11-
using OnnxStack.StableDiffusion.Config;
1212
using SixLabors.ImageSharp;
1313
using SixLabors.ImageSharp.PixelFormats;
1414
using System.Collections.Generic;

0 commit comments

Comments
 (0)