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

Commit a494397

Browse files
saddam213saddam213
committed
Streaming video support
1 parent ba6475f commit a494397

File tree

5 files changed

+131
-122
lines changed

5 files changed

+131
-122
lines changed

‎OnnxStack.Core/Extensions/TensorExtension.cs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using Microsoft.ML.OnnxRuntime.Tensors;
22
using System;
33
using System.Collections.Generic;
4+
using System.Linq;
45

56
namespace OnnxStack.Core
67
{
@@ -75,6 +76,33 @@ public static IEnumerable<DenseTensor<float>> SplitBatch(this DenseTensor<float>
7576
}
7677

7778

79+
/// <summary>
80+
/// Joins the tensors across the 0 axis.
81+
/// </summary>
82+
/// <param name="tensors">The tensors.</param>
83+
/// <param name="axis">The axis.</param>
84+
/// <returns></returns>
85+
/// <exception cref="System.NotImplementedException">Only axis 0 is supported</exception>
86+
public static DenseTensor<float> Join(this IList<DenseTensor<float>> tensors, int axis = 0)
87+
{
88+
if (axis != 0)
89+
throw new NotImplementedException("Only axis 0 is supported");
90+
91+
var tensor = tensors.First();
92+
var dimensions = tensor.Dimensions.ToArray();
93+
dimensions[0] *= tensors.Count;
94+
95+
var newLength = (int)tensor.Length;
96+
var buffer = new float[newLength * tensors.Count].AsMemory();
97+
for (int i = 0; i < tensors.Count(); i++)
98+
{
99+
var start = i * newLength;
100+
tensors[i].Buffer.CopyTo(buffer[start..]);
101+
}
102+
return new DenseTensor<float>(buffer, dimensions);
103+
}
104+
105+
78106
/// <summary>
79107
/// Concatenates the specified tensors along the specified axis.
80108
/// </summary>

‎OnnxStack.Core/Video/OnnxVideo.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ public OnnxVideo(VideoInfo info, IEnumerable<DenseTensor<float>> videoTensors)
8888
/// <summary>
8989
/// Gets the aspect ratio.
9090
/// </summary>
91-
public double AspectRatio => (double)_info.Width / _info.Height;
91+
public double AspectRatio => _info.AspectRatio;
9292

9393
/// <summary>
9494
/// Gets a value indicating whether this instance has video.

‎OnnxStack.Core/Video/VideoHelper.cs‎

Lines changed: 100 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,36 @@ private static async Task WriteVideoFramesAsync(IEnumerable<OnnxImage> onnxImage
8383
}
8484

8585

86+
/// <summary>
87+
/// Writes the video stream to file.
88+
/// </summary>
89+
/// <param name="onnxImages">The onnx image stream.</param>
90+
/// <param name="filename">The filename.</param>
91+
/// <param name="frameRate">The frame rate.</param>
92+
/// <param name="aspectRatio">The aspect ratio.</param>
93+
/// <param name="cancellationToken">The cancellation token.</param>
94+
public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumerable<OnnxImage> videoStream, string filename, CancellationToken cancellationToken = default)
95+
{
96+
if (File.Exists(filename))
97+
File.Delete(filename);
98+
99+
using (var videoWriter = CreateWriter(filename, videoInfo.FrameRate, videoInfo.AspectRatio))
100+
{
101+
// Start FFMPEG
102+
videoWriter.Start();
103+
await foreach (var frame in videoStream)
104+
{
105+
// Write each frame to the input stream of FFMPEG
106+
await videoWriter.StandardInput.BaseStream.WriteAsync(frame.GetImageBytes(), cancellationToken);
107+
}
108+
109+
// Done close stream and wait for app to process
110+
videoWriter.StandardInput.BaseStream.Close();
111+
await videoWriter.WaitForExitAsync(cancellationToken);
112+
}
113+
}
114+
115+
86116
/// <summary>
87117
/// Reads the video information.
88118
/// </summary>
@@ -119,9 +149,16 @@ public static async Task<VideoInfo> ReadVideoInfoAsync(string filename)
119149
/// <returns></returns>
120150
public static async Task<List<OnnxImage>> ReadVideoFramesAsync(byte[] videoBytes, float frameRate = 15, CancellationToken cancellationToken = default)
121151
{
122-
return await CreateFramesInternalAsync(videoBytes, frameRate, cancellationToken)
123-
.Select(x => new OnnxImage(x))
124-
.ToListAsync(cancellationToken);
152+
string tempVideoPath = GetTempFilename();
153+
try
154+
{
155+
await File.WriteAllBytesAsync(tempVideoPath, videoBytes, cancellationToken);
156+
return await ReadVideoStreamAsync(tempVideoPath, frameRate, cancellationToken).ToListAsync(cancellationToken);
157+
}
158+
finally
159+
{
160+
DeleteTempFile(tempVideoPath);
161+
}
125162
}
126163

127164

@@ -134,10 +171,23 @@ public static async Task<List<OnnxImage>> ReadVideoFramesAsync(byte[] videoBytes
134171
/// <returns></returns>
135172
public static async Task<List<OnnxImage>> ReadVideoFramesAsync(string filename, float frameRate = 15, CancellationToken cancellationToken = default)
136173
{
137-
var videoBytes = await File.ReadAllBytesAsync(filename, cancellationToken);
138-
return await CreateFramesInternalAsync(videoBytes, frameRate, cancellationToken)
139-
.Select(x => new OnnxImage(x))
140-
.ToListAsync(cancellationToken);
174+
return await ReadVideoStreamAsync(filename, frameRate, cancellationToken).ToListAsync(cancellationToken);
175+
}
176+
177+
178+
/// <summary>
179+
/// Reads the video frames as a stream.
180+
/// </summary>
181+
/// <param name="filename">The filename.</param>
182+
/// <param name="frameRate">The frame rate.</param>
183+
/// <param name="cancellationToken">The cancellation token.</param>
184+
/// <returns></returns>
185+
public static async IAsyncEnumerable<OnnxImage> ReadVideoStreamAsync(string filename, float frameRate = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)
186+
{
187+
await foreach (var frameBytes in CreateFramesInternalAsync(filename, frameRate, cancellationToken))
188+
{
189+
yield return new OnnxImage(frameBytes);
190+
}
141191
}
142192

143193

@@ -152,76 +202,67 @@ public static async Task<List<OnnxImage>> ReadVideoFramesAsync(string filename,
152202
/// <param name="cancellationToken">The cancellation token.</param>
153203
/// <returns></returns>
154204
/// <exception cref="Exception">Invalid PNG header</exception>
155-
private static async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(byte[] videoData, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)
205+
private static async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(string fileName, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)
156206
{
157-
string tempVideoPath = GetTempFilename();
158-
try
207+
using (var ffmpegProcess = CreateReader(fileName, fps))
159208
{
160-
await File.WriteAllBytesAsync(tempVideoPath, videoData, cancellationToken);
161-
using (var ffmpegProcess = CreateReader(tempVideoPath, fps))
162-
{
163-
// Start FFMPEG
164-
ffmpegProcess.Start();
209+
// Start FFMPEG
210+
ffmpegProcess.Start();
165211

166-
// FFMPEG output stream
167-
var processOutputStream = ffmpegProcess.StandardOutput.BaseStream;
212+
// FFMPEG output stream
213+
var processOutputStream = ffmpegProcess.StandardOutput.BaseStream;
168214

169-
// Buffer to hold the current image
170-
var buffer = new byte[20480000];
215+
// Buffer to hold the current image
216+
var buffer = new byte[20480000];
171217

172-
var currentIndex = 0;
173-
while (!cancellationToken.IsCancellationRequested)
174-
{
175-
// Reset the index new PNG
176-
currentIndex = 0;
218+
var currentIndex = 0;
219+
while (!cancellationToken.IsCancellationRequested)
220+
{
221+
// Reset the index new PNG
222+
currentIndex = 0;
177223

178-
// Read the PNG Header
179-
if (await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, 8), cancellationToken) <= 0)
180-
break;
224+
// Read the PNG Header
225+
if (await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, 8), cancellationToken) <= 0)
226+
break;
181227

182-
currentIndex += 8;// header length
228+
currentIndex += 8;// header length
183229

184-
if (!IsImageHeader(buffer))
185-
throw new Exception("Invalid PNG header");
230+
if (!IsImageHeader(buffer))
231+
throw new Exception("Invalid PNG header");
186232

187-
// loop through each chunk
188-
while (true)
189-
{
190-
// Read the chunk header
191-
await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, 12), cancellationToken);
233+
// loop through each chunk
234+
while (true)
235+
{
236+
// Read the chunk header
237+
await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, 12), cancellationToken);
192238

193-
var chunkIndex = currentIndex;
194-
currentIndex += 12; // Chunk header length
239+
var chunkIndex = currentIndex;
240+
currentIndex += 12; // Chunk header length
195241

196-
// Get the chunk's content size in bytes from the header we just read
197-
var totalSize = buffer[chunkIndex] << 24 | buffer[chunkIndex + 1] << 16 | buffer[chunkIndex + 2] << 8 | buffer[chunkIndex + 3];
198-
if (totalSize > 0)
242+
// Get the chunk's content size in bytes from the header we just read
243+
var totalSize = buffer[chunkIndex] << 24 | buffer[chunkIndex + 1] << 16 | buffer[chunkIndex + 2] << 8 | buffer[chunkIndex + 3];
244+
if (totalSize > 0)
245+
{
246+
var totalRead = 0;
247+
while (totalRead < totalSize)
199248
{
200-
var totalRead = 0;
201-
while (totalRead < totalSize)
202-
{
203-
int read = await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, totalSize - totalRead), cancellationToken);
204-
currentIndex += read;
205-
totalRead += read;
206-
}
207-
continue;
249+
int read = await processOutputStream.ReadAsync(buffer.AsMemory(currentIndex, totalSize - totalRead), cancellationToken);
250+
currentIndex += read;
251+
totalRead += read;
208252
}
209-
210-
// If the size is 0 and is the end of the image
211-
if (totalSize == 0 && IsImageEnd(buffer, chunkIndex))
212-
break;
253+
continue;
213254
}
214255

215-
yield return buffer[..currentIndex];
256+
// If the size is 0 and is the end of the image
257+
if (totalSize == 0 && IsImageEnd(buffer, chunkIndex))
258+
break;
216259
}
217260

218-
if (cancellationToken.IsCancellationRequested)
219-
ffmpegProcess.Kill();
261+
yield return buffer[..currentIndex];
220262
}
221-
}
222-
finally
223-
{
224-
DeleteTempFile(tempVideoPath);
263+
264+
if (cancellationToken.IsCancellationRequested)
265+
ffmpegProcess.Kill();
225266
}
226267
}
227268

‎OnnxStack.Core/Video/VideoInfo.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,5 +11,7 @@ public VideoInfo(int height, int width, TimeSpan duration, float frameRate) : th
1111
}
1212
public int Height { get; set; }
1313
public int Width { get; set; }
14+
15+
public double AspectRatio => (double)Height / Width;
1416
}
1517
}

‎OnnxStack.StableDiffusion/Helpers/TensorHelper.cs‎

Lines changed: 0 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -247,15 +247,6 @@ public static DenseTensor<float> Clip(this DenseTensor<float> tensor, float minV
247247
}
248248

249249

250-
251-
252-
253-
254-
255-
256-
257-
258-
259250
/// <summary>
260251
/// Generate a random Tensor from a normal distribution with mean 0 and variance 1
261252
/// </summary>
@@ -279,58 +270,5 @@ public static DenseTensor<float> GetRandomTensor(Random random, ReadOnlySpan<int
279270
return latents;
280271
}
281272

282-
283-
/// <summary>
284-
/// Splits the Tensor along axis 0.
285-
/// </summary>
286-
/// <param name="tensor">The tensor.</param>
287-
/// <param name="count">The count.</param>
288-
/// <param name="axis">The axis.</param>
289-
/// <returns></returns>
290-
/// <exception cref="System.NotImplementedException">Only axis 0 is supported</exception>
291-
public static DenseTensor<float>[] Split(this DenseTensor<float> tensor, int count, int axis = 0)
292-
{
293-
if (axis != 0)
294-
throw new NotImplementedException("Only axis 0 is supported");
295-
296-
var dimensions = tensor.Dimensions.ToArray();
297-
dimensions[0] /= count;
298-
299-
var newLength = (int)tensor.Length / count;
300-
var results = new DenseTensor<float>[count];
301-
for (int i = 0; i < count; i++)
302-
{
303-
var start = i * newLength;
304-
results[i] = new DenseTensor<float>(tensor.Buffer.Slice(start, newLength), dimensions);
305-
}
306-
return results;
307-
}
308-
309-
310-
/// <summary>
311-
/// Joins the tensors across the 0 axis.
312-
/// </summary>
313-
/// <param name="tensors">The tensors.</param>
314-
/// <param name="axis">The axis.</param>
315-
/// <returns></returns>
316-
/// <exception cref="System.NotImplementedException">Only axis 0 is supported</exception>
317-
public static DenseTensor<float> Join(this IList<DenseTensor<float>> tensors, int axis = 0)
318-
{
319-
if (axis != 0)
320-
throw new NotImplementedException("Only axis 0 is supported");
321-
322-
var tensor = tensors.First();
323-
var dimensions = tensor.Dimensions.ToArray();
324-
dimensions[0] *= tensors.Count;
325-
326-
var newLength = (int)tensor.Length;
327-
var buffer = new float[newLength * tensors.Count].AsMemory();
328-
for (int i = 0; i < tensors.Count(); i++)
329-
{
330-
var start = i * newLength;
331-
tensors[i].Buffer.CopyTo(buffer[start..]);
332-
}
333-
return new DenseTensor<float>(buffer, dimensions);
334-
}
335273
}
336274
}

0 commit comments

Comments
 (0)