Merge pull request #129 from saddam213/Background

Background Removal Pipeline

Author: saddam213

OnnxStack.Console/Examples/BackgroundRemovalImageExample.cs

@@ -0,0 +1,51 @@
+using OnnxStack.Core.Image;
+using OnnxStack.FeatureExtractor.Pipelines;
+using System.Diagnostics;
+
+namespace OnnxStack.Console.Runner
+{
+    public sealed class BackgroundRemovalImageExample : IExampleRunner
+    {
+        private readonly string _outputDirectory;
+
+        public BackgroundRemovalImageExample()
+        {
+            _outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", "BackgroundRemovalExample");
+            Directory.CreateDirectory(_outputDirectory);
+        }
+
+        public int Index => 20;
+
+        public string Name => "Image Background Removal Example";
+
+        public string Description => "Remove a background from an image";
+
+        /// <summary>
+        /// ControlNet Example
+        /// </summary>
+        public async Task RunAsync()
+        {
+            OutputHelpers.WriteConsole("Please enter an image file path and press ENTER", ConsoleColor.Yellow);
+            var imageFile = OutputHelpers.ReadConsole(ConsoleColor.Cyan);
+
+            var timestamp = Stopwatch.GetTimestamp();
+
+            OutputHelpers.WriteConsole($"Load Image", ConsoleColor.Gray);
+            var inputImage = await OnnxImage.FromFileAsync(imageFile);
+
+            OutputHelpers.WriteConsole($"Create Pipeline", ConsoleColor.Gray);
+            var pipeline = BackgroundRemovalPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024);
+
+            OutputHelpers.WriteConsole($"Run Pipeline", ConsoleColor.Gray);
+            var imageFeature = await pipeline.RunAsync(inputImage);
+
+            OutputHelpers.WriteConsole($"Save Image", ConsoleColor.Gray);
+            await imageFeature.SaveAsync(Path.Combine(_outputDirectory, $"{pipeline.Name}.png"));
+
+            OutputHelpers.WriteConsole($"Unload pipeline", ConsoleColor.Gray);
+            await pipeline.UnloadAsync();
+
+            OutputHelpers.WriteConsole($"Elapsed: {Stopwatch.GetElapsedTime(timestamp)}ms", ConsoleColor.Yellow);
+        }
+    }
+}

OnnxStack.Console/Examples/BackgroundRemovalVideoExample.cs

@@ -0,0 +1,54 @@
+using OnnxStack.Core.Video;
+using OnnxStack.FeatureExtractor.Pipelines;
+using System.Diagnostics;
+
+namespace OnnxStack.Console.Runner
+{
+    public sealed class BackgroundRemovalVideoExample : IExampleRunner
+    {
+        private readonly string _outputDirectory;
+
+        public BackgroundRemovalVideoExample()
+        {
+            _outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", "BackgroundRemovalExample");
+            Directory.CreateDirectory(_outputDirectory);
+        }
+
+        public int Index => 21;
+
+        public string Name => "Video Background Removal Example";
+
+        public string Description => "Remove a background from an video";
+
+        public async Task RunAsync()
+        {
+            OutputHelpers.WriteConsole("Please enter an video/gif file path and press ENTER", ConsoleColor.Yellow);
+            var videoFile = OutputHelpers.ReadConsole(ConsoleColor.Cyan);
+
+            var timestamp = Stopwatch.GetTimestamp();
+
+            OutputHelpers.WriteConsole($"Read Video", ConsoleColor.Gray);
+            var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
+
+            OutputHelpers.WriteConsole($"Create Pipeline", ConsoleColor.Gray);
+            var pipeline = BackgroundRemovalPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024);
+
+            OutputHelpers.WriteConsole($"Load Pipeline", ConsoleColor.Gray);
+            await pipeline.LoadAsync();
+
+            OutputHelpers.WriteConsole($"Create Video Stream", ConsoleColor.Gray);
+            var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
+
+            OutputHelpers.WriteConsole($"Create Pipeline Stream", ConsoleColor.Gray);
+            var pipelineStream = pipeline.RunAsync(videoStream);
+
+            OutputHelpers.WriteConsole($"Write Video Stream", ConsoleColor.Gray);
+            await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"), true);
+
+            OutputHelpers.WriteConsole($"Unload", ConsoleColor.Gray);
+            await pipeline.UnloadAsync();
+
+            OutputHelpers.WriteConsole($"Elapsed: {Stopwatch.GetElapsedTime(timestamp)}ms", ConsoleColor.Yellow);
+        }
+    }
+}

OnnxStack.Core/Extensions/TensorExtension.cs

@@ -397,7 +397,7 @@ private static DenseTensor<float> ConcatenateAxis1(DenseTensor<float> tensor1, D
 
             // Copy data from the second tensor
             for (int i = 0; i < dimensions[0]; i++)
-                for (int j = 0; j < tensor1.Dimensions[1]; j++)
+                for (int j = 0; j < tensor2.Dimensions[1]; j++)
                     concatenatedTensor[i, j + tensor1.Dimensions[1]] = tensor2[i, j];
 
             return concatenatedTensor;

OnnxStack.Core/Image/OnnxImage.cs

@@ -64,6 +64,7 @@ public OnnxImage(DenseTensor<float> imageTensor, ImageNormalizeType normalizeTyp
         {
             var height = imageTensor.Dimensions[2];
             var width = imageTensor.Dimensions[3];
+            var hasTransparency = imageTensor.Dimensions[1] == 4;
             _imageData = new Image<Rgba32>(width, height);
             for (var y = 0; y < height; y++)
             {
@@ -74,14 +75,16 @@ public OnnxImage(DenseTensor<float> imageTensor, ImageNormalizeType normalizeTyp
                         _imageData[x, y] = new Rgba32(
                             DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
                             DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
-                            DenormalizeZeroToOneToByte(imageTensor, 2, y, x));
+                            DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
+                            hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
                     }
                     else
                     {
                         _imageData[x, y] = new Rgba32(
                             DenormalizeOneToOneToByte(imageTensor, 0, y, x),
                             DenormalizeOneToOneToByte(imageTensor, 1, y, x),
-                            DenormalizeOneToOneToByte(imageTensor, 2, y, x));
+                            DenormalizeOneToOneToByte(imageTensor, 2, y, x),
+                            hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
                     }
                 }
             }
@@ -337,6 +340,7 @@ private DenseTensor<float> NormalizeToZeroToOne(ReadOnlySpan<int> dimensions)
             var width = dimensions[3];
             var height = dimensions[2];
             var channels = dimensions[1];
+            var hasTransparency = channels == 4;
             var imageArray = new DenseTensor<float>(new[] { 1, channels, height, width });
             _imageData.ProcessPixelRows(img =>
             {
@@ -348,6 +352,8 @@ private DenseTensor<float> NormalizeToZeroToOne(ReadOnlySpan<int> dimensions)
                         imageArray[0, 0, y, x] = (pixelSpan[x].R / 255.0f);
                         imageArray[0, 1, y, x] = (pixelSpan[x].G / 255.0f);
                         imageArray[0, 2, y, x] = (pixelSpan[x].B / 255.0f);
+                        if (hasTransparency)
+                            imageArray[0, 3, y, x] = (pixelSpan[x].A / 255.0f);
                     }
                 }
             });
@@ -366,6 +372,7 @@ private DenseTensor<float> NormalizeToOneToOne(ReadOnlySpan<int> dimensions)
             var width = dimensions[3];
             var height = dimensions[2];
             var channels = dimensions[1];
+            var hasTransparency = channels == 4;
             var imageArray = new DenseTensor<float>(new[] { 1, channels, height, width });
             _imageData.ProcessPixelRows(img =>
             {
@@ -377,6 +384,8 @@ private DenseTensor<float> NormalizeToOneToOne(ReadOnlySpan<int> dimensions)
                         imageArray[0, 0, y, x] = (pixelSpan[x].R / 255.0f) * 2.0f - 1.0f;
                         imageArray[0, 1, y, x] = (pixelSpan[x].G / 255.0f) * 2.0f - 1.0f;
                         imageArray[0, 2, y, x] = (pixelSpan[x].B / 255.0f) * 2.0f - 1.0f;
+                        if (hasTransparency)
+                            imageArray[0, 3, y, x] = (pixelSpan[x].A / 255.0f) * 2.0f - 1.0f;
                     }
                 }
             });

OnnxStack.Core/Video/OnnxVideo.cs

@@ -137,9 +137,9 @@ public void Resize(int height, int width)
         /// <param name="filename">The filename.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public Task SaveAsync(string filename, CancellationToken cancellationToken = default)
+        public Task SaveAsync(string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
         {
-            return VideoHelper.WriteVideoFramesAsync(this, filename, cancellationToken);
+            return VideoHelper.WriteVideoFramesAsync(this, filename, preserveTransparency, cancellationToken);
         }
 
 

OnnxStack.Core/Video/VideoHelper.cs

@@ -32,9 +32,9 @@ public static void SetConfiguration(OnnxStackConfig configuration)
         /// <param name="onnxVideo">The onnx video.</param>
         /// <param name="filename">The filename.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
-        public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filename, CancellationToken cancellationToken = default)
+        public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
         {
-            await WriteVideoFramesAsync(onnxVideo.Frames, filename, onnxVideo.FrameRate, onnxVideo.AspectRatio, cancellationToken);
+            await WriteVideoFramesAsync(onnxVideo.Frames, filename, onnxVideo.FrameRate, onnxVideo.AspectRatio, preserveTransparency, cancellationToken);
         }
 
 
@@ -45,11 +45,11 @@ public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filen
         /// <param name="filename">The filename.</param>
         /// <param name="frameRate">The frame rate.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
-        public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate = 15, CancellationToken cancellationToken = default)
+        public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate = 15, bool preserveTransparency = false, CancellationToken cancellationToken = default)
         {
             var firstImage = onnxImages.First();
             var aspectRatio = (double)firstImage.Width / firstImage.Height;
-            await WriteVideoFramesAsync(onnxImages, filename, frameRate, aspectRatio, cancellationToken);
+            await WriteVideoFramesAsync(onnxImages, filename, frameRate, aspectRatio, preserveTransparency, cancellationToken);
         }
 
 
@@ -61,12 +61,12 @@ public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages
         /// <param name="frameRate">The frame rate.</param>
         /// <param name="aspectRatio">The aspect ratio.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
-        private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate, double aspectRatio, CancellationToken cancellationToken = default)
+        private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate, double aspectRatio, bool preserveTransparency, CancellationToken cancellationToken = default)
         {
             if (File.Exists(filename))
                 File.Delete(filename);
 
-            using (var videoWriter = CreateWriter(filename, frameRate, aspectRatio))
+            using (var videoWriter = CreateWriter(filename, frameRate, aspectRatio, preserveTransparency))
             {
                 // Start FFMPEG
                 videoWriter.Start();
@@ -91,12 +91,12 @@ private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImage
         /// <param name="frameRate">The frame rate.</param>
         /// <param name="aspectRatio">The aspect ratio.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
-        public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumerable<OnnxImage> videoStream, string filename, CancellationToken cancellationToken = default)
+        public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumerable<OnnxImage> videoStream, string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
         {
             if (File.Exists(filename))
                 File.Delete(filename);
 
-            using (var videoWriter = CreateWriter(filename, videoInfo.FrameRate, videoInfo.AspectRatio))
+            using (var videoWriter = CreateWriter(filename, videoInfo.FrameRate, videoInfo.AspectRatio, preserveTransparency))
             {
                 // Start FFMPEG
                 videoWriter.Start();
@@ -323,11 +323,13 @@ private static Process CreateReader(string inputFile, float fps)
         /// <param name="fps">The FPS.</param>
         /// <param name="aspectRatio">The aspect ratio.</param>
         /// <returns></returns>
-        private static Process CreateWriter(string outputFile, float fps, double aspectRatio)
+        private static Process CreateWriter(string outputFile, float fps, double aspectRatio, bool preserveTransparency)
         {
             var ffmpegProcess = new Process();
+            var codec = preserveTransparency ? "png" : "libx264";
+            var format = preserveTransparency ? "yuva420p" : "yuv420p";
             ffmpegProcess.StartInfo.FileName = _configuration.FFmpegPath;
-            ffmpegProcess.StartInfo.Arguments = $"-hide_banner -loglevel error -framerate {fps:F4} -i - -c:v libx264 -movflags +faststart -vf format=yuv420p -aspect {aspectRatio} {outputFile}";
+            ffmpegProcess.StartInfo.Arguments = $"-hide_banner -loglevel error -framerate {fps:F4} -i - -c:v {codec} -movflags +faststart -vf format={format} -aspect {aspectRatio} {outputFile}";
             ffmpegProcess.StartInfo.RedirectStandardInput = true;
             ffmpegProcess.StartInfo.UseShellExecute = false;
             ffmpegProcess.StartInfo.CreateNoWindow = true;