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

Commit 3cab2f6

Browse files
saddam213saddam213
committed
Upscale & Feature Extractor stream support
1 parent a494397 commit 3cab2f6

File tree

7 files changed

+265
-63
lines changed

7 files changed

+265
-63
lines changed

‎OnnxStack.Console/Examples/FeatureExtractorVideoExample.cs‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
using OnnxStack.Core.Video;
2+
using OnnxStack.FeatureExtractor.Pipelines;
3+
using OnnxStack.StableDiffusion.Config;
4+
5+
namespace OnnxStack.Console.Runner
6+
{
7+
public sealed class FeatureExtractorVideoExample : IExampleRunner
8+
{
9+
private readonly string _outputDirectory;
10+
private readonly StableDiffusionConfig _configuration;
11+
12+
public FeatureExtractorVideoExample(StableDiffusionConfig configuration)
13+
{
14+
_configuration = configuration;
15+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(FeatureExtractorVideoExample));
16+
Directory.CreateDirectory(_outputDirectory);
17+
}
18+
19+
public int Index => 13;
20+
21+
public string Name => "Feature Extractor Video Example";
22+
23+
public string Description => "Video exmaple using basic feature extractor";
24+
25+
/// <summary>
26+
/// ControlNet Example
27+
/// </summary>
28+
public async Task RunAsync()
29+
{
30+
// Read Video
31+
var videoFile = "C:\\Users\\Deven\\Pictures\\parrot.mp4";
32+
var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
33+
34+
// Create pipeline
35+
var pipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\canny.onnx");
36+
37+
// Create Video Stream
38+
var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
39+
40+
// Create Pipeline Stream
41+
var pipelineStream = pipeline.RunAsync(videoStream);
42+
43+
// Write Video Stream
44+
await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"));
45+
46+
//Unload
47+
await pipeline.UnloadAsync();
48+
}
49+
}
50+
}

‎OnnxStack.Console/Examples/UpscaleExample.cs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,10 @@ public async Task RunAsync()
2929

3030
// Run pipeline
3131
var result = await pipeline.RunAsync(inputImage);
32-
33-
// Create Image from Tensor result
34-
var image = new OnnxImage(result, ImageNormalizeType.ZeroToOne);
35-
32+
3633
// Save Image File
3734
var outputFilename = Path.Combine(_outputDirectory, $"Upscaled.png");
38-
await image.SaveAsync(outputFilename);
35+
await result.SaveAsync(outputFilename);
3936

4037
// Unload
4138
await pipeline.UnloadAsync();

‎OnnxStack.Console/Examples/UpscaleStreamExample.cs‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
using OnnxStack.Core.Video;
2+
using OnnxStack.FeatureExtractor.Pipelines;
3+
4+
namespace OnnxStack.Console.Runner
5+
{
6+
public sealed class UpscaleStreamExample : IExampleRunner
7+
{
8+
private readonly string _outputDirectory;
9+
10+
public UpscaleStreamExample()
11+
{
12+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(UpscaleStreamExample));
13+
Directory.CreateDirectory(_outputDirectory);
14+
}
15+
16+
public int Index => 10;
17+
18+
public string Name => "Streaming Video Upscale Demo";
19+
20+
public string Description => "Upscales a video stream";
21+
22+
public async Task RunAsync()
23+
{
24+
// Read Video
25+
var videoFile = "C:\\Users\\Deven\\Pictures\\parrot.mp4";
26+
var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
27+
28+
// Create pipeline
29+
var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
30+
31+
// Load pipeline
32+
await pipeline.LoadAsync();
33+
34+
// Create Video Stream
35+
var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
36+
37+
// Create Pipeline Stream
38+
var pipelineStream = pipeline.RunAsync(videoStream);
39+
40+
// Write Video Stream
41+
await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"));
42+
43+
//Unload
44+
await pipeline.UnloadAsync();
45+
}
46+
47+
}
48+
}

‎OnnxStack.Core/Image/OnnxImage.cs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,20 @@ public byte[] GetImageBytes()
144144
}
145145

146146

147+
/// <summary>
148+
/// Gets the image as bytes.
149+
/// </summary>
150+
/// <returns></returns>
151+
public async Task<byte[]> GetImageBytesAsync()
152+
{
153+
using (var memoryStream = new MemoryStream())
154+
{
155+
await _imageData.SaveAsPngAsync(memoryStream);
156+
return memoryStream.ToArray();
157+
}
158+
}
159+
160+
147161
/// <summary>
148162
/// Gets the image as stream.
149163
/// </summary>
@@ -156,6 +170,18 @@ public Stream GetImageStream()
156170
}
157171

158172

173+
/// <summary>
174+
/// Gets the image as stream.
175+
/// </summary>
176+
/// <returns></returns>
177+
public async Task<Stream> GetImageStreamAsync()
178+
{
179+
var memoryStream = new MemoryStream();
180+
await _imageData.SaveAsPngAsync(memoryStream);
181+
return memoryStream;
182+
}
183+
184+
159185
/// <summary>
160186
/// Gets the image as tensor.
161187
/// </summary>

‎OnnxStack.FeatureExtractor/Pipelines/FeatureExtractorPipeline.cs‎

Lines changed: 51 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
using System.Collections.Generic;
99
using System.IO;
1010
using System.Linq;
11+
using System.Runtime.CompilerServices;
1112
using System.Threading;
1213
using System.Threading.Tasks;
1314

@@ -58,29 +59,9 @@ public async Task UnloadAsync()
5859
public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
5960
{
6061
var timestamp = _logger?.LogBegin("Extracting image feature...");
61-
var controlImage = await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
62-
var metadata = await _featureExtractorModel.GetMetadataAsync();
63-
cancellationToken.ThrowIfCancellationRequested();
64-
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
65-
{
66-
inferenceParameters.AddInputTensor(controlImage);
67-
inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
68-
69-
var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
70-
using (var result = results.First())
71-
{
72-
cancellationToken.ThrowIfCancellationRequested();
73-
74-
var resultTensor = result.ToDenseTensor();
75-
if (_featureExtractorModel.Normalize)
76-
resultTensor.NormalizeMinMax();
77-
78-
var maskImage = resultTensor.ToImageMask();
79-
//await maskImage.SaveAsPngAsync("D:\\Mask.png");
80-
_logger?.LogEnd("Extracting image feature complete.", timestamp);
81-
return maskImage;
82-
}
83-
}
62+
var result = await RunInternalAsync(inputImage, cancellationToken);
63+
_logger?.LogEnd("Extracting image feature complete.", timestamp);
64+
return result;
8465
}
8566

8667

@@ -92,34 +73,61 @@ public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken ca
9273
public async Task<OnnxVideo> RunAsync(OnnxVideo video, CancellationToken cancellationToken = default)
9374
{
9475
var timestamp = _logger?.LogBegin("Extracting video features...");
76+
var featureFrames = new List<OnnxImage>();
77+
foreach (var videoFrame in video.Frames)
78+
{
79+
featureFrames.Add(await RunAsync(videoFrame, cancellationToken));
80+
}
81+
_logger?.LogEnd("Extracting video features complete.", timestamp);
82+
return new OnnxVideo(video.Info, featureFrames);
83+
}
84+
85+
86+
/// <summary>
87+
/// Generates the feature extractor video stream
88+
/// </summary>
89+
/// <param name="imageFrames">The image frames.</param>
90+
/// <param name="cancellationToken">The cancellation token.</param>
91+
/// <returns></returns>
92+
public async IAsyncEnumerable<OnnxImage> RunAsync(IAsyncEnumerable<OnnxImage> imageFrames, [EnumeratorCancellation] CancellationToken cancellationToken = default)
93+
{
94+
var timestamp = _logger?.LogBegin("Extracting video stream features...");
95+
await foreach (var imageFrame in imageFrames)
96+
{
97+
yield return await RunInternalAsync(imageFrame, cancellationToken);
98+
}
99+
_logger?.LogEnd("Extracting video stream features complete.", timestamp);
100+
}
101+
102+
103+
/// <summary>
104+
/// Runs the pipeline
105+
/// </summary>
106+
/// <param name="inputImage">The input image.</param>
107+
/// <param name="cancellationToken">The cancellation token.</param>
108+
/// <returns></returns>
109+
private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
110+
{
111+
var controlImage = await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
95112
var metadata = await _featureExtractorModel.GetMetadataAsync();
96113
cancellationToken.ThrowIfCancellationRequested();
97-
98-
var frames = new List<OnnxImage>();
99-
foreach (var videoFrame in video.Frames)
114+
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
100115
{
101-
var controlImage = await videoFrame.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
102-
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
103-
{
104-
inferenceParameters.AddInputTensor(controlImage);
105-
inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
116+
inferenceParameters.AddInputTensor(controlImage);
117+
inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
106118

107-
var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
108-
using (var result = results.First())
109-
{
110-
cancellationToken.ThrowIfCancellationRequested();
119+
var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
120+
using (var result = results.First())
121+
{
122+
cancellationToken.ThrowIfCancellationRequested();
111123

112-
var resultTensor = result.ToDenseTensor();
113-
if (_featureExtractorModel.Normalize)
114-
resultTensor.NormalizeMinMax();
124+
var resultTensor = result.ToDenseTensor();
125+
if (_featureExtractorModel.Normalize)
126+
resultTensor.NormalizeMinMax();
115127

116-
var maskImage = resultTensor.ToImageMask();
117-
frames.Add(maskImage);
118-
}
128+
return resultTensor.ToImageMask();
119129
}
120130
}
121-
_logger?.LogEnd("Extracting video features complete.", timestamp);
122-
return new OnnxVideo(video.Info, frames);
123131
}
124132

125133

‎OnnxStack.ImageUpscaler/Pipelines/ImageUpscalePipeline.cs‎

Lines changed: 65 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,16 @@
44
using OnnxStack.Core.Config;
55
using OnnxStack.Core.Image;
66
using OnnxStack.Core.Model;
7+
using OnnxStack.Core.Video;
78
using OnnxStack.ImageUpscaler.Common;
89
using OnnxStack.ImageUpscaler.Extensions;
910
using OnnxStack.ImageUpscaler.Models;
1011
using SixLabors.ImageSharp;
1112
using SixLabors.ImageSharp.PixelFormats;
13+
using System.Collections.Generic;
1214
using System.IO;
1315
using System.Linq;
16+
using System.Runtime.CompilerServices;
1417
using System.Threading;
1518
using System.Threading.Tasks;
1619

@@ -69,7 +72,60 @@ public async Task UnloadAsync()
6972
/// <param name="inputImage">The input image.</param>
7073
/// <param name="cancellationToken">The cancellation token.</param>
7174
/// <returns></returns>
72-
public async Task<DenseTensor<float>> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
75+
public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
76+
{
77+
var timestamp = _logger?.LogBegin("Upscale image..");
78+
var result = await RunInternalAsync(inputImage, cancellationToken);
79+
_logger?.LogEnd("Upscale image complete.", timestamp);
80+
return result;
81+
}
82+
83+
84+
/// <summary>
85+
/// Runs the pipline on a buffered video.
86+
/// </summary>
87+
/// <param name="inputVideo">The input video.</param>
88+
/// <param name="cancellationToken">The cancellation token.</param>
89+
/// <returns></returns>
90+
public async Task<OnnxVideo> RunAsync(OnnxVideo inputVideo, CancellationToken cancellationToken = default)
91+
{
92+
var timestamp = _logger?.LogBegin("Upscale video..");
93+
var upscaledFrames = new List<OnnxImage>();
94+
foreach (var videoFrame in inputVideo.Frames)
95+
{
96+
upscaledFrames.Add(await RunInternalAsync(videoFrame, cancellationToken));
97+
}
98+
99+
var firstFrame = upscaledFrames.First();
100+
var videoInfo = inputVideo.Info with
101+
{
102+
Width = firstFrame.Width,
103+
Height = firstFrame.Height,
104+
};
105+
106+
_logger?.LogEnd("Upscale video complete.", timestamp);
107+
return new OnnxVideo(videoInfo, upscaledFrames);
108+
}
109+
110+
111+
/// <summary>
112+
/// Runs the pipline on a video stream.
113+
/// </summary>
114+
/// <param name="imageFrames">The image frames.</param>
115+
/// <param name="cancellationToken">The cancellation token.</param>
116+
/// <returns></returns>
117+
public async IAsyncEnumerable<OnnxImage> RunAsync(IAsyncEnumerable<OnnxImage> imageFrames, [EnumeratorCancellation] CancellationToken cancellationToken = default)
118+
{
119+
var timestamp = _logger?.LogBegin("Upscale video stream..");
120+
await foreach (var imageFrame in imageFrames)
121+
{
122+
yield return await RunInternalAsync(imageFrame, cancellationToken);
123+
}
124+
_logger?.LogEnd("Upscale video stream complete.", timestamp);
125+
}
126+
127+
128+
private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
73129
{
74130
var upscaleInput = CreateInputParams(inputImage, _upscaleModel.SampleSize, _upscaleModel.ScaleFactor);
75131
var metadata = await _upscaleModel.GetMetadataAsync();
@@ -93,10 +149,16 @@ public async Task<DenseTensor<float>> RunAsync(OnnxImage inputImage, Cancellatio
93149
}
94150
}
95151
}
96-
return outputTensor;
152+
return new OnnxImage(outputTensor, ImageNormalizeType.ZeroToOne);
97153
}
98154

99-
155+
/// <summary>
156+
/// Creates the input parameters.
157+
/// </summary>
158+
/// <param name="imageSource">The image source.</param>
159+
/// <param name="maxTileSize">Maximum size of the tile.</param>
160+
/// <param name="scaleFactor">The scale factor.</param>
161+
/// <returns></returns>
100162
private static UpscaleInput CreateInputParams(OnnxImage imageSource, int maxTileSize, int scaleFactor)
101163
{
102164
var tiles = imageSource.GenerateTiles(maxTileSize, scaleFactor);

0 commit comments

Comments
 (0)