Retain original color space when encoding, cache ffmpeg/ffprobe outputs

n00mkrad · May 6, 2021 · 55c0f59 · 55c0f59
1 parent 58e3f3e
commit 55c0f59
Show file tree

Hide file tree

Showing 11 changed files with 274 additions and 46 deletions.
diff --git a/Code/Data/ColorInfo.cs b/Code/Data/ColorInfo.cs
@@ -0,0 +1,92 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Flowframes.Data
+{
+    class ColorInfo
+    {
+        public string colorSpace;
+        public string colorRange;
+        public string colorTransfer;
+        public string colorPrimaries;
+
+        public ColorInfo(string ffprobeOutput)
+        {
+            string[] lines = ffprobeOutput.SplitIntoLines();
+
+            foreach (string line in lines)
+            {
+                if (line.Contains("color_range"))
+                {
+                    colorRange = line.Split('=').LastOrDefault();
+                    continue;
+                }
+
+                if (line.Contains("color_space"))
+                {
+                    colorSpace = line.Split('=').LastOrDefault();
+                    continue;
+                }
+
+                if (line.Contains("color_transfer"))
+                {
+                    colorTransfer = line.Split('=').LastOrDefault();
+                    continue;
+                }
+
+                if (line.Contains("color_primaries"))
+                {
+                    colorPrimaries = line.Split('=').LastOrDefault();
+                    continue;
+                }
+            }
+
+            if (colorSpace.Trim() == "unknown")
+                colorSpace = "";
+
+            if (colorRange.Trim() == "unknown")
+                colorRange = "";
+
+            if (colorTransfer.Trim() == "unknown")
+                colorTransfer = "";
+
+            if (colorPrimaries.Trim() == "unknown")
+                colorPrimaries = "";
+        }
+
+        public bool HasAnyValues ()
+        {
+            if (!string.IsNullOrWhiteSpace(colorSpace))
+                return true;
+
+            if (!string.IsNullOrWhiteSpace(colorRange))
+                return true;
+
+            if (!string.IsNullOrWhiteSpace(colorTransfer))
+                return true;
+
+            if (!string.IsNullOrWhiteSpace(colorPrimaries))
+                return true;
+
+            return false;
+        }
+
+        public bool HasAllValues()
+        {
+            if (string.IsNullOrWhiteSpace(colorSpace))
+                return false;
+
+            if (string.IsNullOrWhiteSpace(colorRange))
+                return false;
+
+            if (string.IsNullOrWhiteSpace(colorTransfer))
+                return false;
+
+            if (string.IsNullOrWhiteSpace(colorPrimaries))
+                return false;
+
+            return true;
+        }
+    }
+}
diff --git a/Code/Flowframes.csproj b/Code/Flowframes.csproj
@@ -328,6 +328,7 @@
   <ItemGroup>
     <Compile Include="Data\AI.cs" />
     <Compile Include="Data\AudioTrack.cs" />
+    <Compile Include="Data\ColorInfo.cs" />
     <Compile Include="Data\Fraction.cs" />
     <Compile Include="Data\InterpSettings.cs" />
     <Compile Include="Data\Networks.cs" />
@@ -387,8 +388,10 @@
     <Compile Include="Media\FFmpegUtils.cs" />
     <Compile Include="Media\GetFrameCountCached.cs" />
     <Compile Include="Media\GetMediaResolutionCached.cs" />
+    <Compile Include="Media\GetVideoInfoCached.cs" />
     <Compile Include="MiscUtils\Benchmarker.cs" />
     <Compile Include="MiscUtils\FrameRename.cs" />
+    <Compile Include="MiscUtils\NmkdStopwatch.cs" />
     <Compile Include="OS\AiProcess.cs" />
     <Compile Include="Extensions\ExtensionMethods.cs" />
     <Compile Include="Form1.cs">

diff --git a/Code/Main/AutoEncode.cs b/Code/Main/AutoEncode.cs
@@ -81,9 +81,6 @@ public static async Task MainLoop(string interpFramesPath)
 
                     unencodedFrameLines.Clear();
 
-                    //for (int frameLineNum = lastEncodedFrameNum; frameLineNum < interpFramesLines.Length; frameLineNum++)
-                    //        unencodedFrameLines.Add(frameLineNum);
-
                     bool aiRunning = !AiProcess.lastAiProcess.HasExited;
 
                     for (int frameLineNum = lastEncodedFrameNum; frameLineNum < interpFramesLines.Length; frameLineNum++)
@@ -103,7 +100,9 @@ public static async Task MainLoop(string interpFramesPath)
 
                             if (!File.Exists(lastOfChunk))
                             {
-                                Logger.Log($"[AutoEnc] Last frame of chunk doesn't exist; skipping loop iteration ({lastOfChunk})", true);
+                                if(debug)
+                                    Logger.Log($"[AutoEnc] Last frame of chunk doesn't exist; skipping loop iteration ({lastOfChunk})", true);
+
                                 await Task.Delay(500);
                                 continue;
                             }
@@ -152,7 +151,9 @@ public static async Task MainLoop(string interpFramesPath)
 
         static async Task DeleteOldFramesAsync (string interpFramesPath, List<int> frameLinesToEncode)
         {
-            Logger.Log("[AutoEnc] Starting DeleteOldFramesAsync.", true, false, "ffmpeg");
+            if(debug)
+                Logger.Log("[AutoEnc] Starting DeleteOldFramesAsync.", true, false, "ffmpeg");
+
             Stopwatch sw = new Stopwatch();
             sw.Restart();
 
@@ -165,7 +166,8 @@ static async Task DeleteOldFramesAsync (string interpFramesPath, List<int> frame
                 }
             }
 
-            Logger.Log("[AutoEnc] DeleteOldFramesAsync finished in " + FormatUtils.TimeSw(sw), true, false, "ffmpeg");
+            if (debug)
+                Logger.Log("[AutoEnc] DeleteOldFramesAsync finished in " + FormatUtils.TimeSw(sw), true, false, "ffmpeg");
         }
 
         static bool FrameIsStillNeeded (string frameName, int frameIndex)
@@ -180,7 +182,7 @@ public static bool HasWorkToDo ()
             if (Interpolate.canceled || interpFramesFolder == null) return false;
 
             if(debug)
-                Logger.Log($"HasWorkToDo - Process Running: {(AiProcess.lastAiProcess != null && !AiProcess.lastAiProcess.HasExited)} - encodedFrameLines.Count: {encodedFrameLines.Count} - interpFramesLines.Length: {interpFramesLines.Length}");
+                Logger.Log($"HasWorkToDo - Process Running: {(AiProcess.lastAiProcess != null && !AiProcess.lastAiProcess.HasExited)} - encodedFrameLines.Count: {encodedFrameLines.Count} - interpFramesLines.Length: {interpFramesLines.Length}", true);
 
             return ((AiProcess.lastAiProcess != null && !AiProcess.lastAiProcess.HasExited) || encodedFrameLines.Count < interpFramesLines.Length);
         }

diff --git a/Code/Main/CreateVideo.cs b/Code/Main/CreateVideo.cs
@@ -151,7 +151,8 @@ static async Task Encode(I.OutMode mode, string framesPath, string outPath, Frac
             }
             else
             {
-                await FfmpegEncode.FramesToVideo(framesFile, outPath, mode, fps, resampleFps);
+                ColorInfo colorInfo = await FfmpegCommands.GetColorInfo(I.current.inPath);
+                await FfmpegEncode.FramesToVideo(framesFile, outPath, mode, fps, resampleFps, colorInfo);
                 await MuxOutputVideo(I.current.inPath, outPath);
                 await Loop(currentOutFile, GetLoopTimes());
             }
@@ -212,18 +213,19 @@ public static async Task EncodeChunk(string outPath, I.OutMode mode, int firstFr
             string max = Config.Get("maxFps");
             Fraction maxFps = max.Contains("/") ? new Fraction(max) : new Fraction(max.GetFloat());
             bool fpsLimit = maxFps.GetFloat() != 0 && I.current.outFps.GetFloat() > maxFps.GetFloat();
+            ColorInfo colorInfo = await FfmpegCommands.GetColorInfo(I.current.inPath);
 
             bool dontEncodeFullFpsVid = fpsLimit && Config.GetInt("maxFpsMode") == 0;
 
             if (!dontEncodeFullFpsVid)
-                await FfmpegEncode.FramesToVideoConcat(framesFileChunk, outPath, mode, I.current.outFps, AvProcess.LogMode.Hidden, true);     // Encode
+                await FfmpegEncode.FramesToVideo(framesFileChunk, outPath, mode, I.current.outFps, new Fraction(), colorInfo, AvProcess.LogMode.Hidden, true);     // Encode
 
             if (fpsLimit)
             {
                 string filename = Path.GetFileName(outPath);
                 string newParentDir = outPath.GetParentDir() + Paths.fpsLimitSuffix;
                 outPath = Path.Combine(newParentDir, filename);
-                await FfmpegEncode.FramesToVideo(framesFileChunk, outPath, mode, I.current.outFps, maxFps, AvProcess.LogMode.Hidden, true);     // Encode with limited fps
+                await FfmpegEncode.FramesToVideo(framesFileChunk, outPath, mode, I.current.outFps, maxFps, colorInfo, AvProcess.LogMode.Hidden, true);     // Encode with limited fps
             }
         }
 

diff --git a/Code/Media/AvProcess.cs b/Code/Media/AvProcess.cs
@@ -167,20 +167,6 @@ static void FfmpegOutputHandlerSilent (object sendingProcess, DataReceivedEventA
             }
         }
 
-        public static string GetFfmpegOutput (string args)
-        {
-            Process ffmpeg = OSUtils.NewProcess(true);
-            lastAvProcess = ffmpeg;
-            ffmpeg.StartInfo.Arguments = $"{GetCmdArg()} cd /D {GetAvDir().Wrap()} & ffmpeg.exe -hide_banner -y -stats {args}";
-            Logger.Log($"ffmpeg {args}", true, false, "ffmpeg");
-            ffmpeg.Start();
-            ffmpeg.WaitForExit();
-            string output = ffmpeg.StandardOutput.ReadToEnd();
-            string err = ffmpeg.StandardError.ReadToEnd();
-            if (!string.IsNullOrWhiteSpace(err)) output += "\n" + err;
-            return output;
-        }
-
         public static async Task<string> GetFfmpegOutputAsync(string args, bool setBusy = false, bool progressBar = false)
         {
             timeSinceLastOutput.Restart();
@@ -192,11 +178,7 @@ public static async Task<string> GetFfmpegOutputAsync(string args, bool setBusy
             ffmpeg.StartInfo.Arguments = $"{GetCmdArg()} cd /D {GetAvDir().Wrap()} & ffmpeg.exe -hide_banner -y -stats {args}";
             Logger.Log($"ffmpeg {args}", true, false, "ffmpeg");
             if (setBusy) Program.mainForm.SetWorking(true);
-            ffmpeg.OutputDataReceived += FfmpegOutputHandlerSilent;
-            ffmpeg.ErrorDataReceived += FfmpegOutputHandlerSilent;
-            ffmpeg.Start();
-            ffmpeg.BeginOutputReadLine();
-            ffmpeg.BeginErrorReadLine();
+            lastOutputFfmpeg = await OSUtils.GetOutputAsync(ffmpeg);
             while (!ffmpeg.HasExited) await Task.Delay(50);
             while(timeSinceLastOutput.ElapsedMilliseconds < 200) await Task.Delay(50);
             if (setBusy) Program.mainForm.SetWorking(false);

diff --git a/Code/Media/FfmpegAudioAndMetadata.cs b/Code/Media/FfmpegAudioAndMetadata.cs
@@ -64,9 +64,8 @@ public static async Task ExtractAudioTracks(string inputFile, string outFolder,
         public static async Task<List<AudioTrack>> GetAudioTracks(string inputFile)
         {
             List<AudioTrack> audioTracks = new List<AudioTrack>();
-            string args = $"-i {inputFile.Wrap()}";
             Logger.Log("GetAudioTracks()", true, false, "ffmpeg");
-            string[] outputLines = (await GetFfmpegOutputAsync(args)).SplitIntoLines();
+            string[] outputLines = (await GetVideoInfoCached.GetFfmpegInfoAsync(inputFile)).SplitIntoLines();
 
             for (int i = 0; i < outputLines.Length; i++)
             {
@@ -142,9 +141,8 @@ public static async Task ExtractSubtitles(string inputFile, string outFolder, In
         public static async Task<List<SubtitleTrack>> GetSubtitleTracks(string inputFile)
         {
             List<SubtitleTrack> subtitleTracks = new List<SubtitleTrack>();
-            string args = $"-i {inputFile.Wrap()}";
             Logger.Log("GetSubtitleTracks()", true, false, "ffmpeg");
-            string[] outputLines = (await GetFfmpegOutputAsync(args)).SplitIntoLines();
+            string[] outputLines = (await GetVideoInfoCached.GetFfmpegInfoAsync(inputFile)).SplitIntoLines();
 
             for (int i = 0; i < outputLines.Length; i++)
             {

diff --git a/Code/Media/FfmpegCommands.cs b/Code/Media/FfmpegCommands.cs
@@ -10,7 +10,6 @@
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
-using Microsoft.VisualBasic.Logging;
 using static Flowframes.AvProcess;
 using Utils = Flowframes.Media.FFmpegUtils;
 
@@ -97,9 +96,7 @@ public static async Task<Fraction> GetFramerate(string inputFile)
                     Logger.Log("GetFramerate ffprobe Error: " + ffprobeEx.Message, true, false);
                 }
 
-
-                string ffmpegArgs = $" -i {inputFile.Wrap()}";
-                string ffmpegOutput = await GetFfmpegOutputAsync(ffmpegArgs);
+                string ffmpegOutput = await GetVideoInfoCached.GetFfmpegInfoAsync(inputFile);
                 string[] entries = ffmpegOutput.Split(',');
 
                 foreach (string entry in entries)
@@ -190,7 +187,7 @@ static async Task<int> ReadFrameCountFfprobeAsync(string inputFile, bool readFra
 
         static async Task<int> ReadFrameCountFfmpegAsync (string inputFile)
         {
-            string args = $" -loglevel panic -i {inputFile.Wrap()} -map 0:v:0 -c copy -f null - ";
+            string args = $" -loglevel panic -stats -i {inputFile.Wrap()} -map 0:v:0 -c copy -f null - ";
             string info = await GetFfmpegOutputAsync(args, true, true);
             try
             {
@@ -204,6 +201,14 @@ static async Task<int> ReadFrameCountFfmpegAsync (string inputFile)
             }
         }
 
+        public static async Task<ColorInfo> GetColorInfo(string inputFile)
+        {
+            string ffprobeOutput = await GetVideoInfoCached.GetFfprobeInfoAsync(inputFile, "color");
+            ColorInfo colorInfo = new ColorInfo(ffprobeOutput);
+            Logger.Log($"Created ColorInfo - Range: {colorInfo.colorRange} - Space: {colorInfo.colorSpace} - Transer: {colorInfo.colorTransfer} - Primaries: {colorInfo.colorPrimaries}", true, false, "ffmpeg");
+            return colorInfo;
+        }
+
         public static async Task<bool> IsEncoderCompatible(string enc)
         {
             Logger.Log($"IsEncoderCompatible('{enc}')", true, false, "ffmpeg");

diff --git a/Code/Media/FfmpegEncode.cs b/Code/Media/FfmpegEncode.cs
@@ -17,12 +17,7 @@ namespace Flowframes.Media
 {
     partial class FfmpegEncode : FfmpegCommands
     {
-        public static async Task FramesToVideoConcat(string framesFile, string outPath, Interpolate.OutMode outMode, Fraction fps, LogMode logMode = LogMode.OnlyLastLine, bool isChunk = false)
-        {
-            await FramesToVideo(framesFile, outPath, outMode, fps, new Fraction(), logMode, isChunk);
-        }
-
-        public static async Task FramesToVideo(string framesFile, string outPath, Interpolate.OutMode outMode, Fraction fps, Fraction resampleFps, LogMode logMode = LogMode.OnlyLastLine, bool isChunk = false)
+        public static async Task FramesToVideo(string framesFile, string outPath, Interpolate.OutMode outMode, Fraction fps, Fraction resampleFps, ColorInfo colors, LogMode logMode = LogMode.OnlyLastLine, bool isChunk = false)
         {
             if (logMode != LogMode.Hidden)
                 Logger.Log((resampleFps.GetFloat() <= 0) ? "Encoding video..." : $"Encoding video resampled to {resampleFps.GetString()} FPS...");
@@ -39,9 +34,22 @@ public static async Task FramesToVideo(string framesFile, string outPath, Interp
                     inArg = $"-i {Path.GetFileName(framesFile) + Paths.symlinksSuffix}/%{Padding.interpFrames}d.png";
             }
 
-            string rate = fps.ToString().Replace(",", ".");
-            string vf = (resampleFps.GetFloat() < 0.1f) ? "" : $"-vf fps=fps={resampleFps}";
             string extraArgs = Config.Get("ffEncArgs");
+            string rate = fps.ToString().Replace(",", ".");
+
+            List<string> filters = new List<string>();
+
+            if (resampleFps.GetFloat() >= 0.1f)
+                filters.Add($"fps=fps={resampleFps}");
+
+            if (colors.HasAllValues())
+            {
+                Logger.Log($"Applying color transfer ({colors.colorSpace}).", true, false, "ffmpeg");
+                filters.Add($"scale=out_color_matrix={colors.colorSpace}");
+                extraArgs += $" -colorspace {colors.colorSpace} -color_primaries {colors.colorPrimaries} -color_trc {colors.colorTransfer} -color_range:v \"{colors.colorRange}\"";
+            }
+
+            string vf = filters.Count > 0 ? $"-vf {string.Join(",", filters)}" : "";
             string args = $"-vsync 0 -r {rate} {inArg} {encArgs} {vf} {extraArgs} -threads {Config.GetInt("ffEncThreads")} {outPath.Wrap()}";
             await RunFfmpeg(args, framesFile.GetParentDir(), logMode, "error", TaskType.Encode, !isChunk);
             IOUtils.TryDeleteIfExists(linksDir);