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

Commit 6f263b4

Browse files
saddam213saddam213
authored
Merge pull request #133 from saddam213/Upscale
Upscale Pipeline Improvments
2 parents b69fbf3 + 8695851 commit 6f263b4

21 files changed

+453
-312
lines changed

‎OnnxStack.Console/Examples/ControlNetFeatureExample.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ public async Task RunAsync()
3535
var inputImage = await OnnxImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
3636

3737
// Create Annotation pipeline
38-
var annotationPipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutputTensor: true);
38+
var annotationPipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutput: true);
3939

4040
// Create Depth Image
4141
var controlImage = await annotationPipeline.RunAsync(inputImage);

‎OnnxStack.Console/Examples/FeatureExtractorExample.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ public async Task RunAsync()
3737
{
3838
FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\canny.onnx"),
3939
FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\hed.onnx"),
40-
FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutputTensor: true, inputResizeMode: ImageResizeMode.Stretch),
40+
FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\depth.onnx", sampleSize: 512, normalizeOutput: true, inputResizeMode: ImageResizeMode.Stretch),
4141
FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024, setOutputToInputAlpha: true, inputResizeMode: ImageResizeMode.Stretch)
4242
};
4343

‎OnnxStack.Console/Examples/UpscaleExample.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ public async Task RunAsync()
2525
var inputImage = await OnnxImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
2626

2727
// Create Pipeline
28-
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
28+
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4, 512);
2929

3030
// Run pipeline
3131
var result = await pipeline.RunAsync(inputImage);

‎OnnxStack.Console/Examples/UpscaleStreamExample.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ public async Task RunAsync()
2626
var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
2727

2828
// Create pipeline
29-
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
29+
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4, 512);
3030

3131
// Load pipeline
3232
await pipeline.LoadAsync();

‎OnnxStack.Core/Extensions/TensorExtension.cs

Lines changed: 9 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Numerics.Tensors;
66
using System.Numerics;
77
using OnnxStack.Core.Model;
8+
using System.Threading.Tasks;
89

910
namespace OnnxStack.Core
1011
{
@@ -415,75 +416,22 @@ private static DenseTensor<float> ConcatenateAxis2(DenseTensor<float> tensor1, D
415416

416417

417418
/// <summary>
418-
/// Splits the Tensor into 4 equal tiles.
419+
/// Normalizes the tensor values from range -1 to 1 to 0 to 1.
419420
/// </summary>
420-
/// <param name="sourceTensor">The source tensor.</param>
421-
/// <returns>TODO: Optimize</returns>
422-
public static ImageTiles SplitTiles(this DenseTensor<float> sourceTensor)
421+
/// <param name="imageTensor">The image tensor.</param>
422+
public static void NormalizeOneOneToZeroOne(this DenseTensor<float> imageTensor)
423423
{
424-
int tileWidth = sourceTensor.Dimensions[3] / 2;
425-
int tileHeight = sourceTensor.Dimensions[2] / 2;
426-
427-
return new ImageTiles(
428-
SplitTile(sourceTensor, 0, 0, tileHeight, tileWidth),
429-
SplitTile(sourceTensor, 0, tileWidth, tileHeight, tileWidth * 2),
430-
SplitTile(sourceTensor, tileHeight, 0, tileHeight * 2, tileWidth),
431-
SplitTile(sourceTensor, tileHeight, tileWidth, tileHeight * 2, tileWidth * 2));
432-
}
433-
434-
private static DenseTensor<float> SplitTile(DenseTensor<float> tensor, int startRow, int startCol, int endRow, int endCol)
435-
{
436-
int height = endRow - startRow;
437-
int width = endCol - startCol;
438-
int channels = tensor.Dimensions[1];
439-
var slicedData = new DenseTensor<float>(new[] { 1, channels, height, width });
440-
for (int c = 0; c < channels; c++)
441-
{
442-
for (int i = 0; i < height; i++)
443-
{
444-
for (int j = 0; j < width; j++)
445-
{
446-
slicedData[0, c, i, j] = tensor[0, c, startRow + i, startCol + j];
447-
}
448-
}
449-
}
450-
return slicedData;
424+
Parallel.For(0, (int)imageTensor.Length, (i) => imageTensor.SetValue(i, imageTensor.GetValue(i) / 2f + 0.5f));
451425
}
452426

453427

454428
/// <summary>
455-
/// Rejoins the tiles into a single Tensor.
429+
/// Normalizes the tensor values from range 0 to 1 to -1 to 1.
456430
/// </summary>
457-
/// <param name="tiles">The tiles.</param>
458-
/// <returns>TODO: Optimize</returns>
459-
public static DenseTensor<float> RejoinTiles(this ImageTiles tiles)
431+
/// <param name="imageTensor">The image tensor.</param>
432+
public static void NormalizeZeroOneToOneOne(this DenseTensor<float> imageTensor)
460433
{
461-
int totalHeight = tiles.Tile1.Dimensions[2] + tiles.Tile3.Dimensions[2];
462-
int totalWidth = tiles.Tile1.Dimensions[3] + tiles.Tile2.Dimensions[3];
463-
int channels = tiles.Tile1.Dimensions[1];
464-
var destination = new DenseTensor<float>(new[] { 1, channels, totalHeight, totalWidth });
465-
RejoinTile(destination, tiles.Tile1, 0, 0);
466-
RejoinTile(destination, tiles.Tile2, 0, tiles.Tile1.Dimensions[3]);
467-
RejoinTile(destination, tiles.Tile3, tiles.Tile1.Dimensions[2], 0);
468-
RejoinTile(destination, tiles.Tile4, tiles.Tile1.Dimensions[2], tiles.Tile1.Dimensions[3]);
469-
return destination;
470-
}
471-
472-
private static void RejoinTile(DenseTensor<float> destination, DenseTensor<float> tile, int startRow, int startCol)
473-
{
474-
int channels = tile.Dimensions[1];
475-
int height = tile.Dimensions[2];
476-
int width = tile.Dimensions[3];
477-
for (int c = 0; c < channels; c++)
478-
{
479-
for (int i = 0; i < height; i++)
480-
{
481-
for (int j = 0; j < width; j++)
482-
{
483-
destination[0, c, startRow + i, startCol + j] = tile[0, c, i, j];
484-
}
485-
}
486-
}
434+
Parallel.For(0, (int)imageTensor.Length, (i) => imageTensor.SetValue(i, 2f * imageTensor.GetValue(i) - 1f));
487435
}
488436
}
489437
}

‎OnnxStack.Core/Image/Extensions.cs

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
using Microsoft.ML.OnnxRuntime.Tensors;
2+
using OnnxStack.Core.Model;
23
using SixLabors.ImageSharp;
34
using SixLabors.ImageSharp.PixelFormats;
45
using SixLabors.ImageSharp.Processing;
6+
using System.Threading.Tasks;
57

68
namespace OnnxStack.Core.Image
79
{
@@ -14,6 +16,17 @@ public static class Extensions
1416
/// <param name="imageTensor">The image tensor.</param>
1517
/// <returns></returns>
1618
public static OnnxImage ToImageMask(this DenseTensor<float> imageTensor)
19+
{
20+
return new OnnxImage(imageTensor.FromMaskTensor());
21+
}
22+
23+
24+
/// <summary>
25+
/// Convert from single channle mask tensor to Rgba32 (Greyscale)
26+
/// </summary>
27+
/// <param name="imageTensor">The image tensor.</param>
28+
/// <returns></returns>
29+
public static Image<Rgba32> FromMaskTensor(this DenseTensor<float> imageTensor)
1730
{
1831
var width = imageTensor.Dimensions[3];
1932
var height = imageTensor.Dimensions[2];
@@ -26,7 +39,7 @@ public static OnnxImage ToImageMask(this DenseTensor<float> imageTensor)
2639
result[x, y] = new L8((byte)(imageTensor[0, 0, y, x] * 255.0f));
2740
}
2841
}
29-
return new OnnxImage(result.CloneAs<Rgba32>());
42+
return result.CloneAs<Rgba32>();
3043
}
3144
}
3245

@@ -40,6 +53,104 @@ public static ResizeMode ToResizeMode(this ImageResizeMode resizeMode)
4053
};
4154
}
4255

56+
57+
/// <summary>
58+
/// Splits the Tensor into 4 equal tiles.
59+
/// </summary>
60+
/// <param name="sourceTensor">The source tensor.</param>
61+
/// <returns>TODO: Optimize</returns>
62+
public static ImageTiles SplitImageTiles(this DenseTensor<float> sourceTensor, int overlap = 20)
63+
{
64+
var tileWidth = sourceTensor.Dimensions[3] / 2;
65+
var tileHeight = sourceTensor.Dimensions[2] / 2;
66+
return new ImageTiles(tileWidth, tileHeight, overlap,
67+
SplitImageTile(sourceTensor, 0, 0, tileHeight + overlap, tileWidth + overlap),
68+
SplitImageTile(sourceTensor, 0, tileWidth - overlap, tileHeight + overlap, tileWidth * 2),
69+
SplitImageTile(sourceTensor, tileHeight - overlap, 0, tileHeight * 2, tileWidth + overlap),
70+
SplitImageTile(sourceTensor, tileHeight - overlap, tileWidth - overlap, tileHeight * 2, tileWidth * 2));
71+
}
72+
73+
74+
/// <summary>
75+
/// Splits a tile from the source.
76+
/// </summary>
77+
/// <param name="source">The tensor.</param>
78+
/// <param name="startRow">The start row.</param>
79+
/// <param name="startCol">The start col.</param>
80+
/// <param name="endRow">The end row.</param>
81+
/// <param name="endCol">The end col.</param>
82+
/// <returns></returns>
83+
private static DenseTensor<float> SplitImageTile(DenseTensor<float> source, int startRow, int startCol, int endRow, int endCol)
84+
{
85+
int height = endRow - startRow;
86+
int width = endCol - startCol;
87+
int channels = source.Dimensions[1];
88+
var splitTensor = new DenseTensor<float>(new[] { 1, channels, height, width });
89+
Parallel.For(0, channels, (c) =>
90+
{
91+
Parallel.For(0, height, (i) =>
92+
{
93+
Parallel.For(0, width, (j) =>
94+
{
95+
splitTensor[0, c, i, j] = source[0, c, startRow + i, startCol + j];
96+
});
97+
});
98+
});
99+
return splitTensor;
100+
}
101+
102+
103+
/// <summary>
104+
/// Joins the tiles into a single Tensor.
105+
/// </summary>
106+
/// <param name="tiles">The tiles.</param>
107+
/// <returns>TODO: Optimize</returns>
108+
public static DenseTensor<float> JoinImageTiles(this ImageTiles tiles)
109+
{
110+
var totalWidth = tiles.Width * 2;
111+
var totalHeight = tiles.Height * 2;
112+
var channels = tiles.Tile1.Dimensions[1];
113+
var destination = new DenseTensor<float>(new[] { 1, channels, totalHeight, totalWidth });
114+
JoinImageTile(destination, tiles.Tile1, 0, 0, tiles.Height + tiles.Overlap, tiles.Width + tiles.Overlap);
115+
JoinImageTile(destination, tiles.Tile2, 0, tiles.Width - tiles.Overlap, tiles.Height + tiles.Overlap, totalWidth);
116+
JoinImageTile(destination, tiles.Tile3, tiles.Height - tiles.Overlap, 0, totalHeight, tiles.Width + tiles.Overlap);
117+
JoinImageTile(destination, tiles.Tile4, tiles.Height - tiles.Overlap, tiles.Width - tiles.Overlap, totalHeight, totalWidth);
118+
return destination;
119+
}
120+
121+
122+
/// <summary>
123+
/// Joins the tile to the destination tensor.
124+
/// </summary>
125+
/// <param name="destination">The destination.</param>
126+
/// <param name="tile">The tile.</param>
127+
/// <param name="startRow">The start row.</param>
128+
/// <param name="startCol">The start col.</param>
129+
/// <param name="endRow">The end row.</param>
130+
/// <param name="endCol">The end col.</param>
131+
private static void JoinImageTile(DenseTensor<float> destination, DenseTensor<float> tile, int startRow, int startCol, int endRow, int endCol)
132+
{
133+
int height = endRow - startRow;
134+
int width = endCol - startCol;
135+
int channels = tile.Dimensions[1];
136+
Parallel.For(0, channels, (c) =>
137+
{
138+
Parallel.For(0, height, (i) =>
139+
{
140+
Parallel.For(0, width, (j) =>
141+
{
142+
var value = tile[0, c, i, j];
143+
var existing = destination[0, c, startRow + i, startCol + j];
144+
if (existing > 0)
145+
{
146+
// Blend ovelap
147+
value = (existing + value) / 2f;
148+
}
149+
destination[0, c, startRow + i, startCol + j] = value;
150+
});
151+
});
152+
});
153+
}
43154
}
44155

45156
public enum ImageNormalizeType

‎OnnxStack.Core/Image/OnnxImage.cs

Lines changed: 26 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -64,27 +64,35 @@ public OnnxImage(DenseTensor<float> imageTensor, ImageNormalizeType normalizeTyp
6464
{
6565
var height = imageTensor.Dimensions[2];
6666
var width = imageTensor.Dimensions[3];
67-
var hasTransparency = imageTensor.Dimensions[1] == 4;
68-
_imageData = new Image<Rgba32>(width, height);
69-
for (var y = 0; y < height; y++)
67+
var channels = imageTensor.Dimensions[1];
68+
if (channels == 1)
7069
{
71-
for (var x = 0; x < width; x++)
70+
_imageData = imageTensor.FromMaskTensor();
71+
}
72+
else
73+
{
74+
var hasTransparency = channels == 4;
75+
_imageData = new Image<Rgba32>(width, height);
76+
for (var y = 0; y < height; y++)
7277
{
73-
if (normalizeType == ImageNormalizeType.ZeroToOne)
74-
{
75-
_imageData[x, y] = new Rgba32(
76-
DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
77-
DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
78-
DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
79-
hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
80-
}
81-
else
78+
for (var x = 0; x < width; x++)
8279
{
83-
_imageData[x, y] = new Rgba32(
84-
DenormalizeOneToOneToByte(imageTensor, 0, y, x),
85-
DenormalizeOneToOneToByte(imageTensor, 1, y, x),
86-
DenormalizeOneToOneToByte(imageTensor, 2, y, x),
87-
hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
80+
if (normalizeType == ImageNormalizeType.ZeroToOne)
81+
{
82+
_imageData[x, y] = new Rgba32(
83+
DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
84+
DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
85+
DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
86+
hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
87+
}
88+
else
89+
{
90+
_imageData[x, y] = new Rgba32(
91+
DenormalizeOneToOneToByte(imageTensor, 0, y, x),
92+
DenormalizeOneToOneToByte(imageTensor, 1, y, x),
93+
DenormalizeOneToOneToByte(imageTensor, 2, y, x),
94+
hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
95+
}
8896
}
8997
}
9098
}

‎OnnxStack.Core/Model/ImageTiles.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@
22

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

‎OnnxStack.FeatureExtractor/Common/FeatureExtractorModel.cs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,18 @@ public FeatureExtractorModel(FeatureExtractorModelConfig configuration)
1717

1818
public int OutputChannels => _configuration.OutputChannels;
1919
public int SampleSize => _configuration.SampleSize;
20-
public bool NormalizeOutputTensor => _configuration.NormalizeOutputTensor;
20+
public bool NormalizeOutput => _configuration.NormalizeOutput;
2121
public bool SetOutputToInputAlpha => _configuration.SetOutputToInputAlpha;
2222
public ImageResizeMode InputResizeMode => _configuration.InputResizeMode;
23-
public ImageNormalizeType InputNormalization => _configuration.NormalizeInputTensor;
23+
public ImageNormalizeType NormalizeType => _configuration.NormalizeType;
24+
public bool NormalizeInput => _configuration.NormalizeInput;
2425

2526
public static FeatureExtractorModel Create(FeatureExtractorModelConfig configuration)
2627
{
2728
return new FeatureExtractorModel(configuration);
2829
}
2930

30-
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)
31+
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)
3132
{
3233
var configuration = new FeatureExtractorModelConfig
3334
{
@@ -38,12 +39,12 @@ public static FeatureExtractorModel Create(string modelFile, int sampleSize = 0,
3839
IntraOpNumThreads = 0,
3940
OnnxModelPath = modelFile,
4041

41-
4242
SampleSize = sampleSize,
4343
OutputChannels = outputChannels,
44-
NormalizeOutputTensor = normalizeOutputTensor,
44+
NormalizeType = normalizeType,
45+
NormalizeInput = normalizeInput,
46+
NormalizeOutput = normalizeOutput,
4547
SetOutputToInputAlpha = setOutputToInputAlpha,
46-
NormalizeInputTensor = normalizeInputTensor,
4748
InputResizeMode = inputResizeMode
4849
};
4950
return new FeatureExtractorModel(configuration);

‎OnnxStack.FeatureExtractor/Common/FeatureExtractorModelConfig.cs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,10 @@ public record FeatureExtractorModelConfig : OnnxModelConfig
77
{
88
public int SampleSize { get; set; }
99
public int OutputChannels { get; set; }
10-
public bool NormalizeOutputTensor { get; set; }
10+
public bool NormalizeOutput { get; set; }
1111
public bool SetOutputToInputAlpha { get; set; }
1212
public ImageResizeMode InputResizeMode { get; set; }
13-
public ImageNormalizeType NormalizeInputTensor { get; set; }
13+
public ImageNormalizeType NormalizeType { get; set; }
14+
public bool NormalizeInput { get; set; }
1415
}
1516
}

0 commit comments

Comments
 (0)