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

Commit 26fe4a9

Browse files
saddam213saddam213
authored
Merge pull request #129 from saddam213/Background
Background Removal Pipeline
2 parents d08aaa1 + 4f3b907 commit 26fe4a9

File tree

8 files changed

+340
-15
lines changed

8 files changed

+340
-15
lines changed

‎OnnxStack.Console/Examples/BackgroundRemovalImageExample.cs‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
using OnnxStack.Core.Image;
2+
using OnnxStack.FeatureExtractor.Pipelines;
3+
using System.Diagnostics;
4+
5+
namespace OnnxStack.Console.Runner
6+
{
7+
public sealed class BackgroundRemovalImageExample : IExampleRunner
8+
{
9+
private readonly string _outputDirectory;
10+
11+
public BackgroundRemovalImageExample()
12+
{
13+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", "BackgroundRemovalExample");
14+
Directory.CreateDirectory(_outputDirectory);
15+
}
16+
17+
public int Index => 20;
18+
19+
public string Name => "Image Background Removal Example";
20+
21+
public string Description => "Remove a background from an image";
22+
23+
/// <summary>
24+
/// ControlNet Example
25+
/// </summary>
26+
public async Task RunAsync()
27+
{
28+
OutputHelpers.WriteConsole("Please enter an image file path and press ENTER", ConsoleColor.Yellow);
29+
var imageFile = OutputHelpers.ReadConsole(ConsoleColor.Cyan);
30+
31+
var timestamp = Stopwatch.GetTimestamp();
32+
33+
OutputHelpers.WriteConsole($"Load Image", ConsoleColor.Gray);
34+
var inputImage = await OnnxImage.FromFileAsync(imageFile);
35+
36+
OutputHelpers.WriteConsole($"Create Pipeline", ConsoleColor.Gray);
37+
var pipeline = BackgroundRemovalPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024);
38+
39+
OutputHelpers.WriteConsole($"Run Pipeline", ConsoleColor.Gray);
40+
var imageFeature = await pipeline.RunAsync(inputImage);
41+
42+
OutputHelpers.WriteConsole($"Save Image", ConsoleColor.Gray);
43+
await imageFeature.SaveAsync(Path.Combine(_outputDirectory, $"{pipeline.Name}.png"));
44+
45+
OutputHelpers.WriteConsole($"Unload pipeline", ConsoleColor.Gray);
46+
await pipeline.UnloadAsync();
47+
48+
OutputHelpers.WriteConsole($"Elapsed: {Stopwatch.GetElapsedTime(timestamp)}ms", ConsoleColor.Yellow);
49+
}
50+
}
51+
}

‎OnnxStack.Console/Examples/BackgroundRemovalVideoExample.cs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
using OnnxStack.Core.Video;
2+
using OnnxStack.FeatureExtractor.Pipelines;
3+
using System.Diagnostics;
4+
5+
namespace OnnxStack.Console.Runner
6+
{
7+
public sealed class BackgroundRemovalVideoExample : IExampleRunner
8+
{
9+
private readonly string _outputDirectory;
10+
11+
public BackgroundRemovalVideoExample()
12+
{
13+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", "BackgroundRemovalExample");
14+
Directory.CreateDirectory(_outputDirectory);
15+
}
16+
17+
public int Index => 21;
18+
19+
public string Name => "Video Background Removal Example";
20+
21+
public string Description => "Remove a background from an video";
22+
23+
public async Task RunAsync()
24+
{
25+
OutputHelpers.WriteConsole("Please enter an video/gif file path and press ENTER", ConsoleColor.Yellow);
26+
var videoFile = OutputHelpers.ReadConsole(ConsoleColor.Cyan);
27+
28+
var timestamp = Stopwatch.GetTimestamp();
29+
30+
OutputHelpers.WriteConsole($"Read Video", ConsoleColor.Gray);
31+
var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
32+
33+
OutputHelpers.WriteConsole($"Create Pipeline", ConsoleColor.Gray);
34+
var pipeline = BackgroundRemovalPipeline.CreatePipeline("D:\\Repositories\\RMBG-1.4\\onnx\\model.onnx", sampleSize: 1024);
35+
36+
OutputHelpers.WriteConsole($"Load Pipeline", ConsoleColor.Gray);
37+
await pipeline.LoadAsync();
38+
39+
OutputHelpers.WriteConsole($"Create Video Stream", ConsoleColor.Gray);
40+
var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
41+
42+
OutputHelpers.WriteConsole($"Create Pipeline Stream", ConsoleColor.Gray);
43+
var pipelineStream = pipeline.RunAsync(videoStream);
44+
45+
OutputHelpers.WriteConsole($"Write Video Stream", ConsoleColor.Gray);
46+
await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"Result.mp4"), true);
47+
48+
OutputHelpers.WriteConsole($"Unload", ConsoleColor.Gray);
49+
await pipeline.UnloadAsync();
50+
51+
OutputHelpers.WriteConsole($"Elapsed: {Stopwatch.GetElapsedTime(timestamp)}ms", ConsoleColor.Yellow);
52+
}
53+
}
54+
}

‎OnnxStack.Core/Extensions/TensorExtension.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -397,7 +397,7 @@ private static DenseTensor<float> ConcatenateAxis1(DenseTensor<float> tensor1, D
397397

398398
// Copy data from the second tensor
399399
for (int i = 0; i < dimensions[0]; i++)
400-
for (int j = 0; j < tensor1.Dimensions[1]; j++)
400+
for (int j = 0; j < tensor2.Dimensions[1]; j++)
401401
concatenatedTensor[i, j + tensor1.Dimensions[1]] = tensor2[i, j];
402402

403403
return concatenatedTensor;

‎OnnxStack.Core/Image/OnnxImage.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ 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;
6768
_imageData = new Image<Rgba32>(width, height);
6869
for (var y = 0; y < height; y++)
6970
{
@@ -74,14 +75,16 @@ public OnnxImage(DenseTensor<float> imageTensor, ImageNormalizeType normalizeTyp
7475
_imageData[x, y] = new Rgba32(
7576
DenormalizeZeroToOneToByte(imageTensor, 0, y, x),
7677
DenormalizeZeroToOneToByte(imageTensor, 1, y, x),
77-
DenormalizeZeroToOneToByte(imageTensor, 2, y, x));
78+
DenormalizeZeroToOneToByte(imageTensor, 2, y, x),
79+
hasTransparency ? DenormalizeZeroToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
7880
}
7981
else
8082
{
8183
_imageData[x, y] = new Rgba32(
8284
DenormalizeOneToOneToByte(imageTensor, 0, y, x),
8385
DenormalizeOneToOneToByte(imageTensor, 1, y, x),
84-
DenormalizeOneToOneToByte(imageTensor, 2, y, x));
86+
DenormalizeOneToOneToByte(imageTensor, 2, y, x),
87+
hasTransparency ? DenormalizeOneToOneToByte(imageTensor, 3, y, x) : byte.MaxValue);
8588
}
8689
}
8790
}
@@ -337,6 +340,7 @@ private DenseTensor<float> NormalizeToZeroToOne(ReadOnlySpan<int> dimensions)
337340
var width = dimensions[3];
338341
var height = dimensions[2];
339342
var channels = dimensions[1];
343+
var hasTransparency = channels == 4;
340344
var imageArray = new DenseTensor<float>(new[] { 1, channels, height, width });
341345
_imageData.ProcessPixelRows(img =>
342346
{
@@ -348,6 +352,8 @@ private DenseTensor<float> NormalizeToZeroToOne(ReadOnlySpan<int> dimensions)
348352
imageArray[0, 0, y, x] = (pixelSpan[x].R / 255.0f);
349353
imageArray[0, 1, y, x] = (pixelSpan[x].G / 255.0f);
350354
imageArray[0, 2, y, x] = (pixelSpan[x].B / 255.0f);
355+
if (hasTransparency)
356+
imageArray[0, 3, y, x] = (pixelSpan[x].A / 255.0f);
351357
}
352358
}
353359
});
@@ -366,6 +372,7 @@ private DenseTensor<float> NormalizeToOneToOne(ReadOnlySpan<int> dimensions)
366372
var width = dimensions[3];
367373
var height = dimensions[2];
368374
var channels = dimensions[1];
375+
var hasTransparency = channels == 4;
369376
var imageArray = new DenseTensor<float>(new[] { 1, channels, height, width });
370377
_imageData.ProcessPixelRows(img =>
371378
{
@@ -377,6 +384,8 @@ private DenseTensor<float> NormalizeToOneToOne(ReadOnlySpan<int> dimensions)
377384
imageArray[0, 0, y, x] = (pixelSpan[x].R / 255.0f) * 2.0f - 1.0f;
378385
imageArray[0, 1, y, x] = (pixelSpan[x].G / 255.0f) * 2.0f - 1.0f;
379386
imageArray[0, 2, y, x] = (pixelSpan[x].B / 255.0f) * 2.0f - 1.0f;
387+
if (hasTransparency)
388+
imageArray[0, 3, y, x] = (pixelSpan[x].A / 255.0f) * 2.0f - 1.0f;
380389
}
381390
}
382391
});

‎OnnxStack.Core/Video/OnnxVideo.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,9 +137,9 @@ public void Resize(int height, int width)
137137
/// <param name="filename">The filename.</param>
138138
/// <param name="cancellationToken">The cancellation token.</param>
139139
/// <returns></returns>
140-
public Task SaveAsync(string filename, CancellationToken cancellationToken = default)
140+
public Task SaveAsync(string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
141141
{
142-
return VideoHelper.WriteVideoFramesAsync(this, filename, cancellationToken);
142+
return VideoHelper.WriteVideoFramesAsync(this, filename, preserveTransparency, cancellationToken);
143143
}
144144

145145

‎OnnxStack.Core/Video/VideoHelper.cs‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@ public static void SetConfiguration(OnnxStackConfig configuration)
3232
/// <param name="onnxVideo">The onnx video.</param>
3333
/// <param name="filename">The filename.</param>
3434
/// <param name="cancellationToken">The cancellation token.</param>
35-
public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filename, CancellationToken cancellationToken = default)
35+
public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
3636
{
37-
await WriteVideoFramesAsync(onnxVideo.Frames, filename, onnxVideo.FrameRate, onnxVideo.AspectRatio, cancellationToken);
37+
await WriteVideoFramesAsync(onnxVideo.Frames, filename, onnxVideo.FrameRate, onnxVideo.AspectRatio, preserveTransparency, cancellationToken);
3838
}
3939

4040

@@ -45,11 +45,11 @@ public static async Task WriteVideoFramesAsync(OnnxVideo onnxVideo, string filen
4545
/// <param name="filename">The filename.</param>
4646
/// <param name="frameRate">The frame rate.</param>
4747
/// <param name="cancellationToken">The cancellation token.</param>
48-
public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate = 15, CancellationToken cancellationToken = default)
48+
public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate = 15, bool preserveTransparency = false, CancellationToken cancellationToken = default)
4949
{
5050
var firstImage = onnxImages.First();
5151
var aspectRatio = (double)firstImage.Width / firstImage.Height;
52-
await WriteVideoFramesAsync(onnxImages, filename, frameRate, aspectRatio, cancellationToken);
52+
await WriteVideoFramesAsync(onnxImages, filename, frameRate, aspectRatio, preserveTransparency, cancellationToken);
5353
}
5454

5555

@@ -61,12 +61,12 @@ public static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages
6161
/// <param name="frameRate">The frame rate.</param>
6262
/// <param name="aspectRatio">The aspect ratio.</param>
6363
/// <param name="cancellationToken">The cancellation token.</param>
64-
private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate, double aspectRatio, CancellationToken cancellationToken = default)
64+
private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImages, string filename, float frameRate, double aspectRatio, bool preserveTransparency, CancellationToken cancellationToken = default)
6565
{
6666
if (File.Exists(filename))
6767
File.Delete(filename);
6868

69-
using (var videoWriter = CreateWriter(filename, frameRate, aspectRatio))
69+
using (var videoWriter = CreateWriter(filename, frameRate, aspectRatio, preserveTransparency))
7070
{
7171
// Start FFMPEG
7272
videoWriter.Start();
@@ -91,12 +91,12 @@ private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImage
9191
/// <param name="frameRate">The frame rate.</param>
9292
/// <param name="aspectRatio">The aspect ratio.</param>
9393
/// <param name="cancellationToken">The cancellation token.</param>
94-
public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumerable<OnnxImage> videoStream, string filename, CancellationToken cancellationToken = default)
94+
public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumerable<OnnxImage> videoStream, string filename, bool preserveTransparency = false, CancellationToken cancellationToken = default)
9595
{
9696
if (File.Exists(filename))
9797
File.Delete(filename);
9898

99-
using (var videoWriter = CreateWriter(filename, videoInfo.FrameRate, videoInfo.AspectRatio))
99+
using (var videoWriter = CreateWriter(filename, videoInfo.FrameRate, videoInfo.AspectRatio, preserveTransparency))
100100
{
101101
// Start FFMPEG
102102
videoWriter.Start();
@@ -323,11 +323,13 @@ private static Process CreateReader(string inputFile, float fps)
323323
/// <param name="fps">The FPS.</param>
324324
/// <param name="aspectRatio">The aspect ratio.</param>
325325
/// <returns></returns>
326-
private static Process CreateWriter(string outputFile, float fps, double aspectRatio)
326+
private static Process CreateWriter(string outputFile, float fps, double aspectRatio, bool preserveTransparency)
327327
{
328328
var ffmpegProcess = new Process();
329+
var codec = preserveTransparency ? "png" : "libx264";
330+
var format = preserveTransparency ? "yuva420p" : "yuv420p";
329331
ffmpegProcess.StartInfo.FileName = _configuration.FFmpegPath;
330-
ffmpegProcess.StartInfo.Arguments = $"-hide_banner -loglevel error -framerate {fps:F4} -i - -c:v libx264 -movflags +faststart -vf format=yuv420p -aspect {aspectRatio} {outputFile}";
332+
ffmpegProcess.StartInfo.Arguments = $"-hide_banner -loglevel error -framerate {fps:F4} -i - -c:v {codec} -movflags +faststart -vf format={format} -aspect {aspectRatio} {outputFile}";
331333
ffmpegProcess.StartInfo.RedirectStandardInput = true;
332334
ffmpegProcess.StartInfo.UseShellExecute = false;
333335
ffmpegProcess.StartInfo.CreateNoWindow = true;

0 commit comments

Comments
 (0)