Skip to content

Commit

Permalink
[cli] Add new save-qp command (#388)
Browse files Browse the repository at this point in the history
  • Loading branch information
Breakthrough committed Oct 20, 2024
1 parent 3fe6ffc commit 30dce12
Show file tree
Hide file tree
Showing 5 changed files with 137 additions and 16 deletions.
9 changes: 9 additions & 0 deletions scenedetect.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,15 @@
#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


#
# BACKEND OPTIONS
#
Expand Down
72 changes: 56 additions & 16 deletions scenedetect/_cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -348,7 +348,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 @@ -991,6 +991,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 @@ -1013,7 +1016,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 @@ -1086,7 +1089,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 @@ -1261,7 +1264,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 @@ -1447,30 +1450,67 @@ 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.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),
"disable_shift": 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)
25 changes: 25 additions & 0 deletions scenedetect/_cli/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,31 @@ def export_html(
)


def save_qp(
context: CliContext,
scenes: SceneList,
cuts: CutList,
output_dir: str,
filename_format: str,
disable_shift: 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,
)
# We assume that if a start time was set, the user will encode the video starting from that
# point, so we shift all frame numbers such that the QP file always starts from frame 0.
with open(qp_path, "wt") as qp_file:
qp_file.write("0 I -1\n")
# Place another I frame at each detected cut.
offset = 0 if context.start_time is None else context.start_time.frame_num
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
4 changes: 4 additions & 0 deletions scenedetect/_cli/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,10 @@ def format(self, timecode: FrameTimecode) -> str:
"image-width": 0,
"no-images": False,
},
"save-qp": {
"filename": "$VIDEO_NAME.qp",
"output": None,
},
"list-scenes": {
"cut-format": "timecode",
"display-cuts": True,
Expand Down
43 changes: 43 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,49 @@ 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_shift(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:]


@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 30dce12

Please sign in to comment.