Skip to content

Commit

Permalink
[pp:ugoira] initial support for "original" frames (#6056)
Browse files Browse the repository at this point in the history
  • Loading branch information
mikf committed Sep 1, 2024
1 parent 9d1e5f3 commit 57da9eb
Showing 1 changed file with 102 additions and 59 deletions.
161 changes: 102 additions & 59 deletions gallery_dl/postprocessor/ugoira.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def __init__(self, job, options):
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
self.uniform = self._convert_zip = self._convert_files = False

ffmpeg = options.get("ffmpeg-location")
self.ffmpeg = util.expand_path(ffmpeg) if ffmpeg else "ffmpeg"
Expand Down Expand Up @@ -90,87 +90,127 @@ def __init__(self, job, options):
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)
job.register_hooks({
"prepare": self.prepare,
"file" : self.convert_zip,
"after" : self.convert_files,
}, options)

def prepare(self, pathfmt):
self._frames = None

if pathfmt.extension != "zip":
if "frames" not in pathfmt.kwdict:
self._frames = None
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"]
self._frames = pathfmt.kwdict["frames"]
if pathfmt.extension == "zip":
self._convert_zip = True
if self.delete:
pathfmt.set_extension(self.extension)
pathfmt.build_path()
else:
return

if self.delete:
pathfmt.set_extension(self.extension)
pathfmt.build_path()
num = pathfmt.kwdict["num"]
if not num:
self._files = [pathfmt.realpath]
else:
self._files.append(pathfmt.realpath)
if num+1 >= len(self._frames):
self._convert_files = True

def convert(self, pathfmt):
if not self._frames:
def convert_zip(self, pathfmt):
if not self._convert_zip:
return
self._convert_zip = False

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)
self.log.debug("", exc_info=True)
pathfmt.realpath = pathfmt.temppath
else:
if self.mtime:
mtime = pathfmt.kwdict.get("_mtime")
if mtime:
util.set_mtime(pathfmt.realpath, mtime)
if self.convert(pathfmt, tempdir):
if self.delete:
pathfmt.delete = True
else:
self.log.info(pathfmt.filename)
pathfmt.set_extension("zip")
pathfmt.build_path()

def convert_files(self, pathfmt):
if not self._convert_files:
return
self._convert_files = False

with tempfile.TemporaryDirectory() as tempdir:
for frame, path in zip(self._frames, self._files):

# update frame filename extension
frame["file"] = name = "{}.{}".format(
frame["file"].partition(".")[0],
path.rpartition(".")[2])

# move frame into tempdir
try:
self._copy_file(path, tempdir + "/" + name)
except OSError as exc:
self.log.debug("Unable to copy frame %s (%s: %s)",
name, exc.__class__.__name__, exc)
return

pathfmt.kwdict["num"] = 0
if self.convert(pathfmt, tempdir):
self.log.info(pathfmt.filename)
if self.delete:
self.log.debug("Deleting frames")
for path in self._files:
util.remove_file(path)

def convert(self, pathfmt, tempdir):
pathfmt.set_extension(self.extension)
pathfmt.build_path()

# process frames and collect command-line arguments
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)
self.log.debug("", exc_info=True)
pathfmt.realpath = pathfmt.temppath
else:
if self.mtime:
mtime = pathfmt.kwdict.get("_mtime")
if mtime:
util.set_mtime(pathfmt.realpath, mtime)
return True

def _exec(self, args):
self.log.debug(args)
out = None if self.output else subprocess.DEVNULL
Expand All @@ -182,6 +222,9 @@ def _exec(self, args):
raise ValueError()
return retcode

def _copy_file(self, src, dst):
shutil.copyfile(src, dst)

def _process_concat(self, pathfmt, tempdir):
rate_in, rate_out = self.calculate_framerate(self._frames)
args = [self.ffmpeg, "-f", "concat"]
Expand Down

0 comments on commit 57da9eb

Please sign in to comment.