Upscale & Feature Extractor stream support

saddam213 · saddam213 · commit 3cab2f656dae · 2024-02-18T14:31:55.000+13:00
diff --git a/OnnxStack.Console/Examples/FeatureExtractorVideoExample.cs b/OnnxStack.Console/Examples/FeatureExtractorVideoExample.cs
@@ -0,0 +1,50 @@
+﻿using OnnxStack.Core.Video;
+using OnnxStack.FeatureExtractor.Pipelines;
+using OnnxStack.StableDiffusion.Config;
+
+namespace OnnxStack.Console.Runner
+{
+    public sealed class FeatureExtractorVideoExample : IExampleRunner
+    {
+        private readonly string _outputDirectory;
+        private readonly StableDiffusionConfig _configuration;
+
+        public FeatureExtractorVideoExample(StableDiffusionConfig configuration)
+        {
+            _configuration = configuration;
+            _outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(FeatureExtractorVideoExample));
+            Directory.CreateDirectory(_outputDirectory);
+        }
+
+        public int Index => 13;
+
+        public string Name => "Feature Extractor Video Example";
+
+        public string Description => "Video exmaple using basic feature extractor";
+
+        /// <summary>
+        /// ControlNet Example
+        /// </summary>
+        public async Task RunAsync()
+        {
+            // Read Video
+            var videoFile = "C:\\Users\\Deven\\Pictures\\parrot.mp4";
+            var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
+
+            // Create pipeline
+            var pipeline = FeatureExtractorPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators\\canny.onnx");
+
+            // Create Video Stream
+            var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
+
+            // Create Pipeline Stream
+            var pipelineStream = pipeline.RunAsync(videoStream);
+
+            // Write Video Stream
+            await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"));
+
+            //Unload
+            await pipeline.UnloadAsync();
+        }
+    }
+}
diff --git a/OnnxStack.Console/Examples/UpscaleExample.cs b/OnnxStack.Console/Examples/UpscaleExample.cs
@@ -29,13 +29,10 @@ public async Task RunAsync()
 
             // Run pipeline
             var result = await pipeline.RunAsync(inputImage);
-
-            // Create Image from Tensor result
-            var image = new OnnxImage(result, ImageNormalizeType.ZeroToOne);
-
+         
             // Save Image File
             var outputFilename = Path.Combine(_outputDirectory, $"Upscaled.png");
-            await image.SaveAsync(outputFilename);
+            await result.SaveAsync(outputFilename);
 
             // Unload
             await pipeline.UnloadAsync();
diff --git a/OnnxStack.Console/Examples/UpscaleStreamExample.cs b/OnnxStack.Console/Examples/UpscaleStreamExample.cs
@@ -0,0 +1,48 @@
+﻿using OnnxStack.Core.Video;
+using OnnxStack.FeatureExtractor.Pipelines;
+
+namespace OnnxStack.Console.Runner
+{
+    public sealed class UpscaleStreamExample : IExampleRunner
+    {
+        private readonly string _outputDirectory;
+
+        public UpscaleStreamExample()
+        {
+            _outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(UpscaleStreamExample));
+            Directory.CreateDirectory(_outputDirectory);
+        }
+
+        public int Index => 10;
+
+        public string Name => "Streaming Video Upscale Demo";
+
+        public string Description => "Upscales a video stream";
+
+        public async Task RunAsync()
+        {
+            // Read Video
+            var videoFile = "C:\\Users\\Deven\\Pictures\\parrot.mp4";
+            var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
+
+            // Create pipeline
+            var pipeline = ImageUpscalePipeline.CreatePipeline("D:\\Repositories\\upscaler\\SwinIR\\003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx", 4);
+
+            // Load pipeline
+            await pipeline.LoadAsync();
+
+            // Create Video Stream
+            var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
+
+            // Create Pipeline Stream
+            var pipelineStream = pipeline.RunAsync(videoStream);
+
+            // Write Video Stream
+            await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"));
+
+            //Unload
+            await pipeline.UnloadAsync();
+        }
+
+    }
+}
diff --git a/OnnxStack.Core/Image/OnnxImage.cs b/OnnxStack.Core/Image/OnnxImage.cs
@@ -144,6 +144,20 @@ public byte[] GetImageBytes()
         }
 
 
+        /// <summary>
+        /// Gets the image as bytes.
+        /// </summary>
+        /// <returns></returns>
+        public async Task<byte[]> GetImageBytesAsync()
+        {
+            using (var memoryStream = new MemoryStream())
+            {
+                await _imageData.SaveAsPngAsync(memoryStream);
+                return memoryStream.ToArray();
+            }
+        }
+
+
         /// <summary>
         /// Gets the image as stream.
         /// </summary>
@@ -156,6 +170,18 @@ public Stream GetImageStream()
         }
 
 
+        /// <summary>
+        /// Gets the image as stream.
+        /// </summary>
+        /// <returns></returns>
+        public async Task<Stream> GetImageStreamAsync()
+        {
+            var memoryStream = new MemoryStream();
+            await _imageData.SaveAsPngAsync(memoryStream);
+            return memoryStream;
+        }
+
+
         /// <summary>
         /// Gets the image as tensor.
         /// </summary>
diff --git a/OnnxStack.FeatureExtractor/Pipelines/FeatureExtractorPipeline.cs b/OnnxStack.FeatureExtractor/Pipelines/FeatureExtractorPipeline.cs
@@ -8,6 +8,7 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -58,29 +59,9 @@ public async Task UnloadAsync()
         public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
         {
             var timestamp = _logger?.LogBegin("Extracting image feature...");
-            var controlImage = await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
-            var metadata = await _featureExtractorModel.GetMetadataAsync();
-            cancellationToken.ThrowIfCancellationRequested();
-            using (var inferenceParameters = new OnnxInferenceParameters(metadata))
-            {
-                inferenceParameters.AddInputTensor(controlImage);
-                inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
-
-                var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
-                using (var result = results.First())
-                {
-                    cancellationToken.ThrowIfCancellationRequested();
-
-                    var resultTensor = result.ToDenseTensor();
-                    if (_featureExtractorModel.Normalize)
-                        resultTensor.NormalizeMinMax();
-
-                    var maskImage = resultTensor.ToImageMask();
-                    //await maskImage.SaveAsPngAsync("D:\\Mask.png");
-                    _logger?.LogEnd("Extracting image feature complete.", timestamp);
-                    return maskImage;
-                }
-            }
+            var result = await RunInternalAsync(inputImage, cancellationToken);
+            _logger?.LogEnd("Extracting image feature complete.", timestamp);
+            return result;
         }
 
 
@@ -92,34 +73,61 @@ public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken ca
         public async Task<OnnxVideo> RunAsync(OnnxVideo video, CancellationToken cancellationToken = default)
         {
             var timestamp = _logger?.LogBegin("Extracting video features...");
+            var featureFrames = new List<OnnxImage>();
+            foreach (var videoFrame in video.Frames)
+            {
+                featureFrames.Add(await RunAsync(videoFrame, cancellationToken));
+            }
+            _logger?.LogEnd("Extracting video features complete.", timestamp);
+            return new OnnxVideo(video.Info, featureFrames);
+        }
+
+
+        /// <summary>
+        /// Generates the feature extractor video stream
+        /// </summary>
+        /// <param name="imageFrames">The image frames.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public async IAsyncEnumerable<OnnxImage> RunAsync(IAsyncEnumerable<OnnxImage> imageFrames, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            var timestamp = _logger?.LogBegin("Extracting video stream features...");
+            await foreach (var imageFrame in imageFrames)
+            {
+                yield return await RunInternalAsync(imageFrame, cancellationToken);
+            }
+            _logger?.LogEnd("Extracting video stream features complete.", timestamp);
+        }
+
+
+        /// <summary>
+        /// Runs the pipeline
+        /// </summary>
+        /// <param name="inputImage">The input image.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
+        {
+            var controlImage = await inputImage.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
             var metadata = await _featureExtractorModel.GetMetadataAsync();
             cancellationToken.ThrowIfCancellationRequested();
-
-            var frames = new List<OnnxImage>();
-            foreach (var videoFrame in video.Frames)
+            using (var inferenceParameters = new OnnxInferenceParameters(metadata))
             {
-                var controlImage = await videoFrame.GetImageTensorAsync(_featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize, ImageNormalizeType.ZeroToOne);
-                using (var inferenceParameters = new OnnxInferenceParameters(metadata))
-                {
-                    inferenceParameters.AddInputTensor(controlImage);
-                    inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
+                inferenceParameters.AddInputTensor(controlImage);
+                inferenceParameters.AddOutputBuffer(new[] { 1, _featureExtractorModel.Channels, _featureExtractorModel.SampleSize, _featureExtractorModel.SampleSize });
 
-                    var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
-                    using (var result = results.First())
-                    {
-                        cancellationToken.ThrowIfCancellationRequested();
+                var results = await _featureExtractorModel.RunInferenceAsync(inferenceParameters);
+                using (var result = results.First())
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
 
-                        var resultTensor = result.ToDenseTensor();
-                        if (_featureExtractorModel.Normalize)
-                            resultTensor.NormalizeMinMax();
+                    var resultTensor = result.ToDenseTensor();
+                    if (_featureExtractorModel.Normalize)
+                        resultTensor.NormalizeMinMax();
 
-                        var maskImage = resultTensor.ToImageMask();
-                        frames.Add(maskImage);
-                    }
+                    return resultTensor.ToImageMask();
                 }
             }
-            _logger?.LogEnd("Extracting video features complete.", timestamp);
-            return new OnnxVideo(video.Info, frames);
         }
 
 
diff --git a/OnnxStack.ImageUpscaler/Pipelines/ImageUpscalePipeline.cs b/OnnxStack.ImageUpscaler/Pipelines/ImageUpscalePipeline.cs
@@ -4,13 +4,16 @@
 using OnnxStack.Core.Config;
 using OnnxStack.Core.Image;
 using OnnxStack.Core.Model;
+using OnnxStack.Core.Video;
 using OnnxStack.ImageUpscaler.Common;
 using OnnxStack.ImageUpscaler.Extensions;
 using OnnxStack.ImageUpscaler.Models;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -69,7 +72,60 @@ public async Task UnloadAsync()
         /// <param name="inputImage">The input image.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public async Task<DenseTensor<float>> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
+        public async Task<OnnxImage> RunAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
+        {
+            var timestamp = _logger?.LogBegin("Upscale image..");
+            var result = await RunInternalAsync(inputImage, cancellationToken);
+            _logger?.LogEnd("Upscale image complete.", timestamp);
+            return result;
+        }
+
+
+        /// <summary>
+        /// Runs the pipline on a buffered video.
+        /// </summary>
+        /// <param name="inputVideo">The input video.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public async Task<OnnxVideo> RunAsync(OnnxVideo inputVideo, CancellationToken cancellationToken = default)
+        {
+            var timestamp = _logger?.LogBegin("Upscale video..");
+            var upscaledFrames = new List<OnnxImage>();
+            foreach (var videoFrame in inputVideo.Frames)
+            {
+                upscaledFrames.Add(await RunInternalAsync(videoFrame, cancellationToken));
+            }
+
+            var firstFrame = upscaledFrames.First();
+            var videoInfo = inputVideo.Info with
+            {
+                Width = firstFrame.Width,
+                Height = firstFrame.Height,
+            };
+
+            _logger?.LogEnd("Upscale video complete.", timestamp);
+            return new OnnxVideo(videoInfo, upscaledFrames);
+        }
+
+
+        /// <summary>
+        /// Runs the pipline on a video stream.
+        /// </summary>
+        /// <param name="imageFrames">The image frames.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public async IAsyncEnumerable<OnnxImage> RunAsync(IAsyncEnumerable<OnnxImage> imageFrames, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            var timestamp = _logger?.LogBegin("Upscale video stream..");
+            await foreach (var imageFrame in imageFrames)
+            {
+                yield return await RunInternalAsync(imageFrame, cancellationToken);
+            }
+            _logger?.LogEnd("Upscale video stream complete.", timestamp);
+        }
+
+
+        private async Task<OnnxImage> RunInternalAsync(OnnxImage inputImage, CancellationToken cancellationToken = default)
         {
             var upscaleInput = CreateInputParams(inputImage, _upscaleModel.SampleSize, _upscaleModel.ScaleFactor);
             var metadata = await _upscaleModel.GetMetadataAsync();
@@ -93,10 +149,16 @@ public async Task<DenseTensor<float>> RunAsync(OnnxImage inputImage, Cancellatio
                     }
                 }
             }
-            return outputTensor;
+            return new OnnxImage(outputTensor, ImageNormalizeType.ZeroToOne);
         }
 
-
+        /// <summary>
+        /// Creates the input parameters.
+        /// </summary>
+        /// <param name="imageSource">The image source.</param>
+        /// <param name="maxTileSize">Maximum size of the tile.</param>
+        /// <param name="scaleFactor">The scale factor.</param>
+        /// <returns></returns>
         private static UpscaleInput CreateInputParams(OnnxImage imageSource, int maxTileSize, int scaleFactor)
         {
             var tiles = imageSource.GenerateTiles(maxTileSize, scaleFactor);
diff --git a/OnnxStack.UI/Services/UpscaleService.cs b/OnnxStack.UI/Services/UpscaleService.cs
