From 546597ccede85dfff4ff5d389bb5e4b6a0df79a0 Mon Sep 17 00:00:00 2001 From: moi15moi <80980684+moi15moi@users.noreply.github.com> Date: Wed, 27 Dec 2023 20:26:00 -0500 Subject: [PATCH] [timestamps] from_video_file - Use mkvextract if the file is a mkv ffprobe can be slow for big mkv file (1 gb or more) mkvextract is way faster, so if it is available, we use it --- pyonfx/timestamps.py | 156 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 141 insertions(+), 15 deletions(-) diff --git a/pyonfx/timestamps.py b/pyonfx/timestamps.py index db75f959..57c6599d 100644 --- a/pyonfx/timestamps.py +++ b/pyonfx/timestamps.py @@ -26,6 +26,8 @@ from decimal import Decimal from fractions import Fraction from io import StringIO, TextIOWrapper +from platform import system +from tempfile import gettempdir from typing import Any, Callable, List, Optional, Tuple, Union @@ -403,7 +405,9 @@ def from_timestamps_file( To extract the timestamps file, you have 2 options: - Open the video with Aegisub. "Video" --> "Save Timecodes File"; - - Using [gMKVExtractGUI](https://sourceforge.net/projects/gmkvextractgui/) (warning: it will produce one timestamp too many at the end of the file, and you will need to manually remove it). + - Using [gMKVExtractGUI](https://sourceforge.net/projects/gmkvextractgui/) + Warning: it will produce one timestamp too many at the end of the file, and you will need to manually remove it + See: https://gitlab.com/mbunkus/mkvtoolnix/-/issues/3075 Parameters: path_to_timestamps_file_or_content (str | os.PathLike[str]): @@ -436,12 +440,12 @@ def from_timestamps_file( ) @classmethod - def from_video_file( + def _from_ffprobe( cls: Timestamps, - video_path: str, - index: Optional[int] = 0, - normalize: Optional[bool] = True, - rounding_method: Optional[RoundingMethod] = RoundingMethod.ROUND, + video_path: os.PathLike[str], + index: int, + normalize: bool, + rounding_method: RoundingMethod, ) -> Timestamps: """Create timestamps based on the ``video_path`` provided. @@ -449,7 +453,7 @@ def from_video_file( This method requires the ``ffprobe`` program to be available. Parameters: - video_path (str): A Video path. + video_path (os.PathLike[str]): A Video path. index (int, optional): Stream index of the video. normalize (bool, optional): If True, it will shift the timestamps to make them start from 0. If false, the option does nothing. rounding_method (RoundingMethod, optional): A rounding method. See the comment in Timestamps description about floor vs round. @@ -480,14 +484,6 @@ def get_timestamps(packets) -> Tuple[Fraction, Fraction, List[int]]: if shutil.which("ffprobe") is None: raise Exception("ffprobe is not in the environment variable.") - if not os.path.isfile(video_path): - raise FileNotFoundError(f'Invalid path for the video file: "{video_path}"') - - # Getting video absolute path and checking for its existance - if not os.path.isabs(video_path): - dirname = os.path.dirname(os.path.abspath(sys.argv[0])) - video_path = os.path.join(dirname, video_path) - cmd = [ "ffprobe", "-select_streams", @@ -535,6 +531,136 @@ def get_timestamps(packets) -> Tuple[Fraction, Fraction, List[int]]: last_frame_time=last_frame_time, ) + @staticmethod + def _from_mkvextract( + mkvextract_path: os.PathLike[str], + mkvmerge_path: os.PathLike[str], + video_path: os.PathLike[str], + index: int, + normalize: bool, + rounding_method: RoundingMethod, + ) -> Timestamps: + """Create timestamps based on the ``video_path`` provided. + + Note: + This method requires the ``mkvextract`` and ``mkvmerge`` programs to be available. + + Parameters: + mkvextract_path (os.PathLike[str]): The path to mkvextract. + mkvmerge_path (os.PathLike[str]): The path to mkvmerge. + video_path (os.PathLike[str]): A Video path. + index (int, optional): Stream index of the video. + normalize (bool, optional): If True, it will shift the timestamps to make them start from 0. If false, the option does nothing. + rounding_method (RoundingMethod, optional): A rounding method. See the comment in Timestamps description about floor vs round. + Returns: + An Timestamps instance. + """ + + cmd = [ + mkvmerge_path, + video_path, + "-J", + ] + mkvmerge_output = subprocess.run(cmd, capture_output=True, text=True) + if mkvmerge_output.returncode == 2: + raise ValueError(f"mkvmerge reported this error: {mkvmerge_output.stdout}") + mkvmerge_output_dict = json.loads(mkvmerge_output.stdout) + + is_index_in_video = False + for track in mkvmerge_output_dict["tracks"]: + if track["id"] == index: + if track["type"] != "video": + raise ValueError( + f'The index {index} is not a video stream. It is an "{track["type"]}" stream.' + ) + is_index_in_video = True + break + + if not is_index_in_video: + raise ValueError(f"The index {index} is not in the file {video_path}.") + + temp_dir = gettempdir() + timestamps_file_path = os.path.join(temp_dir, "temp_timestamps.txt") + cmd = [ + mkvextract_path, + video_path, + "timestamps_v2", + "0:" + timestamps_file_path + ] + + mkvextract_output = subprocess.run(cmd, capture_output=True, text=True) + if mkvextract_output.returncode == 2: + raise ValueError(f"mkvextract reported this error: {mkvextract_output.stdout}") + + with open(timestamps_file_path, "r") as file: + lines = file.readlines() + + os.remove(timestamps_file_path) + + # Ignore the last line due to this issue: https://gitlab.com/mbunkus/mkvtoolnix/-/issues/3075 + content = "".join(lines[:-1]) + + return Timestamps.from_timestamps_file( + content, + normalize=normalize, + rounding_method=rounding_method + ) + + @classmethod + def from_video_file( + cls: Timestamps, + video_path: os.PathLike[str], + index: Optional[int] = 0, + normalize: Optional[bool] = True, + rounding_method: Optional[RoundingMethod] = RoundingMethod.ROUND, + ) -> Timestamps: + """Create timestamps based on the ``video_path`` provided. + + Note: + This method requires the ``ffprobe`` or ``mkvextract/mkvmerge`` (if your if file is a mkv) programs to be available. + + Parameters: + video_path (os.PathLike[str]): A Video path. + index (int, optional): Stream index of the video. + normalize (bool, optional): If True, it will shift the timestamps to make them start from 0. If false, the option does nothing. + rounding_method (RoundingMethod, optional): A rounding method. See the comment in Timestamps description about floor vs round. + Returns: + An Timestamps instance. + """ + + if not os.path.isfile(video_path): + raise FileNotFoundError(f'Invalid path for the video file: "{video_path}"') + + # Getting video absolute path and checking for its existance + if not os.path.isabs(video_path): + dirname = os.path.dirname(os.path.abspath(sys.argv[0])) + video_path = os.path.join(dirname, video_path) + + with open(video_path, "rb") as f: + # From https://en.wikipedia.org/wiki/List_of_file_signatures + is_mkv = f.read(4) == b"\x1a\x45\xdf\xa3" + + mkvextract_path = None + if is_mkv: + # Verify if mkvextract and mkvmerge are installed + mkvextract_path = shutil.which("mkvextract") + mkvmerge_path = shutil.which("mkvmerge") + if mkvextract_path is None and system() == "Windows": + for possible_path in ["C:\\Program Files\\MKVToolNix\\mkvextract.exe", "C:\\Program Files (x86)\\MKVToolNix\\mkvextract.exe"]: + if os.path.isfile(os.path.join(possible_path)): + mkvextract_path = possible_path + break + if mkvmerge_path is None and system() == "Windows": + for possible_path in ["C:\\Program Files\\MKVToolNix\\mkvmerge.exe", "C:\\Program Files (x86)\\MKVToolNix\\mkvmerge.exe"]: + if os.path.isfile(os.path.join(possible_path)): + mkvmerge_path = possible_path + break + + if is_mkv and mkvextract_path is not None and mkvmerge_path is not None: + return Timestamps._from_mkvextract(mkvextract_path, mkvmerge_path, video_path, index, normalize, rounding_method) + else: + return Timestamps._from_ffprobe(video_path, index, normalize, rounding_method) + @staticmethod def normalize( timestamps: List[int], last_frame_time: Fraction