Add more flags

Author: WOSSFOSS

src/muxtools/create_offset_mkv.py

@@ -1,25 +1,40 @@
 import asyncio
 from typing import Annotated
 
-from typer import Argument
+from typer import Argument, Option
 from ..app import app
-from .find_change_from_start import find_change_from_start_inner
+from .find_change_from_start import (
+    find_change_from_start_inner,
+    find_change_from_first_frame,
+)
 
 
-async def create_offset_mkv_inner(bd_path: str, target_path: str, output_path: str):
-    bd_offset_frame, bd_fps = find_change_from_start_inner(bd_path)
+async def create_offset_mkv_inner(
+    bd_path: str,
+    target_path: str,
+    output_path: str,
+    output_offset: bool,
+    silent: bool,
+    *,
+    keep_video: bool = False,
+    output_offset_frame_number: bool = False,
+):
+    bd_offset_frame, bd_fps, first_frame = find_change_from_start_inner(bd_path)
     if bd_offset_frame == -1:
         raise ValueError("No significant change found in BD video.")
     bd_offset_seconds = bd_offset_frame / bd_fps
-    target_offset_frame, target_fps = find_change_from_start_inner(target_path)
+    target_offset_frame, target_fps = find_change_from_first_frame(
+        target_path, first_frame
+    )
     if target_offset_frame == -1:
         raise ValueError("No significant change found in target video.")
     target_offset_seconds = target_offset_frame / target_fps
     # Positive if BD starts later, negative if earlier
     offset_seconds = round(bd_offset_seconds - target_offset_seconds, 5)
-    print(
-        f"Calculated offset: {offset_seconds} seconds, adjusting all non-video streams accordingly."
-    )
+    if not silent:
+        print(
+            f"Calculated offset: {offset_seconds} seconds, adjusting all non-video streams accordingly."
+        )
     ffmpeg_command = [
         "ffmpeg",
         "-y",
@@ -31,8 +46,13 @@ async def create_offset_mkv_inner(bd_path: str, target_path: str, output_path: s
         target_path,
         "-map",
         "0",
-        "-map",
-        "-0:v",  # Exclude original video stream
+    ]
+    if not keep_video:
+        ffmpeg_command += [
+            "-map",
+            "-0:v",  # Exclude original video stream
+        ]
+    ffmpeg_command += [
         "-map_metadata",
         "0",
         "-c",
@@ -45,6 +65,12 @@ async def create_offset_mkv_inner(bd_path: str, target_path: str, output_path: s
     await proc.communicate()
     if proc.returncode != 0:
         raise RuntimeError(f"ffmpeg command failed with return code {proc.returncode}")
+    if output_offset:
+        if output_offset_frame_number:
+            offset_frames = round(offset_seconds * target_fps)
+            print(offset_frames)
+        else:
+            print(offset_seconds)
 
 
 @app.command()
@@ -77,15 +103,53 @@ def create_offset_mkv(
             help="Path to save the output MKV file with adjusted timing.",
         ),
     ],
+    output_offset: Annotated[
+        bool,
+        Option(
+            help="If true, output the adjusted timing in seconds.",
+            is_flag=True,
+        ),
+    ] = False,
+    output_offset_frame_number: Annotated[
+        bool,
+        Option(
+            help="If true, output the adjusted timing in frame numbers instead of seconds.",
+            is_flag=True,
+        ),
+    ] = False,
+    silent: Annotated[
+        bool,
+        Option(
+            help="If true, suppress output messages (other than the offset if requested).",
+            is_flag=True,
+        ),
+    ] = False,
+    keep_video: Annotated[
+        bool,
+        Option(
+            help="If true, keep the original video stream in the output file.",
+            is_flag=True,
+        ),
+    ] = False,
 ):
     """Create an MKV file with adjusted timing based on offset from Blu-ray video.
 
     :param bd_path: Path to the Blu-ray video file to analyze for offset.
     :param target_path: Path to the target video file to adjust.
     :param output_path: Path to save the output MKV file with adjusted timing.
+    :param output_offset: If true, output the adjusted timing.
+    :param output_offset_frame_number: If true, output the adjusted timing in frame numbers instead of seconds.
+    :param silent: If true, suppress output messages (other than the offset if requested).
+    :param keep_video: If true, keep the original video stream in the output file.
     """
     asyncio.run(
         create_offset_mkv_inner(
-            bd_path=bd_path, target_path=target_path, output_path=output_path
+            bd_path=bd_path,
+            target_path=target_path,
+            output_path=output_path,
+            output_offset=output_offset,
+            silent=silent,
+            keep_video=keep_video,
+            output_offset_frame_number=output_offset_frame_number,
         )
     )

src/muxtools/find_change_from_start.py

@@ -1,42 +1,32 @@
 import cv2
 
-from typing import Annotated
+from typing import Annotated, TypedDict, TYPE_CHECKING
 
 from typer import Argument, Option
 from ..app import app
 
+if TYPE_CHECKING:
+    from cv2.typing import MatLike
+    from cv2 import UMat
 
-def find_change_from_start_inner(
-    video_path, threshold=0, min_changed_pixels=50
-) -> tuple[int, float]:
-    """
-    Consumes video one frame at a time to find where it diverges from Frame 1.
-
-    Args:
-        video_path (str): Path to video file.
-        threshold (int): Sensitivity (0-255). Lower = detects subtle changes.
-                         Higher = ignores compression artifacts.
-        min_changed_pixels (int): How many pixels must change to trigger detection.
+    Mat = MatLike | UMat
 
-    Returns:
-        tuple[int, float]: Frame number where change is detected and FPS of the video.
-                           Returns -1 if no significant change is found.
-    """
-
-    # 1. Open the video stream
-    cap = cv2.VideoCapture(video_path)
 
-    try:
+def find_change_from_first_frame(
+    video_path: str | cv2.VideoCapture,
+    frame: "Mat",
+    threshold: int = 0,
+    min_changed_pixels: int = 50,
+) -> tuple[int, float]:
+    if isinstance(video_path, str):
+        cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise FileNotFoundError(f"Could not open video file: {video_path}")
-
-        # 2. Read Frame 1 (The Reference)
-        ret, frame1 = cap.read()
-        if not ret:
-            raise ValueError("Video file is empty or unreadable.")
-
+    else:
+        cap = video_path
+    try:
         # Convert to grayscale to reduce complexity and ignore color noise
-        frame1_gray = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
+        frame1_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
 
         # Optional: Apply slight blur to reduce compression artifact noise
         frame1_gray = cv2.GaussianBlur(frame1_gray, (21, 21), 0)
@@ -83,6 +73,42 @@ def find_change_from_start_inner(
         cap.release()
 
 
+def find_change_from_start_inner(
+    video_path, threshold=0, min_changed_pixels=50
+) -> tuple[int, float, "Mat"]:
+    """
+    Consumes video one frame at a time to find where it diverges from Frame 1.
+
+    Args:
+        video_path (str): Path to video file.
+        threshold (int): Sensitivity (0-255). Lower = detects subtle changes.
+                         Higher = ignores compression artifacts.
+        min_changed_pixels (int): How many pixels must change to trigger detection.
+
+    Returns:
+        tuple[int, float, Mat]:
+        - Frame number where significant change is detected, or -1 if none found.
+        - Frames per second of the video.
+        - The first frame of the video as a Mat object.
+    """
+
+    # 1. Open the video stream
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        raise FileNotFoundError(f"Could not open video file: {video_path}")
+
+    # 2. Read Frame 1 (The Reference)
+    ret, frame1 = cap.read()
+    if not ret:
+        raise ValueError("Video file is empty or unreadable.")
+
+    frame_num, fps = find_change_from_first_frame(
+        cap, frame1, threshold=threshold, min_changed_pixels=min_changed_pixels
+    )
+    return frame_num + 1, fps, frame1  # +1 to account for the first frame read earlier
+
+
 @app.command()
 def find_change_from_start(
     video_path: Annotated[
@@ -111,7 +137,7 @@ def find_change_from_start(
     :param threshold: Sensitivity (0-255). Lower = detects subtle changes. Higher = ignores compression artifacts.
     :param min_changed_pixels: How many pixels must change to trigger detection.
     """
-    frame_diff_at, fps = find_change_from_start_inner(
+    frame_diff_at, fps, _ = find_change_from_start_inner(
         video_path=video_path,
         threshold=threshold,
         min_changed_pixels=min_changed_pixels,