Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Snap start and end times of ASS subtitles to scene changes using WWXD
#!/usr/bin/env python3.6
"""Snap start and end times of ASS subtitles to scene changes using WWXD.
Does not work with variable frame rate (VFR).
usage: snap_scenechanges.py [-h] [--epsilon EPSILON] [-o OUTPUT] [-v]
sub_path video_path
positional arguments:
sub_path Path to subtitle file.
video_path Path to video file.
optional arguments:
-h, --help show this help message and exit
--epsilon EPSILON Number of frames to search for scene changes around
each frame.
-o OUTPUT, --output OUTPUT
Output path. By default, the input path with
`_snapped` appended.
-v, --verbose Increase verbosity.
Requires:
- Vapoursynth
- Vapoursynth WWXD plugin (https://github.com/dubhater/vapoursynth-wwxd)
- ffms2
*or*
L-SMASH plugin
- python-ass (https://github.com/rfw/python-ass)
"""
# Copyright 2017 FichteFoll <fichtefoll2@googlemail.com>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
# OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
# TODO
# - improve caching a lot
# use a list of len(clip) with tri-state logic
# and dump/load from or to file (disableable)
import argparse
from collections import defaultdict
from itertools import cycle, islice
import logging
import math
from pathlib import Path
import sys
from datetime import timedelta
from typing import List, Optional, Iterable
import ass
import vapoursynth
__version__ = "0.2.2"
l = logging.getLogger(__name__)
core = vapoursynth.get_core()
# itertools recipe
def roundrobin(*iterables: Iterable) -> Iterable:
"roundrobin('ABC', 'D', 'EF') --> A D E B F C"
# Recipe credited to George Sakkis
pending = len(iterables)
nexts = cycle(iter(it).__next__ for it in iterables)
while pending:
try:
for next_ in nexts:
yield next_()
except StopIteration:
pending -= 1
nexts = cycle(islice(nexts, pending))
def negate(iterable: Iterable) -> Iterable:
for i in iterable:
yield -i
def parse_args(args: Optional[List[str]]) -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Snap start and end times of subtitles to scene changes using WWXD."
)
parser.add_argument("sub_path", type=Path,
help="Path to subtitle file.")
parser.add_argument("video_path", type=Path,
help="Path to video file.")
parser.add_argument("--epsilon", type=int, default=7,
help="Number of frames to search for scene changes around each frame.")
parser.add_argument("-o", "--output", type=Path,
help="Output path. By default, the input path with `_snapped` appended.")
parser.add_argument("-v", "--verbose", action='store_true',
help="Increase verbosity.")
return parser.parse_args(args)
def time_to_frame(time: timedelta, framerate: float, floor: bool = True) -> int:
frame = time.total_seconds() * framerate
if floor:
return math.floor(frame)
else:
return math.ceil(frame)
def frame_to_time(frame: int, framerate: float, floor: bool = True) -> timedelta:
if floor:
middle_frame = max(0, frame - 0.5)
else:
middle_frame = frame + 0.5
secs = middle_frame / framerate
secs = round(secs, 2) # round to centiseconds because python-ass floors by default
return timedelta(seconds=secs)
# cache scene-changing frames for speedup (by clip)
_scene_change_cache = defaultdict(set)
def find_nearest_scenechange(clip, frame: int, epsilon: int, prefer_forward: bool = False) \
-> Optional[int]:
# We want an iterator like [-1, 1, -2, 2, -3, 3, -4, 4, -5, 5, -6, 6, -7, 7]
base_range = range(1, epsilon + 1)
offsets = roundrobin(base_range, negate(base_range))
if not prefer_forward:
offsets = negate(offsets)
else:
offsets = roundrobin(negate(base_range), negate(base_range))
last_frame = len(clip) - 1
for offset in offsets:
test_frame = frame + offset
if not (0 <= test_frame <= last_frame):
continue
if test_frame in _scene_change_cache[clip]:
return test_frame
is_scenechange = clip.get_frame(test_frame).props.Scenechange
if is_scenechange:
_scene_change_cache[clip].add(test_frame)
return test_frame
return None
def snap_keyframes(script, input_clip, epsilon: int):
# We resize to speed-up the wwxd filter.
# The target resultion is kinda arbitrary,
# but seems to produce the "best results".
# Additionally, we ensure that we have 420P8
# because scxvid requires that (for obvious reasons).
clip = input_clip.resize.Bilinear(640, 360, format=vapoursynth.YUV420P8)
clip = clip.wwxd.WWXD()
framerate = clip.fps_num / 1001
l.info("Video frame rate: %.3f", framerate)
l.info("Snapping...")
# Iterate over dialog lines and adjust timestamps
for event in script.events:
if event.TYPE != "Dialogue":
continue
l.debug("Checking line: %s", event.dump())
# Timestamps (should) always point between two frames.
#
# Our goal here is to have the first frame *after* the start time
# and the first frame *after* the end time to be a scene change.
# Thus, we ceil the timestamps when loading
# but floor the timestamps of scene changes we found when dumping.
start_frame = time_to_frame(event.start, framerate, floor=False)
end_frame = time_to_frame(event.end, framerate, floor=False)
nearest_start = find_nearest_scenechange(clip, start_frame, epsilon)
nearest_end = find_nearest_scenechange(clip, end_frame, epsilon)
if nearest_start is not None:
event.start = frame_to_time(nearest_start, framerate, floor=True)
l.debug(" Adjusted start by %+d frames", nearest_start - start_frame)
if nearest_end is not None:
event.end = frame_to_time(nearest_end, framerate, floor=True)
l.debug(" Adjusted end by %+d frames", nearest_end - end_frame)
l.debug(" New line: %s", event.dump())
l.debug("List of frames with scene changes: %s",
list(sorted(_scene_change_cache[clip])))
def main(args: List[str] = None) -> int:
params = parse_args(args)
# configure logging
l.addHandler(logging.StreamHandler())
log_level = logging.DEBUG if params.verbose else logging.INFO
l.setLevel(log_level)
l.debug("snap_scenechanges.py version %s", __version__)
# verify args
if not params.sub_path.suffix == ".ass":
l.error("Only works with Advanced Substation Alpha subtitles (.ass)")
return 1
if not params.output:
params.output = params.sub_path.with_name(
f"{params.sub_path.stem}_snapped{params.sub_path.suffix}"
)
l.info("Parsing subtitles...")
with params.sub_path.open(encoding='utf-8-sig') as f:
# I think ASS files universally use a BOM
script = ass.parse(f)
l.info("Opening video...")
if hasattr(core, 'lsmas'):
clip = core.lsmas.LWLibavSource(str(params.video_path))
elif hasattr(core, 'ffms2'):
l.info("Falling back to ffms2")
clip = core.ffms2.Source(str(params.video_path))
else:
l.error("ffms2 or L-SMASH plugins are not available")
return 2
# do wörk
snap_keyframes(script, clip, params.epsilon)
# save new file
l.info("Writing subtitles...")
with params.output.open('w', encoding='utf-8-sig') as f:
script.dump_file(f)
return 0
if __name__ == '__main__':
sys.exit(main())
@FichteFoll

This comment has been minimized.

Copy link
Owner Author

@FichteFoll FichteFoll commented Jan 29, 2017

The cache could be a lot better.

If you use this and would like to have better caching (see TODO comment), reach out to me in some way that are not these comments, because I am not notified for those. You'll figure something out.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.