Add first draft of keyframe extraction for Matroska

10 files changed, 481 insertions, 0 deletions

diff --git a/src/Jellyfin.MediaEncoding.Keyframes/FfProbe/FfProbeKeyframeExtractor.cs b/src/Jellyfin.MediaEncoding.Keyframes/FfProbe/FfProbeKeyframeExtractor.cs
new file mode 100644
index 000000000..249608ef9
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/FfProbe/FfProbeKeyframeExtractor.cs
@@ -0,0 +1,10 @@
+using System;
+
+namespace Jellyfin.MediaEncoding.Keyframes.FfProbe
+{
+    public static class FfProbeKeyframeExtractor
+    {
+        // TODO
+        public static KeyframeData GetKeyframeData(string ffProbePath, string filePath) => throw new NotImplementedException();
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/FfTool/FfToolKeyframeExtractor.cs b/src/Jellyfin.MediaEncoding.Keyframes/FfTool/FfToolKeyframeExtractor.cs
new file mode 100644
index 000000000..89c149ff4
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/FfTool/FfToolKeyframeExtractor.cs
@@ -0,0 +1,10 @@
+using System;
+
+namespace Jellyfin.MediaEncoding.Keyframes.FfTool
+{
+    public static class FfToolKeyframeExtractor
+    {
+        // TODO
+        public static KeyframeData GetKeyframeData(string ffProbePath, string filePath) => throw new NotImplementedException();
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Jellyfin.MediaEncoding.Keyframes.csproj b/src/Jellyfin.MediaEncoding.Keyframes/Jellyfin.MediaEncoding.Keyframes.csproj
new file mode 100644
index 000000000..7a984658b
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Jellyfin.MediaEncoding.Keyframes.csproj
@@ -0,0 +1,24 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>net5.0</TargetFramework>
+    <RootNamespace>Jellyfin.MediaEncoding.Keyframes</RootNamespace>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="NEbml" Version="0.11.0" />
+  </ItemGroup>
+
+  <!-- Code Analyzers-->
+  <ItemGroup Condition=" '$(Configuration)' == 'Debug' ">
+    <PackageReference Include="SerilogAnalyzer" Version="0.15.0" PrivateAssets="All" />
+    <PackageReference Include="StyleCop.Analyzers" Version="1.1.118" PrivateAssets="All" />
+    <PackageReference Include="SmartAnalyzers.MultithreadingAnalyzer" Version="1.1.31" PrivateAssets="All" />
+  </ItemGroup>
+  <ItemGroup>
+    <Reference Include="Microsoft.Extensions.Logging.Abstractions, Version=5.0.0.0, Culture=neutral, PublicKeyToken=adb9793829ddae60">
+      <HintPath>..\..\..\..\..\..\Program Files\dotnet\packs\Microsoft.AspNetCore.App.Ref\5.0.0\ref\net5.0\Microsoft.Extensions.Logging.Abstractions.dll</HintPath>
+    </Reference>
+  </ItemGroup>
+
+</Project>
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/KeyframeData.cs b/src/Jellyfin.MediaEncoding.Keyframes/KeyframeData.cs
new file mode 100644
index 000000000..3122f827c
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/KeyframeData.cs
@@ -0,0 +1,28 @@
+using System.Collections.Generic;
+
+namespace Jellyfin.MediaEncoding.Keyframes
+{
+    public class KeyframeData
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="KeyframeData"/> class.
+        /// </summary>
+        /// <param name="totalDuration">The total duration of the video stream in ticks.</param>
+        /// <param name="keyframeTicks">The video keyframes in ticks.</param>
+        public KeyframeData(long totalDuration, IReadOnlyList<long> keyframeTicks)
+        {
+            TotalDuration = totalDuration;
+            KeyframeTicks = keyframeTicks;
+        }
+
+        /// <summary>
+        /// Gets the total duration of the stream in ticks.
+        /// </summary>
+        public long TotalDuration { get; }
+
+        /// <summary>
+        /// Gets the keyframes in ticks.
+        /// </summary>
+        public IReadOnlyList<long> KeyframeTicks { get; }
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/KeyframeExtractor.cs b/src/Jellyfin.MediaEncoding.Keyframes/KeyframeExtractor.cs
new file mode 100644
index 000000000..2ee6b43e6
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/KeyframeExtractor.cs
@@ -0,0 +1,56 @@
+using System;
+using System.IO;
+using Jellyfin.MediaEncoding.Keyframes.FfProbe;
+using Jellyfin.MediaEncoding.Keyframes.FfTool;
+using Jellyfin.MediaEncoding.Keyframes.Matroska;
+using Microsoft.Extensions.Logging;
+
+namespace Jellyfin.MediaEncoding.Keyframes
+{
+    /// <summary>
+    /// Manager class for the set of keyframe extractors.
+    /// </summary>
+    public class KeyframeExtractor
+    {
+        private readonly ILogger<KeyframeExtractor> _logger;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="KeyframeExtractor"/> class.
+        /// </summary>
+        /// <param name="logger">An instance of the <see cref="ILogger{KeyframeExtractor}"/> interface.</param>
+        public KeyframeExtractor(ILogger<KeyframeExtractor> logger)
+        {
+            _logger = logger;
+        }
+
+        /// <summary>
+        /// Extracts the keyframe positions from a video file.
+        /// </summary>
+        /// <param name="filePath">Absolute file path to the media file.</param>
+        /// <param name="ffProbePath">Absolute file path to the ffprobe executable.</param>
+        /// <param name="ffToolPath">Absolute file path to the fftool executable.</param>
+        /// <returns></returns>
+        public KeyframeData GetKeyframeData(string filePath, string ffProbePath, string ffToolPath)
+        {
+            var extension = Path.GetExtension(filePath);
+            if (string.Equals(extension, ".mkv", StringComparison.OrdinalIgnoreCase))
+            {
+                try
+                {
+                    return MatroskaKeyframeExtractor.GetKeyframeData(filePath);
+                }
+                catch (InvalidOperationException ex)
+                {
+                    _logger.LogError(ex, "{MatroskaKeyframeExtractor} failed to extract keyframes", nameof(MatroskaKeyframeExtractor));
+                }
+            }
+
+            if (!string.IsNullOrEmpty(ffToolPath))
+            {
+                return FfToolKeyframeExtractor.GetKeyframeData(ffToolPath, filePath);
+            }
+
+            return FfProbeKeyframeExtractor.GetKeyframeData(ffProbePath, filePath);
+        }
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs
new file mode 100644
index 000000000..0de0f996c
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs
@@ -0,0 +1,181 @@
+using System;
+using Jellyfin.MediaEncoding.Keyframes.Matroska.Models;
+using NEbml.Core;
+
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska.Extensions
+{
+    /// <summary>
+    /// Extension methods for the <see cref="EbmlReader"/> class.
+    /// </summary>
+    internal static class EbmlReaderExtensions
+    {
+        /// <summary>
+        /// Traverses the current container to find the element with <paramref name="identifier"/> identifier.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <param name="identifier">The element identifier.</param>
+        /// <returns>A value indicating whether the element was found.</returns>
+        internal static bool FindElement(this EbmlReader reader, ulong identifier)
+        {
+            while (reader.ReadNext())
+            {
+                if (reader.ElementId.EncodedValue == identifier)
+                {
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Reads the current position in the file as an unsigned integer converted from binary.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <returns>The unsigned integer.</returns>
+        internal static uint ReadUIntFromBinary(this EbmlReader reader)
+        {
+            var buffer = new byte[4];
+            reader.ReadBinary(buffer, 0, 4);
+            if (BitConverter.IsLittleEndian)
+            {
+                Array.Reverse(buffer);
+            }
+
+            return BitConverter.ToUInt32(buffer);
+        }
+
+        /// <summary>
+        /// Reads from the start of the file to retrieve the SeekHead segment.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <returns>Instance of <see cref="SeekHead"/></returns>
+        internal static SeekHead ReadSeekHead(this EbmlReader reader)
+        {
+            reader = reader ?? throw new ArgumentNullException(nameof(reader));
+
+            if (reader.ElementPosition != 0)
+            {
+                throw new InvalidOperationException("File position must be at 0");
+            }
+
+            // Skip the header
+            if (!reader.FindElement(MatroskaConstants.SegmentContainer))
+            {
+                throw new InvalidOperationException("Expected a segment container");
+            }
+
+            reader.EnterContainer();
+
+            long? tracksPosition = null;
+            long? cuesPosition = null;
+            long? infoPosition = null;
+            // The first element should be a SeekHead otherwise we'll have to search manually
+            if (!reader.FindElement(MatroskaConstants.SeekHead))
+            {
+                throw new InvalidOperationException("Expected a SeekHead");
+            }
+
+            reader.EnterContainer();
+            while (reader.FindElement(MatroskaConstants.Seek))
+            {
+                reader.EnterContainer();
+                reader.ReadNext();
+                var type = (ulong)reader.ReadUIntFromBinary();
+                switch (type)
+                {
+                    case MatroskaConstants.Tracks:
+                        reader.ReadNext();
+                        tracksPosition = (long)reader.ReadUInt();
+                        break;
+                    case MatroskaConstants.Cues:
+                        reader.ReadNext();
+                        cuesPosition = (long)reader.ReadUInt();
+                        break;
+                    case MatroskaConstants.Info:
+                        reader.ReadNext();
+                        infoPosition = (long)reader.ReadUInt();
+                        break;
+                }
+
+                reader.LeaveContainer();
+
+                if (tracksPosition.HasValue && cuesPosition.HasValue && infoPosition.HasValue)
+                {
+                    break;
+                }
+            }
+
+            reader.LeaveContainer();
+
+            if (!tracksPosition.HasValue || !cuesPosition.HasValue || !infoPosition.HasValue)
+            {
+                throw new InvalidOperationException("SeekHead is missing or does not contain Info, Tracks and Cues positions");
+            }
+
+            return new SeekHead(infoPosition.Value, tracksPosition.Value, cuesPosition.Value);
+        }
+
+        /// <summary>
+        /// Reads from SegmentContainer to retrieve the Info segment.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <returns>Instance of <see cref="Info"/></returns>
+        internal static Info ReadInfo(this EbmlReader reader, long position)
+        {
+            reader.ReadAt(position);
+
+            double? duration = null;
+            reader.EnterContainer();
+            // Mandatory element
+            reader.FindElement(MatroskaConstants.TimestampScale);
+            var timestampScale = reader.ReadUInt();
+
+            if (reader.FindElement(MatroskaConstants.Duration))
+            {
+                duration = reader.ReadFloat();
+            }
+
+            reader.LeaveContainer();
+
+            return new Info((long)timestampScale, duration);
+        }
+
+        /// <summary>
+        /// Enters the Tracks segment and reads all tracks to find the specified type.
+        /// </summary>
+        /// <param name="reader">Instance of <see cref="EbmlReader"/>.</param>
+        /// <param name="tracksPosition">The relative position of the tracks segment.</param>
+        /// <param name="type">The track type identifier.</param>
+        /// <returns>The first track number with the specified type.</returns>
+        /// <exception cref="InvalidOperationException">Stream type is not found.</exception>
+        internal static ulong FindFirstTrackNumberByType(this EbmlReader reader, long tracksPosition, ulong type)
+        {
+            reader.ReadAt(tracksPosition);
+
+            reader.EnterContainer();
+            while (reader.FindElement(MatroskaConstants.TrackEntry))
+            {
+                reader.EnterContainer();
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.TrackNumber);
+                var trackNumber = reader.ReadUInt();
+
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.TrackType);
+                var trackType = reader.ReadUInt();
+
+                reader.LeaveContainer();
+                if (trackType == MatroskaConstants.TrackTypeVideo)
+                {
+                    reader.LeaveContainer();
+                    return trackNumber;
+                }
+            }
+
+            reader.LeaveContainer();
+
+            throw new InvalidOperationException($"No stream with type {type} found");
+        }
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaConstants.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaConstants.cs
new file mode 100644
index 000000000..d18418d45
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaConstants.cs
@@ -0,0 +1,31 @@
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska
+{
+    /// <summary>
+    /// Constants for the Matroska identifiers.
+    /// </summary>
+    public static class MatroskaConstants
+    {
+        internal const ulong SegmentContainer = 0x18538067;
+
+        internal const ulong SeekHead = 0x114D9B74;
+        internal const ulong Seek = 0x4DBB;
+
+        internal const ulong Info = 0x1549A966;
+        internal const ulong TimestampScale = 0x2AD7B1;
+        internal const ulong Duration = 0x4489;
+
+        internal const ulong Tracks = 0x1654AE6B;
+        internal const ulong TrackEntry = 0xAE;
+        internal const ulong TrackNumber = 0xD7;
+        internal const ulong TrackType = 0x83;
+
+        internal const ulong TrackTypeVideo = 0x1;
+        internal const ulong TrackTypeSubtitle = 0x11;
+
+        internal const ulong Cues = 0x1C53BB6B;
+        internal const ulong CueTime = 0xB3;
+        internal const ulong CuePoint = 0xBB;
+        internal const ulong CueTrackPositions = 0xB7;
+        internal const ulong CuePointTrackNumber = 0xF7;
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaKeyframeExtractor.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaKeyframeExtractor.cs
new file mode 100644
index 000000000..10d017d2a
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/MatroskaKeyframeExtractor.cs
@@ -0,0 +1,76 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using Jellyfin.MediaEncoding.Keyframes.Matroska.Extensions;
+using NEbml.Core;
+
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska
+{
+    /// <summary>
+    /// The keyframe extractor for the matroska container.
+    /// </summary>
+    public static class MatroskaKeyframeExtractor
+    {
+        /// <summary>
+        /// Extracts the keyframes in ticks (scaled using the container timestamp scale) from the matroska container.
+        /// </summary>
+        /// <param name="filePath">The file path.</param>
+        /// <returns>An instance of <see cref="KeyframeData"/>.</returns>
+        public static KeyframeData GetKeyframeData(string filePath)
+        {
+            using var stream = File.OpenRead(filePath);
+            using var reader = new EbmlReader(stream);
+
+            var seekHead = reader.ReadSeekHead();
+            var info = reader.ReadInfo(seekHead.InfoPosition);
+            var videoTrackNumber = reader.FindFirstTrackNumberByType(seekHead.TracksPosition, MatroskaConstants.TrackTypeVideo);
+
+            var keyframes = new List<long>();
+            reader.ReadAt(seekHead.CuesPosition);
+            reader.EnterContainer();
+
+            while (reader.FindElement(MatroskaConstants.CuePoint))
+            {
+                reader.EnterContainer();
+                ulong? trackNumber = null;
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.CueTime);
+                var cueTime = reader.ReadUInt();
+
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.CueTrackPositions);
+                reader.EnterContainer();
+                if (reader.FindElement(MatroskaConstants.CuePointTrackNumber))
+                {
+                    trackNumber = reader.ReadUInt();
+                }
+
+                reader.LeaveContainer();
+
+                if (trackNumber == videoTrackNumber)
+                {
+                    keyframes.Add(ScaleToNanoseconds(cueTime, info.TimestampScale));
+                }
+
+                reader.LeaveContainer();
+            }
+
+            reader.LeaveContainer();
+
+            var result = new KeyframeData(ScaleToNanoseconds(info.Duration ?? 0, info.TimestampScale), keyframes);
+            return result;
+        }
+
+        private static long ScaleToNanoseconds(ulong unscaledValue, long timestampScale)
+        {
+            // TimestampScale is in nanoseconds, scale it to get the value in ticks, 1 tick == 100 ns
+            return (long)unscaledValue * timestampScale / 100;
+        }
+
+        private static long ScaleToNanoseconds(double unscaledValue, long timestampScale)
+        {
+            // TimestampScale is in nanoseconds, scale it to get the value in ticks, 1 tick == 100 ns
+            return Convert.ToInt64(unscaledValue * timestampScale / 100);
+        }
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/Info.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/Info.cs
new file mode 100644
index 000000000..02c6741ec
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/Info.cs
@@ -0,0 +1,29 @@
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska.Models
+{
+    /// <summary>
+    /// The matroska Info segment.
+    /// </summary>
+    internal class Info
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="Info"/> class.
+        /// </summary>
+        /// <param name="timestampScale">The timestamp scale in nanoseconds.</param>
+        /// <param name="duration">The duration of the entire file.</param>
+        public Info(long timestampScale, double? duration)
+        {
+            TimestampScale = timestampScale;
+            Duration = duration;
+        }
+
+        /// <summary>
+        /// Gets the timestamp scale in nanoseconds.
+        /// </summary>
+        public long TimestampScale { get; }
+
+        /// <summary>
+        /// Gets the total duration of the file.
+        /// </summary>
+        public double? Duration { get; }
+    }
+}
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/SeekHead.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/SeekHead.cs
new file mode 100644
index 000000000..d9e346c03
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Models/SeekHead.cs
@@ -0,0 +1,36 @@
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska.Models
+{
+    /// <summary>
+    /// The matroska SeekHead segment. All positions are relative to the Segment container.
+    /// </summary>
+    internal class SeekHead
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="SeekHead"/> class.
+        /// </summary>
+        /// <param name="infoPosition">The relative file position of the info segment.</param>
+        /// <param name="tracksPosition">The relative file position of the tracks segment.</param>
+        /// <param name="cuesPosition">The relative file position of the cues segment.</param>
+        public SeekHead(long infoPosition, long tracksPosition, long cuesPosition)
+        {
+            InfoPosition = infoPosition;
+            TracksPosition = tracksPosition;
+            CuesPosition = cuesPosition;
+        }
+
+        /// <summary>
+        /// Gets relative file position of the info segment.
+        /// </summary>
+        public long InfoPosition { get; }
+
+        /// <summary>
+        /// Gets the relative file position of the tracks segment.
+        /// </summary>
+        public long TracksPosition { get; }
+
+        /// <summary>
+        /// Gets the relative file position of the cues segment.
+        /// </summary>
+        public long CuesPosition { get; }
+    }
+}