[postprocessor:ugoira] support setting timecodes with mkvmerge

by selecting the "mkvmerge" demuxer

(#1550)

gallery_dl/postprocessor/ugoira.py

@@ -32,14 +32,26 @@ class UgoiraPP(PostProcessor):
         ffmpeg = options.get("ffmpeg-location")
         self.ffmpeg = util.expand_path(ffmpeg) if ffmpeg else "ffmpeg"
 
+        mkvmerge = options.get("mkvmerge-location")
+        self.mkvmerge = util.expand_path(mkvmerge) if mkvmerge else "mkvmerge"
+
         rate = options.get("framerate", "auto")
         if rate != "auto":
             self.calculate_framerate = lambda _: (None, rate)
 
-        if options.get("ffmpeg-demuxer") == "image2":
-            self._process = self._image2
+        demuxer = options.get("ffmpeg-demuxer")
+        if demuxer == "image2":
+            self._process = self._process_image2
+            self._finalize = None
+            self.log.debug("using image2 demuxer")
+        elif demuxer == "mkvmerge":
+            self._process = self._process_mkvmerge
+            self._finalize = self._finalize_mkvmerge
+            self.log.debug("using image2+mkvmerge demuxer")
         else:
-            self._process = self._concat
+            self._process = self._process_concat
+            self._finalize = None
+            self.log.debug("using concat demuxer")
 
         if options.get("libx264-prevent-odd", True):
             # get last video-codec argument
@@ -89,13 +101,12 @@ class UgoiraPP(PostProcessor):
                 return
 
             # process frames and collect command-line arguments
-            args = self._process(tempdir)
+            pathfmt.set_extension(self.extension)
+            args = self._process(pathfmt, tempdir)
             if self.args:
                 args += self.args
-            self.log.debug("ffmpeg args: %s", args)
 
             # invoke ffmpeg
-            pathfmt.set_extension(self.extension)
             try:
                 if self.twopass:
                     if "-f" not in self.args:
@@ -106,6 +117,8 @@ class UgoiraPP(PostProcessor):
                 else:
                     args.append(pathfmt.realpath)
                     self._exec(args)
+                if self._finalize:
+                    self._finalize(pathfmt, tempdir)
             except OSError as exc:
                 print()
                 self.log.error("Unable to invoke FFmpeg (%s: %s)",
@@ -121,37 +134,28 @@ class UgoiraPP(PostProcessor):
                 else:
                     pathfmt.set_extension("zip")
 
-    def _concat(self, path):
-        ffconcat = path + "/ffconcat.txt"
-
-        content = ["ffconcat version 1.0"]
-        append = content.append
-        for frame in self._frames:
-            append("file '{}'\nduration {}".format(
-                frame["file"], frame["delay"] / 1000))
-        if self.repeat:
-            append("file '{}'".format(frame["file"]))
-        append("")
-
-        with open(ffconcat, "w") as file:
-            file.write("\n".join(content))
+    def _exec(self, args):
+        self.log.debug(args)
+        out = None if self.output else subprocess.DEVNULL
+        return subprocess.Popen(args, stdout=out, stderr=out).wait()
 
+    def _process_concat(self, pathfmt, tempdir):
         rate_in, rate_out = self.calculate_framerate(self._frames)
         args = [self.ffmpeg, "-f", "concat"]
         if rate_in:
             args += ("-r", str(rate_in))
-        args += ("-i", ffconcat)
+        args += ("-i", self._write_ffmpeg_concat(tempdir))
         if rate_out:
             args += ("-r", str(rate_out))
         return args
 
-    def _image2(self, path):
-        path += "/"
+    def _process_image2(self, pathfmt, tempdir):
+        tempdir += "/"
 
         # adjust frame mtime values
         ts = 0
         for frame in self._frames:
-            os.utime(path + frame["file"], ns=(ts, ts))
+            os.utime(tempdir + frame["file"], ns=(ts, ts))
             ts += frame["delay"] * 1000000
 
         return [
@@ -160,12 +164,69 @@ class UgoiraPP(PostProcessor):
             "-ts_from_file", "2",
             "-pattern_type", "sequence",
             "-i", "{}%06d.{}".format(
-                path.replace("%", "%%"), frame["file"].rpartition(".")[2]),
+                tempdir.replace("%", "%%"),
+                frame["file"].rpartition(".")[2]
+            ),
         ]
 
-    def _exec(self, args):
-        out = None if self.output else subprocess.DEVNULL
-        return subprocess.Popen(args, stdout=out, stderr=out).wait()
+    def _process_mkvmerge(self, pathfmt, tempdir):
+        self._realpath = pathfmt.realpath
+        pathfmt.realpath = tempdir + "/temp." + self.extension
+
+        return [
+            self.ffmpeg,
+            "-f", "image2",
+            "-pattern_type", "sequence",
+            "-i", "{}/%06d.{}".format(
+                tempdir.replace("%", "%%"),
+                self._frames[0]["file"].rpartition(".")[2]
+            ),
+        ]
+
+    def _finalize_mkvmerge(self, pathfmt, tempdir):
+        args = [
+            self.mkvmerge,
+            "-o", self._realpath,
+            "--timecodes", "0:" + self._write_mkvmerge_timecodes(tempdir),
+        ]
+        if self.extension == "webm":
+            args.append("--webm")
+        args.append(pathfmt.realpath)
+
+        pathfmt.realpath = self._realpath
+        self._exec(args)
+
+    def _write_ffmpeg_concat(self, tempdir):
+        content = ["ffconcat version 1.0"]
+        append = content.append
+
+        for frame in self._frames:
+            append("file '{}'\nduration {}".format(
+                frame["file"], frame["delay"] / 1000))
+        if self.repeat:
+            append("file '{}'".format(frame["file"]))
+        append("")
+
+        ffconcat = tempdir + "/ffconcat.txt"
+        with open(ffconcat, "w") as file:
+            file.write("\n".join(content))
+        return ffconcat
+
+    def _write_mkvmerge_timecodes(self, tempdir):
+        content = ["# timecode format v2"]
+        append = content.append
+
+        delay_sum = 0
+        for frame in self._frames:
+            append(str(delay_sum))
+            delay_sum += frame["delay"]
+        append(str(delay_sum))
+        append("")
+
+        timecodes = tempdir + "/timecodes.tc"
+        with open(timecodes, "w") as file:
+            file.write("\n".join(content))
+        return timecodes
 
     @staticmethod
     def calculate_framerate(framelist):