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

Commit c1025c2

Browse files
saddam213saddam213
committed
Upscale tensor support, tiled tensor support
1 parent b69fbf3 commit c1025c2

File tree

11 files changed

+211
-242
lines changed

11 files changed

+211
-242
lines changed

‎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: 1 addition & 73 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
{
@@ -412,78 +413,5 @@ private static DenseTensor<float> ConcatenateAxis2(DenseTensor<float> tensor1, D
412413

413414
return concatenatedTensor;
414415
}
415-
416-
417-
/// <summary>
418-
/// Splits the Tensor into 4 equal tiles.
419-
/// </summary>
420-
/// <param name="sourceTensor">The source tensor.</param>
421-
/// <returns>TODO: Optimize</returns>
422-
public static ImageTiles SplitTiles(this DenseTensor<float> sourceTensor)
423-
{
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;
451-
}
452-
453-
454-
/// <summary>
455-
/// Rejoins the tiles into a single Tensor.
456-
/// </summary>
457-
/// <param name="tiles">The tiles.</param>
458-
/// <returns>TODO: Optimize</returns>
459-
public static DenseTensor<float> RejoinTiles(this ImageTiles tiles)
460-
{
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-
}
487-
}
488416
}
489417
}

‎OnnxStack.Core/Image/Extensions.cs‎

Lines changed: 100 additions & 0 deletions
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
{
@@ -40,6 +42,104 @@ public static ResizeMode ToResizeMode(this ImageResizeMode resizeMode)
4042
};
4143
}
4244

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

45145
public enum ImageNormalizeType

‎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.ImageUpscaler/Common/UpscaleModel.cs‎

Lines changed: 12 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,39 @@
11
using Microsoft.ML.OnnxRuntime;
22
using OnnxStack.Core.Config;
33
using OnnxStack.Core.Model;
4+
using System;
45

56
namespace OnnxStack.ImageUpscaler.Common
67
{
78
public class UpscaleModel : OnnxModelSession
89
{
9-
private readonly int _channels;
10-
private readonly int _sampleSize;
11-
private readonly int _scaleFactor;
10+
private readonly UpscaleModelConfig _configuration;
1211

1312
public UpscaleModel(UpscaleModelConfig configuration) : base(configuration)
1413
{
15-
_channels = configuration.Channels;
16-
_sampleSize = configuration.SampleSize;
17-
_scaleFactor = configuration.ScaleFactor;
14+
_configuration = configuration;
1815
}
1916

20-
public int Channels => _channels;
21-
public int SampleSize => _sampleSize;
22-
public int ScaleFactor => _scaleFactor;
23-
17+
public int Channels => _configuration.Channels;
18+
public int SampleSize => _configuration.SampleSize;
19+
public int ScaleFactor => _configuration.ScaleFactor;
20+
public int TileSize => _configuration.TileSize;
21+
public int TileOverlap => _configuration.TileOverlap;
2422

2523
public static UpscaleModel Create(UpscaleModelConfig configuration)
2624
{
2725
return new UpscaleModel(configuration);
2826
}
2927

30-
public static UpscaleModel Create(string modelFile, int scaleFactor, int sampleSize = 512, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
28+
public static UpscaleModel Create(string modelFile, int scaleFactor, int sampleSize, int tileSize = 0, int tileOverlap = 20, int channels = 3, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
3129
{
3230
var configuration = new UpscaleModelConfig
3331
{
34-
Channels = 3,
32+
Channels = channels,
3533
SampleSize = sampleSize,
3634
ScaleFactor = scaleFactor,
35+
TileOverlap = tileOverlap,
36+
TileSize = Math.Min(sampleSize, tileSize > 0 ? tileSize : sampleSize),
3737
DeviceId = deviceId,
3838
ExecutionProvider = executionProvider,
3939
ExecutionMode = ExecutionMode.ORT_SEQUENTIAL,
@@ -44,12 +44,4 @@ public static UpscaleModel Create(string modelFile, int scaleFactor, int sampleS
4444
return new UpscaleModel(configuration);
4545
}
4646
}
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-
}
5547
}

‎OnnxStack.ImageUpscaler/Common/UpscaleModelConfig.cs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
using OnnxStack.Core.Config;
2+
3+
namespace OnnxStack.ImageUpscaler.Common
4+
{
5+
public record UpscaleModelConfig : OnnxModelConfig
6+
{
7+
public int Channels { get; set; }
8+
public int SampleSize { get; set; }
9+
public int ScaleFactor { get; set; }
10+
11+
public int TileSize { get; set; }
12+
public int TileOverlap { get; set; }
13+
}
14+
}

‎OnnxStack.ImageUpscaler/Extensions/ImageExtensions.cs‎

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

0 commit comments

Comments
 (0)