1
0
mirror of https://github.com/mikf/gallery-dl.git synced 2024-11-26 04:32:51 +01:00
gallery-dl/gallery_dl/postprocessor/ugoira.py
2024-04-06 17:00:48 +02:00

317 lines
10 KiB
Python

# -*- coding: utf-8 -*-
# Copyright 2018-2023 Mike Fährmann
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
"""Convert Pixiv Ugoira to WebM"""
from .common import PostProcessor
from .. import util
import subprocess
import tempfile
import zipfile
import shutil
import os
try:
from math import gcd
except ImportError:
def gcd(a, b):
while b:
a, b = b, a % b
return a
class UgoiraPP(PostProcessor):
def __init__(self, job, options):
PostProcessor.__init__(self, job)
self.extension = options.get("extension") or "webm"
self.args = options.get("ffmpeg-args") or ()
self.twopass = options.get("ffmpeg-twopass", False)
self.output = options.get("ffmpeg-output", "error")
self.delete = not options.get("keep-files", False)
self.repeat = options.get("repeat-last-frame", True)
self.mtime = options.get("mtime", True)
self.uniform = False
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"
demuxer = options.get("ffmpeg-demuxer")
if demuxer is None or demuxer == "auto":
if self.extension in ("webm", "mkv") and (
mkvmerge or shutil.which("mkvmerge")):
demuxer = "mkvmerge"
else:
demuxer = "concat"
if demuxer == "mkvmerge":
self._process = self._process_mkvmerge
self._finalize = self._finalize_mkvmerge
elif demuxer == "image2":
self._process = self._process_image2
self._finalize = None
else:
self._process = self._process_concat
self._finalize = None
self.log.debug("using %s demuxer", demuxer)
rate = options.get("framerate", "auto")
if rate == "uniform":
self.uniform = True
elif rate != "auto":
self.calculate_framerate = lambda _: (None, rate)
if options.get("libx264-prevent-odd", True):
# get last video-codec argument
vcodec = None
for index, arg in enumerate(self.args):
arg, _, stream = arg.partition(":")
if arg == "-vcodec" or arg in ("-c", "-codec") and (
not stream or stream.partition(":")[0] in ("v", "V")):
vcodec = self.args[index + 1]
# use filter when using libx264/5
self.prevent_odd = (
vcodec in ("libx264", "libx265") or
not vcodec and self.extension.lower() in ("mp4", "mkv"))
else:
self.prevent_odd = False
self.args_pp = args = []
if isinstance(self.output, str):
args += ("-hide_banner", "-loglevel", self.output)
if self.prevent_odd:
args += ("-vf", "crop=iw-mod(iw\\,2):ih-mod(ih\\,2)")
job.register_hooks(
{"prepare": self.prepare, "file": self.convert}, options)
def prepare(self, pathfmt):
self._frames = None
if pathfmt.extension != "zip":
return
kwdict = pathfmt.kwdict
if "frames" in kwdict:
self._frames = kwdict["frames"]
elif "pixiv_ugoira_frame_data" in kwdict:
self._frames = kwdict["pixiv_ugoira_frame_data"]["data"]
else:
return
if self.delete:
pathfmt.set_extension(self.extension)
pathfmt.build_path()
def convert(self, pathfmt):
if not self._frames:
return
with tempfile.TemporaryDirectory() as tempdir:
# extract frames
try:
with zipfile.ZipFile(pathfmt.temppath) as zfile:
zfile.extractall(tempdir)
except FileNotFoundError:
pathfmt.realpath = pathfmt.temppath
return
# process frames and collect command-line arguments
pathfmt.set_extension(self.extension)
pathfmt.build_path()
args = self._process(pathfmt, tempdir)
if self.args_pp:
args += self.args_pp
if self.args:
args += self.args
# ensure target directory exists
os.makedirs(pathfmt.realdirectory, exist_ok=True)
# invoke ffmpeg
try:
if self.twopass:
if "-f" not in self.args:
args += ("-f", self.extension)
args += ("-passlogfile", tempdir + "/ffmpeg2pass", "-pass")
self._exec(args + ["1", "-y", os.devnull])
self._exec(args + ["2", pathfmt.realpath])
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)",
exc.__class__.__name__, exc)
pathfmt.realpath = pathfmt.temppath
except Exception as exc:
print()
self.log.error("%s: %s", exc.__class__.__name__, exc)
pathfmt.realpath = pathfmt.temppath
else:
if self.mtime:
mtime = pathfmt.kwdict.get("_mtime")
if mtime:
util.set_mtime(pathfmt.realpath, mtime)
if self.delete:
pathfmt.delete = True
else:
pathfmt.set_extension("zip")
pathfmt.build_path()
def _exec(self, args):
self.log.debug(args)
out = None if self.output else subprocess.DEVNULL
retcode = util.Popen(args, stdout=out, stderr=out).wait()
if retcode:
print()
self.log.error("Non-zero exit status when running %s (%s)",
args, retcode)
raise ValueError()
return retcode
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", self._write_ffmpeg_concat(tempdir))
if rate_out:
args += ("-r", str(rate_out))
return args
def _process_image2(self, pathfmt, tempdir):
tempdir += "/"
frames = self._frames
# add extra frame if necessary
if self.repeat and not self._delay_is_uniform(frames):
last = frames[-1]
delay_gcd = self._delay_gcd(frames)
if last["delay"] - delay_gcd > 0:
last["delay"] -= delay_gcd
self.log.debug("non-uniform delays; inserting extra frame")
last_copy = last.copy()
frames.append(last_copy)
name, _, ext = last_copy["file"].rpartition(".")
last_copy["file"] = "{:>06}.{}".format(int(name)+1, ext)
shutil.copyfile(tempdir + last["file"],
tempdir + last_copy["file"])
# adjust frame mtime values
ts = 0
for frame in frames:
os.utime(tempdir + frame["file"], ns=(ts, ts))
ts += frame["delay"] * 1000000
return [
self.ffmpeg,
"-f", "image2",
"-ts_from_file", "2",
"-pattern_type", "sequence",
"-i", "{}%06d.{}".format(
tempdir.replace("%", "%%"),
frame["file"].rpartition(".")[2]
),
]
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", pathfmt.path, # mkvmerge does not support "raw" paths
"--timecodes", "0:" + self._write_mkvmerge_timecodes(tempdir),
]
if self.extension == "webm":
args.append("--webm")
args += ("=", 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
def calculate_framerate(self, frames):
if self._delay_is_uniform(frames):
return ("1000/{}".format(frames[0]["delay"]), None)
if not self.uniform:
gcd = self._delay_gcd(frames)
if gcd >= 10:
return (None, "1000/{}".format(gcd))
return (None, None)
@staticmethod
def _delay_gcd(frames):
result = frames[0]["delay"]
for f in frames:
result = gcd(result, f["delay"])
return result
@staticmethod
def _delay_is_uniform(frames):
delay = frames[0]["delay"]
for f in frames:
if f["delay"] != delay:
return False
return True
__postprocessor__ = UgoiraPP