From 770d5144efdbec47fbef1819ec69c18361dd1e46 Mon Sep 17 00:00:00 2001 From: Brandon Castellano Date: Sun, 10 Nov 2024 20:19:59 -0500 Subject: [PATCH] [cli] Add new save-qp command (#448) * [cli] Add new `save-qp` command (#388) * [cli] Ensure `save-qp` command shifts frame numbers and add tests --- scenedetect.cfg | 12 ++++++ scenedetect/_cli/__init__.py | 81 +++++++++++++++++++++++++++++------- scenedetect/_cli/commands.py | 23 ++++++++++ scenedetect/_cli/config.py | 5 +++ tests/test_cli.py | 61 +++++++++++++++++++++++++++ 5 files changed, 166 insertions(+), 16 deletions(-) diff --git a/scenedetect.cfg b/scenedetect.cfg index bde791de..c736138a 100644 --- a/scenedetect.cfg +++ b/scenedetect.cfg @@ -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 # diff --git a/scenedetect/_cli/__init__.py b/scenedetect/_cli/__init__.py index bc8a311e..f1a03f77 100644 --- a/scenedetect/_cli/__init__.py +++ b/scenedetect/_cli/__init__.py @@ -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)) @@ -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) @@ -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( @@ -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( @@ -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( @@ -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) diff --git a/scenedetect/_cli/commands.py b/scenedetect/_cli/commands.py index 83a23bf2..2271f429 100644 --- a/scenedetect/_cli/commands.py +++ b/scenedetect/_cli/commands.py @@ -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, diff --git a/scenedetect/_cli/config.py b/scenedetect/_cli/config.py index c854a6a2..7e7c06f7 100644 --- a/scenedetect/_cli/config.py +++ b/scenedetect/_cli/config.py @@ -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, diff --git a/tests/test_cli.py b/tests/test_cli.py index fcadb9bd..2bcd2435 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -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."""