From abf03c279aa34c4b75b4fc2904ac7b5bc3bfa95a Mon Sep 17 00:00:00 2001
From: Breakthrough <brandon248@gmail.com>
Date: Sun, 29 Sep 2024 20:08:48 -0400
Subject: [PATCH] [cli] Add new `export-qp` command (#388)

---
 scenedetect.cfg              |  9 +++++
 scenedetect/_cli/__init__.py | 64 +++++++++++++++++++++++++++---------
 scenedetect/_cli/commands.py | 25 ++++++++++++++
 scenedetect/_cli/config.py   |  4 +++
 tests/test_cli.py            | 14 ++++++++
 5 files changed, 100 insertions(+), 16 deletions(-)

diff --git a/scenedetect.cfg b/scenedetect.cfg
index bde791de..28d2135f 100644
--- a/scenedetect.cfg
+++ b/scenedetect.cfg
@@ -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
 #
diff --git a/scenedetect/_cli/__init__.py b/scenedetect/_cli/__init__.py
index 208e4d1d..e81d264f 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))
@@ -1011,7 +1011,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(
@@ -1081,7 +1081,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(
@@ -1256,7 +1256,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(
@@ -1442,30 +1442,62 @@ 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],
+):
+    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),
+    }
+    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..09f2f31d 100644
--- a/scenedetect/_cli/commands.py
+++ b/scenedetect/_cli/commands.py
@@ -64,6 +64,31 @@ def export_html(
     )
 
 
+def save_qp(
+    context: CliContext,
+    scenes: SceneList,
+    cuts: CutList,
+    output_dir: str,
+    filename_format: str,
+):
+    """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,
+    )
+    with open(qp_path, "w") as qp_file:
+        # TODO(#388): Instead of setting start time, should we always start at 0 and shift each
+        # cut by the amount that was seeked?
+        first_frame = 0 if context.start_time is None else context.start_time.frame_num
+        # Place an initial I frame at the first frame.
+        qp_file.writelines([f"{first_frame} I -1"])
+        # Place another I frame at each detected cut.
+        qp_file.writelines(f"{cut.frame_num} I -1" 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 3ea5babe..9dda354e 100644
--- a/scenedetect/_cli/config.py
+++ b/scenedetect/_cli/config.py
@@ -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,
diff --git a/tests/test_cli.py b/tests/test_cli.py
index dbd9ef90..5eb258a9 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -431,6 +431,20 @@ 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."""
+    base_command = "-i {VIDEO} time {TIME} {DETECTOR} {COMMAND}"
+    assert invoke_scenedetect(base_command, COMMAND="save-qp", output_dir=tmp_path) == 0
+    assert (
+        invoke_scenedetect(
+            base_command, COMMAND="save-qp --filename custom.txt", output_dir=tmp_path
+        )
+        == 0
+    )
+    assert os.path.exists(tmp_path.joinpath(f"{DEFAULT_VIDEO_NAME}.qp"))
+    assert os.path.exists(tmp_path.joinpath("custom.txt"))
+
+
 @pytest.mark.parametrize("backend_type", ALL_BACKENDS)
 def test_cli_backend(backend_type: str):
     """Test setting the `-b`/`--backend` argument."""