Skip to content

Commit

Permalink
[cli] Add new save-qp command (#448)
Browse files Browse the repository at this point in the history
* [cli] Add new `save-qp` command (#388)

* [cli] Ensure `save-qp` command shifts frame numbers and add tests
  • Loading branch information
Breakthrough authored Nov 11, 2024
1 parent 64f1bae commit 770d514
Show file tree
Hide file tree
Showing 5 changed files with 166 additions and 16 deletions.
12 changes: 12 additions & 0 deletions scenedetect.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,18 @@
#start-col-name = Start Frame


[save-qp]

# Filename format of QP file. Can use $VIDEO_NAME macro.
#filename = $VIDEO_NAME.qp

# Folder to output QP file to. Overrides [global] output option.
#output = /usr/tmp/images

# Disable shifting frame numbers by start time (yes/no).
#disable-shift = no


#
# BACKEND OPTIONS
#
Expand Down
81 changes: 65 additions & 16 deletions scenedetect/_cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,7 @@ def scenedetect(
)
@click.pass_context
def help_command(ctx: click.Context, command_name: str):
"""Print help for command (`help [command]`)."""
"""Print full help reference."""
assert isinstance(ctx.parent.command, click.MultiCommand)
parent_command = ctx.parent.command
all_commands = set(parent_command.list_commands(ctx))
Expand Down Expand Up @@ -989,6 +989,9 @@ def export_html_command(
image_height: ty.Optional[int],
):
"""Export scene list to HTML file. Requires save-images unless --no-images is specified."""
# TODO: Rename this command to save-html to align with other export commands. This will require
# that we allow `export-html` as an alias on the CLI and via the config file for a few versions
# as to not break existing workflows.
ctx = ctx.obj
assert isinstance(ctx, CliContext)

Expand All @@ -1011,7 +1014,7 @@ def export_html_command(
"-o",
metavar="DIR",
type=click.Path(exists=False, dir_okay=True, writable=True, resolve_path=False),
help="Output directory to save videos to. Overrides global option -o/--output if set.%s"
help="Output directory to save videos to. Overrides global option -o/--output.%s"
% (USER_CONFIG.get_help_string("list-scenes", "output", show_default=False)),
)
@click.option(
Expand Down Expand Up @@ -1084,7 +1087,7 @@ def list_scenes_command(
"-o",
metavar="DIR",
type=click.Path(exists=False, dir_okay=True, writable=True, resolve_path=False),
help="Output directory to save videos to. Overrides global option -o/--output if set.%s"
help="Output directory to save videos to. Overrides global option -o/--output.%s"
% (USER_CONFIG.get_help_string("split-video", "output", show_default=False)),
)
@click.option(
Expand Down Expand Up @@ -1259,7 +1262,7 @@ def split_video_command(
"-o",
metavar="DIR",
type=click.Path(exists=False, dir_okay=True, writable=True, resolve_path=False),
help="Output directory for images. Overrides global option -o/--output if set.%s"
help="Output directory for images. Overrides global option -o/--output.%s"
% (USER_CONFIG.get_help_string("save-images", "output", show_default=False)),
)
@click.option(
Expand Down Expand Up @@ -1445,30 +1448,76 @@ def save_images_command(
ctx.save_images = True


@click.command("save-qp", cls=_Command)
@click.option(
"--filename",
"-f",
metavar="NAME",
default=None,
type=click.STRING,
help="Filename format to use.%s" % (USER_CONFIG.get_help_string("save-qp", "filename")),
)
@click.option(
"--output",
"-o",
metavar="DIR",
type=click.Path(exists=False, dir_okay=True, writable=True, resolve_path=False),
help="Output directory to save QP file to. Overrides global option -o/--output.%s"
% (USER_CONFIG.get_help_string("save-qp", "output", show_default=False)),
)
@click.option(
"--disable-shift",
"-d",
is_flag=True,
flag_value=True,
default=None,
help="Disable shifting frame numbers by start time.%s"
% (USER_CONFIG.get_help_string("save-qp", "disable-shift")),
)
@click.pass_context
def save_qp_command(
ctx: click.Context,
filename: ty.Optional[ty.AnyStr],
output: ty.Optional[ty.AnyStr],
disable_shift: ty.Optional[bool],
):
"""Save cuts as keyframes (I-frames) for video encoding.
The resulting QP file can be used with the `--qpfile` argument in x264/x265."""
ctx = ctx.obj
assert isinstance(ctx, CliContext)

save_qp_args = {
"filename_format": ctx.config.get_value("save-qp", "filename", filename),
"output_dir": ctx.config.get_value("save-qp", "output", output),
"shift_start": not ctx.config.get_value("save-qp", "disable-shift", disable_shift),
}
ctx.add_command(cli_commands.save_qp, save_qp_args)


# ----------------------------------------------------------------------
# Commands Omitted From Help List
# CLI Sub-Command Registration
# ----------------------------------------------------------------------

# Info Commands
# Informational
scenedetect.add_command(about_command)
scenedetect.add_command(help_command)
scenedetect.add_command(version_command)

# ----------------------------------------------------------------------
# Commands Added To Help List
# ----------------------------------------------------------------------

# Input / Output
scenedetect.add_command(export_html_command)
scenedetect.add_command(list_scenes_command)
# Input
scenedetect.add_command(load_scenes_command)
scenedetect.add_command(save_images_command)
scenedetect.add_command(split_video_command)
scenedetect.add_command(time_command)

# Detection Algorithms
# Detectors
scenedetect.add_command(detect_adaptive_command)
scenedetect.add_command(detect_content_command)
scenedetect.add_command(detect_hash_command)
scenedetect.add_command(detect_hist_command)
scenedetect.add_command(detect_threshold_command)

# Output
scenedetect.add_command(export_html_command)
scenedetect.add_command(save_qp_command)
scenedetect.add_command(list_scenes_command)
scenedetect.add_command(save_images_command)
scenedetect.add_command(split_video_command)
23 changes: 23 additions & 0 deletions scenedetect/_cli/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,29 @@ def export_html(
)


def save_qp(
context: CliContext,
scenes: SceneList,
cuts: CutList,
output_dir: str,
filename_format: str,
shift_start: bool,
):
"""Handler for the `save-qp` command."""
del scenes # We only use cuts for this handler.
qp_path = get_and_create_path(
Template(filename_format).safe_substitute(VIDEO_NAME=context.video_stream.name),
output_dir,
)
start_frame = context.start_time.frame_num if context.start_time else 0
offset = start_frame if shift_start else 0
with open(qp_path, "wt") as qp_file:
qp_file.write(f"{0 if shift_start else start_frame} I -1\n")
# Place another I frame at each detected cut.
qp_file.writelines(f"{cut.frame_num - offset} I -1\n" for cut in cuts)
logger.info(f"QP file written to: {qp_path}")


def list_scenes(
context: CliContext,
scenes: SceneList,
Expand Down
5 changes: 5 additions & 0 deletions scenedetect/_cli/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,11 @@ def format(self, timecode: FrameTimecode) -> str:
"scale-method": Interpolation.LINEAR,
"width": 0,
},
"save-qp": {
"disable-shift": False,
"filename": "$VIDEO_NAME.qp",
"output": None,
},
"split-video": {
"args": DEFAULT_FFMPEG_ARGS,
"copy": False,
Expand Down
61 changes: 61 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,67 @@ def test_cli_export_html(tmp_path: Path):
# TODO: Check for existence of HTML & image files.


def test_cli_save_qp(tmp_path: Path):
"""Test `save-qp` command with and without a custom filename format."""
EXPECTED_QP_CONTENTS = """
0 I -1
90 I -1
"""
for filename in (None, "custom.txt"):
filename_format = f"--filename {filename}" if filename else ""
assert (
invoke_scenedetect(
f"-i {{VIDEO}} time -e 95 {{DETECTOR}} save-qp {filename_format}",
output_dir=tmp_path,
)
== 0
)
output_path = tmp_path.joinpath(filename if filename else f"{DEFAULT_VIDEO_NAME}.qp")
assert os.path.exists(output_path)
assert output_path.read_text() == EXPECTED_QP_CONTENTS[1:]


def test_cli_save_qp_start_offset(tmp_path: Path):
"""Test `save-qp` command but using a shifted start time."""
# The QP file should always start from frame 0, so we expect a similar result to the above, but
# with the frame numbers shifted by the start frame. Note that on the command-line, the first
# frame is frame 1, but the first frame in a QP file is indexed by 0.
#
# Since we are starting at frame 51, we must shift all cuts by 50 frames.
EXPECTED_QP_CONTENTS = """
0 I -1
40 I -1
"""
assert (
invoke_scenedetect(
"-i {VIDEO} time -s 51 -e 95 {DETECTOR} save-qp",
output_dir=tmp_path,
)
== 0
)
output_path = tmp_path.joinpath(f"{DEFAULT_VIDEO_NAME}.qp")
assert os.path.exists(output_path)
assert output_path.read_text() == EXPECTED_QP_CONTENTS[1:]


def test_cli_save_qp_no_shift(tmp_path: Path):
"""Test `save-qp` command with start time shifting disabled."""
EXPECTED_QP_CONTENTS = """
50 I -1
90 I -1
"""
assert (
invoke_scenedetect(
"-i {VIDEO} time -s 51 -e 95 {DETECTOR} save-qp --disable-shift",
output_dir=tmp_path,
)
== 0
)
output_path = tmp_path.joinpath(f"{DEFAULT_VIDEO_NAME}.qp")
assert os.path.exists(output_path)
assert output_path.read_text() == EXPECTED_QP_CONTENTS[1:]


@pytest.mark.parametrize("backend_type", ALL_BACKENDS)
def test_cli_backend(backend_type: str):
"""Test setting the `-b`/`--backend` argument."""
Expand Down

0 comments on commit 770d514

Please sign in to comment.