From 2590faf20f233a6cf528a1a1900e7d0b7c137038 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Tronel?= Date: Sun, 30 Aug 2026 08:58:23 +0200 Subject: [PATCH] The full pipeline is now in its own module. --- src/tscut/cli.py | 515 ++-------------------------------- src/tscut/config.py | 21 ++ src/tscut/matroska/codec.py | 7 +- src/tscut/pipeline.py | 514 +++++++++++++++++++++++++++++++++ src/tscut/tools/mkvtoolnix.py | 6 +- 5 files changed, 566 insertions(+), 497 deletions(-) create mode 100644 src/tscut/config.py create mode 100644 src/tscut/pipeline.py diff --git a/src/tscut/cli.py b/src/tscut/cli.py index 9cfc290..5a4dc6a 100644 --- a/src/tscut/cli.py +++ b/src/tscut/cli.py @@ -6,46 +6,17 @@ import argparse import logging -import os.path -from datetime import datetime, timedelta -from os import unlink -from shutil import copyfile, move -from sys import exit +from pathlib import Path +from datetime import timedelta from functools import cmp_to_key -import coloredlogs -import hexdump -from .tscut import ( - SupportedFormat, - change_codec_private_data, - check_required_tools, - compare_time_interval, - concatenate_h264_parts, - concatenate_h264_ts_parts, - do_ocr, - dump_codec_private_data, - extract_all_streams, - extract_mkv_part, - extract_srt, - extract_track_from_mkv, - ffmpeg_convert, - find_subtitles_tracks, - get_avc_config_from_h264, - get_codec_private_data_from_mkv, - get_format, - get_frame_rate, - get_movie_duration, - get_nearest_iframe, - get_streams, - get_tesseract_supported_lang, - get_ts_frame, - merge_mkvs, - parse_codec_private, - parse_time_interval, - remove_video_tracks_from_mkv, - remux_srt_subtitles, -) +import coloredlogs + +from tscut.config import ProcessingOptions +from tscut.tools.discovery import check_required_tools +from tscut.tools.timeframe import parse_time_interval, compare_time_interval +from tscut.pipeline import process_recording logger = logging.getLogger(__name__) @@ -123,460 +94,20 @@ def main() -> int: exit(-1) prevts = ts2 - nb_parts = len(parts) - temporaries = [] - - basename = os.path.splitext(os.path.basename(args.input_file))[0] - mp4filename = basename+'.mp4' - mkvfilename = basename+'.mkv' - - try: - input_file = open(args.input_file, mode='rb') - logger.debug("Type of input file: %s", type(input_file)) - except OSError: - logger.error("Impossible to open %s", args.input_file) - exit(-1) - - format_of_file = get_format(paths['ffprobe'], input_file) - - if format_of_file is None: - exit(-1) - - duration = timedelta(seconds=float(format_of_file['duration'])) - logger.info("Durée de l'enregistrement: %s", duration) - - if args.framerate is None: - framerate = get_frame_rate(paths['ffprobe'], input_file) - if framerate is None: - logger.error('Impossible to estimate frame rate !') - exit(-1) - else: - framerate = args.framerate - - logger.info('Frame rate: %.1f fps', framerate) - - found = False - for f in SupportedFormat: - if 'format_name' in format_of_file: - if format_of_file['format_name'] == str(f): - found = True - format_of_file = f - break - - if not found: - logger.error('Unsupported format of file') - - if format_of_file == SupportedFormat.TS: - logger.info("Converting TS to MP4 (to fix timestamps).") - try: - with open(mp4filename, 'wb+') as mp4: - ffmpeg_convert(paths['ffmpeg'], paths['ffprobe'], input_file, 'mpegts', mp4, 'mp4', - duration) - temporaries.append(mp4) - logger.info("Converting MP4 to MKV.") - try: - mkv = open(mkvfilename, 'wb+') - except OSError: - logger.error('') - - ffmpeg_convert(paths['ffmpeg'], paths['ffprobe'], mp4, 'mp4', mkv, 'matroska', - duration) - if nb_parts > 0: - temporaries.append(mkv) - except OSError: - logger.error('') - - elif format_of_file == SupportedFormat.MP4: - logger.info("Converting MP4 to MKV") - try: - mkv = open(mkvfilename, 'wb+') - except OSError: - logger.error('') - ffmpeg_convert(paths['ffmpeg'], paths['ffprobe'], input_file, 'mp4', mkv, 'matroska', - duration) - if nb_parts > 0: - temporaries.append(mkv) - else: - logger.info("Already in MKV") - mkv = input_file - - streams = get_streams(paths['ffprobe'], mkv) - - logger.debug('Streams: %s', streams) - main_video = None - nb_videos = 0 - for stream in streams: - if stream['codec_type'] == 'video': - if stream['disposition']['default'] == 1: - main_video = stream - width = stream['width'] - height = stream['height'] - break - nb_videos+=1 - if nb_videos == 1: - main_video = stream - width = stream['width'] - height = stream['height'] - else: - main_video = None - - if main_video is None: - logger.error('Impossible to find main video stream.') - exit(-1) - - # We retrieve the main private codec data - _, main_codec_private_data = get_codec_private_data_from_mkv(mkvinfo_path=paths['mkvinfo'], - input_file=mkv) - logger.debug('Main video stream has following private data: %s', - hexdump.dump(main_codec_private_data, sep=':')) - - # We parse them - main_avc_config = parse_codec_private(main_codec_private_data) - logger.debug('AVC configuration: %s', main_avc_config) - - # We check if the parse and dump operations are idempotent. - private_data = dump_codec_private_data(main_avc_config) - logger.debug('Redump AVC configuration: %s', hexdump.dump(private_data, sep=':')) - # In rare occasion, the PPS has trailing zeroes that do not seem to be related to useful data - # but they differ from the private data we generate that do not contain them. - # In that case we try to redecode our own private data to see if both AVC configurations are - # the same. - if main_codec_private_data != private_data: - logger.warning('Difference detected in bitstream !!') - iso_avc_config = parse_codec_private(private_data) - logger.debug('Reread AVC configuration: %s', iso_avc_config) - # If there exists a difference between our own reconstructed AVC configuration and the - # original one, we abandon - if iso_avc_config != main_avc_config: - logger.error('AVC configurations are different: %s\n%s\n', main_avc_config, - iso_avc_config) - exit(-1) - - # Pour chaque portion - partnum = 0 - mkvparts = [] - h264parts = [] - h264_ts = [] - checks = [] - pos = timedelta() - - other_avc_configs = [] - - for ts1, ts2 in parts: - # TODO: translate comment in english - # Trouver l'estampille de la trame 'I' la plus proche (mais postérieure) au début - # de la portion. - # Trouver l'estampille de la trame 'I' la plus proche (mais antérieure) à la fin - # de la portion. - # On a alors - # debut ----- trame --------- trame --------- fin fin+1 - # 'B/P' 'B/P'* 'I' 'I' 'B/P'* 'B/P' 'I/B/P' - # Si la trame de début est déjà 'I', il n'y a rien à faire. - # Sinon on extrait les trames 'B' ou 'P' depuis le début jusqu'à la trame 'I' non incluse. - # Si la trame de fin précède une trame I, on n'a rien à faire. - # Sinon on extrait toutes les trames depuis la dernière trame I jusqu'à la trame de fin. - - partnum = partnum + 1 - - # Get the nearest I-frame whose timestamp is greater or equal to the beginning. - head_frames = get_nearest_iframe(paths['ffprobe'], mkv, ts1, before=False) - if head_frames is None: - logger.error('Impossible to retrieve I-frame') - exit(-1) - - # Get the nearest I-frame whose timestamp ... - # TODO: wrong here ... - tail_frames = get_nearest_iframe(paths['ffprobe'], mkv, ts2, before=True) - if tail_frames is None: - logger.error('Impossible to retrieve I-frame') - exit(-1) - - nb_head_frames, head_iframe = head_frames - nb_tail_frames, tail_iframe = tail_frames - - logger.info("Found %d frames between beginning of current part and first I-frame", - nb_head_frames) - logger.info("Found %d frames between last I-frame and end of current part", - nb_tail_frames) - - head_iframe_ts = get_ts_frame(head_iframe) - if head_iframe_ts is None: - exit(-1) - tail_iframe_ts = get_ts_frame(tail_iframe) - if tail_iframe_ts is None: - exit(-1) - - checks.append(pos+head_iframe_ts-ts1) - - subparts = [] - - # TODO: separate pipeline processing between coarse and not fine grain options. - - # if args.coarse: - # do_coarse_processing(ffmpeg=paths['ffmpeg'], ffprobe=paths['ffprobe'], input_file=mkv, - # begin=ts1, end=head_iframe_ts, nb_frames=nb_head_frames-1, - # frameRate=frameRate, files_prefix='part-%d-head' % (partnum), - # streams=streams, width=width, height=height, - # temporaries=temporaries, dump_mem_fd=args.dump) - # else: - # doFineGrainProcessing(ffmpeg=paths['ffmpeg'], ffprobe=paths['ffprobe'], - # input_file=mkv, begin=ts1, end=head_iframe_ts, - # nb_frames=nb_head_frames-1, frameRate=frameRate, - # files_prefix='part-%d-head' % (partnum), streams=streams, - # width=width, height=height, temporaries=temporaries, - # dump_mem_fd=args.dump) - - if (not args.coarse) and (nb_head_frames > args.threshold): - # We extract all frames between the beginning upto the frame that immediately preceeds - # the I-frame. - h264_head, h264_head_ts, mkv_head = extract_all_streams(ffmpeg_path=paths['ffmpeg'], - ffprobe_path=paths['ffprobe'], - input_file=mkv, begin=ts1, - end=head_iframe_ts, - nb_frames=nb_head_frames-1, - framerate=framerate, - files_prefix=f'part-{partnum:d}-head', - streams=streams, width=width, - height=height, - temporaries=temporaries, - dump_mem_fd=args.dump) - - # If we are not at an exact boundary: - if mkv_head is not None: - subparts.append(mkv_head) - if h264_head is not None: - avcconfig = get_avc_config_from_h264(h264_head) - other_avc_configs.append(avcconfig) - h264parts.append(h264_head) - if h264_head_ts is not None: - h264_ts.append(h264_head_ts) - - # Creating MKV file that corresponds to current part between I-frames - # Internal video with all streams (video, audio and subtitles) - internal_mkv_name = f'part-{partnum:d}-internal.mkv' - # Internal video stream as a raw H264 stream - internal_h264_name = f'part-{partnum:d}-internal.h264' - # Internal video timestamps - internal_h264_ts_name = f'part-{partnum:d}-internal-ts.txt' - # Internal video with only audio and subtitles streams - internal_novideo_mkv_name = f'part-{partnum:d}-internal-novideo.mkv' - - try: - internal_mkv = open(internal_mkv_name, 'wb+') - except OSError: - logger.error('Impossible to create file: %s', internal_mkv_name) - exit(-1) - - try: - internal_novideo_mkv = open(internal_novideo_mkv_name, 'wb+') - except OSError: - logger.error('Impossible to create file: %s', internal_novideo_mkv_name) - exit(-1) - - try: - internal_h264 = open(internal_h264_name, 'wb+') - except OSError: - logger.error('Impossible to create file: %s', internal_h264_name) - exit(-1) - - try: - internal_h264_ts = open(internal_h264_ts_name, 'w+', encoding='utf8') - except OSError: - logger.error('Impossible to create file: %s', internal_h264_ts_name) - exit(-1) - - # logger.info('Merge header, middle and trailer subpart into: %s' % internal_mkv_name) - # Extract internal part of MKV - extract_mkv_part(mkvmerge_path=paths['mkvmerge'], input_file=mkv, output_file=internal_mkv, - begin=head_iframe_ts, end=tail_iframe_ts) - - # Extract video stream of internal part as a raw H264 and its timestamps. - logger.info('Extract video track as raw H264 file.') - extract_track_from_mkv(mkvextract_path=paths['mkvextract'], input_file=internal_mkv, - index=0, output_file=internal_h264, timestamps=internal_h264_ts) - - # Remove video track from internal part of MKV - logger.info('Remove video track from %s', internal_mkv_name) - remove_video_tracks_from_mkv(mkvmerge_path=paths['mkvmerge'], input_file=internal_mkv, - output_file=internal_novideo_mkv) - - temporaries.append(internal_mkv) - temporaries.append(internal_h264) - temporaries.append(internal_h264_ts) - temporaries.append(internal_novideo_mkv) - - h264parts.append(internal_h264) - h264_ts.append(internal_h264_ts) - subparts.append(internal_novideo_mkv) - - if (not args.coarse) and (nb_tail_frames > args.threshold): - # We extract all frames between the I-frame (including it) upto the end. - h264_tail, h264_tail_ts, mkv_tail = extract_all_streams(ffmpeg_path=paths['ffmpeg'], - ffprobe_path=paths['ffprobe'], - input_file=mkv, begin=tail_iframe_ts, - end=ts2, nb_frames=nb_tail_frames, - framerate=framerate, - files_prefix=f'part-{partnum:d}-tail', - streams=streams, - width=width, height=height, - temporaries=temporaries, - dump_mem_fd=args.dump) - - if mkv_tail is not None: - subparts.append(mkv_tail) - if h264_tail is not None: - avcconfig = get_avc_config_from_h264(h264_tail) - other_avc_configs.append(avcconfig) - h264parts.append(h264_tail) - if h264_tail_ts is not None: - h264_ts.append(h264_tail_ts) - - logger.info('Merging MKV: %s', subparts) - - part = merge_mkvs(mkvmerge_path=paths['mkvmerge'], inputs=subparts, - output_name=f'part-{partnum:d}.mkv', concatenate=True) - mkvparts.append(part) - temporaries.append(part) - - pos = pos+tail_iframe_ts-ts1 - - # We need to check the end also - checks.append(pos) - - # When using coarse option there is a single AVC configuration. - for avc_config in other_avc_configs: - main_avc_config.merge(avc_config) - logger.debug('Merged AVC configuration: %s', main_avc_config) - - nb_mkv_parts = len(mkvparts) - if nb_mkv_parts > 0: - try: - full_h264 = open(f'{basename}-full.h264', 'wb+') - except OSError: - logger.error('Impossible to create file full H264 stream.') - exit(-1) - - logger.info('Merging all H264 tracks') - concatenate_h264_parts(h264parts=h264parts, output=full_h264) - temporaries.append(full_h264) - - try: - full_h264_ts = open(f'{basename}-ts.txt', 'w+', encoding='utf8') - except OSError: - logger.error('Impossible to create file containing all video timestamps.') - exit(-1) - - logger.info('Merging H264 timestamps') - concatenate_h264_ts_parts(h264_ts_parts=h264_ts, output=full_h264_ts) - temporaries.append(full_h264_ts) - - final_novideo_name = f'{basename}-novideo.mkv' - final_with_video_name = f'{basename}-video.mkv' - - if nb_mkv_parts > 1: - logger.info('Merging all audio and subtitles parts: %s', mkvparts) - merge_mkvs(mkvmerge_path=paths['mkvmerge'], inputs=mkvparts, output_name=final_novideo_name, - concatenate=True) - elif nb_mkv_parts == 1: - copyfile('part-1.mkv', final_novideo_name) - else: - logger.info("Nothing else to do.") - copyfile(mkvfilename, final_with_video_name) - - if nb_mkv_parts >=1 : - try: - final_novideo = open(final_novideo_name, 'rb') - except OSError: - logger.error('Impossible to open file: %s.', final_novideo_name) - exit(-1) - - temporaries.append(final_novideo) - - full_h264_ts.seek(0) - - logger.info('Merging final video track and all other tracks together') - final_with_video = merge_mkvs(mkvmerge_path=paths['mkvmerge'], inputs=[full_h264, - final_novideo], - output_name=final_with_video_name, concatenate=False, - timestamps={0: full_h264_ts}) - final_codec_private_data = dump_codec_private_data(main_avc_config) - logger.debug('Final codec private data: %s', hexdump.dump(final_codec_private_data, - sep=':')) - logger.info('Changing codec private data with the new one.') - change_codec_private_data(paths['mkvinfo'], final_with_video, final_codec_private_data) - - if args.srt: - if not all_optional_tools: - logger.warning("Missing tools for extracting subtitles.") - move(final_with_video_name, args.output_file) - else: - # Final cut is not any more the final step. - temporaries.append(final_with_video) - duration = get_movie_duration(paths['ffprobe'], final_with_video) - supported_langs = get_tesseract_supported_lang(paths['tesseract']) - logger.info('Supported lang: %s', supported_langs) - logger.info('Find subtitles tracks and language.') - subtitles = find_subtitles_tracks(paths['ffprobe'], final_with_video) - logger.info(subtitles) - sts = {} - for subtitle in subtitles: - index = subtitle['index'] - if 'tags' in subtitle: - if 'language' in subtitle['tags']: - lang = subtitle['tags']['language'] - if lang in sts: - sts[lang].append(index) - else: - sts[lang] = [index] - else: - logger.error("Dropping subtitle: %s because it is missing language\ - indication", subtitle) - else: - logger.error("Dropping subtitle: %s because it is missing language indication", - subtitle) - - logger.info(sts) - if len(sts) > 0: - logger.info('Supported languages: %s', supported_langs) - list_of_subtitles = extract_srt(paths['mkvextract'], final_with_video_name, sts, - supported_langs) - logger.info(list_of_subtitles) - for idx_name, sub_name, _, _ in list_of_subtitles: - try: - idx = open(idx_name,'rb') - except OSError: - logger.error("Impossible to open %s.", idx_name) - exit(-1) - try: - sub = open(sub_name,'rb') - except OSError: - logger.error("Impossible to open %s.", sub_name) - exit(-1) - - temporaries.append(idx) - temporaries.append(sub) - - ocr = do_ocr(paths['vobsubocr'], list_of_subtitles, duration, temporaries, - args.dump) - logger.info(ocr) - - # Remux SRT subtitles - remux_srt_subtitles(paths['mkvmerge'], final_with_video, args.output_file, ocr) - else: - copyfile(final_with_video_name, args.output_file) - else: - move(final_with_video_name, args.output_file) - - if not args.keep: - logger.info("Cleaning temporary files") - for f in temporaries: - path = os.path.realpath(f.name) - logger.info("Removing: %s", path) - f.close() - unlink(path) - - d = datetime(1,1,1) - for c in checks: - logger.info("Please check cut smoothness at %s", (c+d).strftime("%H:%M:%S")) + config = ProcessingOptions( + input_file=Path(args.input_file), + output_file=Path(args.output_file), + parts=parts, + verbose=args.verbose, + coarse=args.coarse, + threshold=args.threshold, + keep_temporaries=args.keep, + subtitles_ocr=args.srt, + dump_memory=args.dump, + tools_paths = paths, + all_optional_tools = all_optional_tools) + + logger.debug("Configuration: %s", config) + process_recording(config) return 0 diff --git a/src/tscut/config.py b/src/tscut/config.py new file mode 100644 index 0000000..6c209b2 --- /dev/null +++ b/src/tscut/config.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: GPL-2.0-or-later +# +# Copyright (C) 2026 Frédéric Tronel + +from dataclasses import dataclass +from pathlib import Path + +@dataclass +class ProcessingOptions: + input_file: Path + output_file: Path + parts: list + tools_paths: dict[str, str] + all_optional_tools: bool + framerate: int|None = None + threshold: int = 0 + coarse: bool = False + subtitles_ocr: bool = True + verbose: bool = False + dump_memory: bool = False + keep_temporaries: bool = False diff --git a/src/tscut/matroska/codec.py b/src/tscut/matroska/codec.py index 10cf2ff..7925deb 100644 --- a/src/tscut/matroska/codec.py +++ b/src/tscut/matroska/codec.py @@ -14,12 +14,13 @@ from tscut.h264.avc import ( AVCDecoderConfiguration, get_avc_config_from_h264 ) +from tscut.exceptions import InvalidMediaError logger = logging.getLogger(__name__) @typechecked -def dump_codec_private_data(avc_decoder_configuration: AVCDecoderConfiguration) -> bytearray | None: +def dump_codec_private_data(avc_decoder_configuration: AVCDecoderConfiguration) -> bytearray: # Rebuild a Matroska Codec Private Element res = bytearray() # Code private element @@ -30,7 +31,7 @@ def dump_codec_private_data(avc_decoder_configuration: AVCDecoderConfiguration) embl_length = get_ebml_length(len(buf)) if embl_length is None: - return None + raise InvalidMediaError logger.debug('EMBL encoded length: %s', hexdump.dump(embl_length, sep=':')) res.extend(embl_length) res.extend(buf) @@ -38,7 +39,7 @@ def dump_codec_private_data(avc_decoder_configuration: AVCDecoderConfiguration) return res @typechecked -def get_codec_private_data_from_h264(input_file: IO[bytes]) -> bytearray | None: +def get_codec_private_data_from_h264(input_file: IO[bytes]) -> bytearray: avcconfig = get_avc_config_from_h264(input_file) res = dump_codec_private_data(avcconfig) diff --git a/src/tscut/pipeline.py b/src/tscut/pipeline.py new file mode 100644 index 0000000..04ec9bc --- /dev/null +++ b/src/tscut/pipeline.py @@ -0,0 +1,514 @@ +# SPDX-License-Identifier: GPL-2.0-or-later +# +# Copyright (C) 2026 Frédéric Tronel + +from datetime import datetime, timedelta +import os +from os import unlink +from shutil import copyfile, move +from sys import exit +import logging + +import hexdump + +from .config import ProcessingOptions +from .tscut import ( + SupportedFormat, + change_codec_private_data, + concatenate_h264_parts, + concatenate_h264_ts_parts, + do_ocr, + dump_codec_private_data, + extract_all_streams, + extract_mkv_part, + extract_srt, + extract_track_from_mkv, + ffmpeg_convert, + find_subtitles_tracks, + get_avc_config_from_h264, + get_codec_private_data_from_mkv, + get_format, + get_frame_rate, + get_movie_duration, + get_nearest_iframe, + get_streams, + get_tesseract_supported_lang, + get_ts_frame, + merge_mkvs, + parse_codec_private, + remove_video_tracks_from_mkv, + remux_srt_subtitles, +) + +logger = logging.getLogger(__name__) + +def process_recording(options: ProcessingOptions) -> None: + nb_parts = len(options.parts) + temporaries = [] + + basename = os.path.splitext(os.path.basename(options.input_file))[0] + mp4filename = basename+'.mp4' + mkvfilename = basename+'.mkv' + + try: + input_file = open(options.input_file, mode='rb') + logger.debug("Type of input file: %s", type(input_file)) + except OSError: + logger.error("Impossible to open %s", options.input_file) + exit(-1) + + format_of_file = get_format(options.tools_paths['ffprobe'], input_file) + + if format_of_file is None: + exit(-1) + + duration = timedelta(seconds=float(format_of_file['duration'])) + logger.info("Durée de l'enregistrement: %s", duration) + + if options.framerate is None: + framerate = get_frame_rate(options.tools_paths['ffprobe'], input_file) + if framerate is None: + logger.error('Impossible to estimate frame rate !') + exit(-1) + else: + framerate = options.framerate + + logger.info('Frame rate: %.1f fps', framerate) + + found = False + for f in SupportedFormat: + if 'format_name' in format_of_file: + if format_of_file['format_name'] == str(f): + found = True + format_of_file = f + break + + if not found: + logger.error('Unsupported format of file') + + if format_of_file == SupportedFormat.TS: + logger.info("Converting TS to MP4 (to fix timestamps).") + try: + with open(mp4filename, 'wb+') as mp4: + ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'], + input_file, 'mpegts', mp4, 'mp4', duration) + temporaries.append(mp4) + logger.info("Converting MP4 to MKV.") + try: + mkv = open(mkvfilename, 'wb+') + except OSError: + logger.error('') + + ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'], + mp4, 'mp4', mkv, 'matroska', duration) + if nb_parts > 0: + temporaries.append(mkv) + except OSError: + logger.error('') + + elif format_of_file == SupportedFormat.MP4: + logger.info("Converting MP4 to MKV") + try: + mkv = open(mkvfilename, 'wb+') + except OSError: + logger.error('') + ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'], + input_file, 'mp4', mkv, 'matroska', duration) + if nb_parts > 0: + temporaries.append(mkv) + else: + logger.info("Already in MKV") + mkv = input_file + + streams = get_streams(options.tools_paths['ffprobe'], mkv) + if streams is None: + logger.error("No streams found in file: %s", mkv) + exit(-1) + + logger.debug('Streams: %s', streams) + main_video = None + nb_videos = 0 + for stream in streams: + if stream['codec_type'] == 'video': + if stream['disposition']['default'] == 1: + main_video = stream + width = stream['width'] + height = stream['height'] + break + nb_videos+=1 + if nb_videos == 1: + main_video = stream + width = stream['width'] + height = stream['height'] + else: + main_video = None + + if main_video is None: + logger.error('Impossible to find main video stream.') + exit(-1) + + # We retrieve the main private codec data + _, main_codec_private_data = get_codec_private_data_from_mkv( + mkvinfo_path=options.tools_paths['mkvinfo'], + input_file=mkv + ) + logger.debug('Main video stream has following private data: %s', + hexdump.dump(main_codec_private_data, sep=':')) + + if main_codec_private_data is None: + logger.error("Impossible to retrieve private data from MKV file %s", mkv) + exit(-1) + + # We parse them + main_avc_config = parse_codec_private(main_codec_private_data) + logger.debug('AVC configuration: %s', main_avc_config) + + # We check if the parse and dump operations are idempotent. + private_data = dump_codec_private_data(main_avc_config) + logger.debug('Redump AVC configuration: %s', hexdump.dump(private_data, sep=':')) + # In rare occasion, the PPS has trailing zeroes that do not seem to be related to useful data + # but they differ from the private data we generate that do not contain them. + # In that case we try to redecode our own private data to see if both AVC configurations are + # the same. + if main_codec_private_data != private_data: + logger.warning('Difference detected in bitstream !!') + iso_avc_config = parse_codec_private(private_data) + logger.debug('Reread AVC configuration: %s', iso_avc_config) + + # If there exists a difference between our own reconstructed AVC configuration and the + # original one, we abandon + if iso_avc_config != main_avc_config: + logger.error('AVC configurations are different: %s\n%s\n', main_avc_config, + iso_avc_config) + exit(-1) + + # Pour chaque portion + partnum = 0 + mkvparts = [] + h264parts = [] + h264_ts = [] + checks = [] + pos = timedelta() + + other_avc_configs = [] + + for ts1, ts2 in options.parts: + # TODO: translate comment in english + # Trouver l'estampille de la trame 'I' la plus proche (mais postérieure) au début + # de la portion. + # Trouver l'estampille de la trame 'I' la plus proche (mais antérieure) à la fin + # de la portion. + # On a alors + # debut ----- trame --------- trame --------- fin fin+1 + # 'B/P' 'B/P'* 'I' 'I' 'B/P'* 'B/P' 'I/B/P' + # Si la trame de début est déjà 'I', il n'y a rien à faire. + # Sinon on extrait les trames 'B' ou 'P' depuis le début jusqu'à la trame 'I' non incluse. + # Si la trame de fin précède une trame I, on n'a rien à faire. + # Sinon on extrait toutes les trames depuis la dernière trame I jusqu'à la trame de fin. + + partnum = partnum + 1 + + # Get the nearest I-frame whose timestamp is greater or equal to the beginning. + head_frames = get_nearest_iframe(options.tools_paths['ffprobe'], mkv, ts1, before=False) + if head_frames is None: + logger.error('Impossible to retrieve I-frame') + exit(-1) + + # Get the nearest I-frame whose timestamp ... + # TODO: wrong here ... + tail_frames = get_nearest_iframe(options.tools_paths['ffprobe'], mkv, ts2, before=True) + if tail_frames is None: + logger.error('Impossible to retrieve I-frame') + exit(-1) + + nb_head_frames, head_iframe = head_frames + nb_tail_frames, tail_iframe = tail_frames + + logger.info("Found %d frames between beginning of current part and first I-frame", + nb_head_frames) + logger.info("Found %d frames between last I-frame and end of current part", + nb_tail_frames) + + head_iframe_ts = get_ts_frame(head_iframe) + if head_iframe_ts is None: + exit(-1) + tail_iframe_ts = get_ts_frame(tail_iframe) + if tail_iframe_ts is None: + exit(-1) + + checks.append(pos+head_iframe_ts-ts1) + + subparts = [] + + # TODO: separate pipeline processing between coarse and not fine grain options. + + # if args.coarse: + # do_coarse_processing(ffmpeg=paths['ffmpeg'], ffprobe=paths['ffprobe'], input_file=mkv, + # begin=ts1, end=head_iframe_ts, nb_frames=nb_head_frames-1, + # frameRate=frameRate, files_prefix='part-%d-head' % (partnum), + # streams=streams, width=width, height=height, + # temporaries=temporaries, dump_mem_fd=args.dump) + # else: + # doFineGrainProcessing(ffmpeg=paths['ffmpeg'], ffprobe=paths['ffprobe'], + # input_file=mkv, begin=ts1, end=head_iframe_ts, + # nb_frames=nb_head_frames-1, frameRate=frameRate, + # files_prefix='part-%d-head' % (partnum), streams=streams, + # width=width, height=height, temporaries=temporaries, + # dump_mem_fd=args.dump) + + if (not options.coarse) and (nb_head_frames > options.threshold): + # We extract all frames between the beginning upto the frame that immediately preceeds + # the I-frame. + h264_head, h264_head_ts, mkv_head = extract_all_streams( + ffmpeg_path=options.tools_paths['ffmpeg'], + ffprobe_path=options.tools_paths['ffprobe'], + input_file=mkv, begin=ts1, + end=head_iframe_ts, + nb_frames=nb_head_frames-1, + framerate=framerate, + files_prefix=f'part-{partnum:d}-head', + streams=streams, width=width, + height=height, + temporaries=temporaries, + dump_mem_fd=options.dump_memory) + + # If we are not at an exact boundary: + if mkv_head is not None: + subparts.append(mkv_head) + if h264_head is not None: + avcconfig = get_avc_config_from_h264(h264_head) + other_avc_configs.append(avcconfig) + h264parts.append(h264_head) + if h264_head_ts is not None: + h264_ts.append(h264_head_ts) + + # Creating MKV file that corresponds to current part between I-frames + # Internal video with all streams (video, audio and subtitles) + internal_mkv_name = f'part-{partnum:d}-internal.mkv' + # Internal video stream as a raw H264 stream + internal_h264_name = f'part-{partnum:d}-internal.h264' + # Internal video timestamps + internal_h264_ts_name = f'part-{partnum:d}-internal-ts.txt' + # Internal video with only audio and subtitles streams + internal_novideo_mkv_name = f'part-{partnum:d}-internal-novideo.mkv' + + try: + internal_mkv = open(internal_mkv_name, 'wb+') + except OSError: + logger.error('Impossible to create file: %s', internal_mkv_name) + exit(-1) + + try: + internal_novideo_mkv = open(internal_novideo_mkv_name, 'wb+') + except OSError: + logger.error('Impossible to create file: %s', internal_novideo_mkv_name) + exit(-1) + + try: + internal_h264 = open(internal_h264_name, 'wb+') + except OSError: + logger.error('Impossible to create file: %s', internal_h264_name) + exit(-1) + + try: + internal_h264_ts = open(internal_h264_ts_name, 'w+', encoding='utf8') + except OSError: + logger.error('Impossible to create file: %s', internal_h264_ts_name) + exit(-1) + + # logger.info('Merge header, middle and trailer subpart into: %s' % internal_mkv_name) + # Extract internal part of MKV + extract_mkv_part(mkvmerge_path=options.tools_paths['mkvmerge'], input_file=mkv, + output_file=internal_mkv, begin=head_iframe_ts, end=tail_iframe_ts) + + # Extract video stream of internal part as a raw H264 and its timestamps. + logger.info('Extract video track as raw H264 file.') + extract_track_from_mkv(mkvextract_path=options.tools_paths['mkvextract'], + input_file=internal_mkv, index=0, output_file=internal_h264, + timestamps=internal_h264_ts) + + # Remove video track from internal part of MKV + logger.info('Remove video track from %s', internal_mkv_name) + remove_video_tracks_from_mkv(mkvmerge_path=options.tools_paths['mkvmerge'], + input_file=internal_mkv, output_file=internal_novideo_mkv) + + temporaries.append(internal_mkv) + temporaries.append(internal_h264) + temporaries.append(internal_h264_ts) + temporaries.append(internal_novideo_mkv) + + h264parts.append(internal_h264) + h264_ts.append(internal_h264_ts) + subparts.append(internal_novideo_mkv) + + if (not options.coarse) and (nb_tail_frames > options.threshold): + # We extract all frames between the I-frame (including it) upto the end. + h264_tail, h264_tail_ts, mkv_tail = extract_all_streams(ffmpeg_path=options.tools_paths['ffmpeg'], + ffprobe_path=options.tools_paths['ffprobe'], + input_file=mkv, begin=tail_iframe_ts, + end=ts2, nb_frames=nb_tail_frames, + framerate=framerate, + files_prefix=f'part-{partnum:d}-tail', + streams=streams, + width=width, height=height, + temporaries=temporaries, + dump_mem_fd=options.dump_memory) + + if mkv_tail is not None: + subparts.append(mkv_tail) + if h264_tail is not None: + avcconfig = get_avc_config_from_h264(h264_tail) + other_avc_configs.append(avcconfig) + h264parts.append(h264_tail) + if h264_tail_ts is not None: + h264_ts.append(h264_tail_ts) + + logger.info('Merging MKV: %s', subparts) + + part = merge_mkvs(mkvmerge_path=options.tools_paths['mkvmerge'], inputs=subparts, + output_name=f'part-{partnum:d}.mkv', concatenate=True) + mkvparts.append(part) + temporaries.append(part) + + pos = pos+tail_iframe_ts-ts1 + + # We need to check the end also + checks.append(pos) + + # When using coarse option there is a single AVC configuration. + for avc_config in other_avc_configs: + main_avc_config.merge(avc_config) + logger.debug('Merged AVC configuration: %s', main_avc_config) + + nb_mkv_parts = len(mkvparts) + if nb_mkv_parts > 0: + try: + full_h264 = open(f'{basename}-full.h264', 'wb+') + except OSError: + logger.error('Impossible to create file full H264 stream.') + exit(-1) + + logger.info('Merging all H264 tracks') + concatenate_h264_parts(h264parts=h264parts, output=full_h264) + temporaries.append(full_h264) + + try: + full_h264_ts = open(f'{basename}-ts.txt', 'w+', encoding='utf8') + except OSError: + logger.error('Impossible to create file containing all video timestamps.') + exit(-1) + + logger.info('Merging H264 timestamps') + concatenate_h264_ts_parts(h264_ts_parts=h264_ts, output=full_h264_ts) + temporaries.append(full_h264_ts) + + final_novideo_name = f'{basename}-novideo.mkv' + final_with_video_name = f'{basename}-video.mkv' + + if nb_mkv_parts > 1: + logger.info('Merging all audio and subtitles parts: %s', mkvparts) + merge_mkvs(mkvmerge_path=options.tools_paths['mkvmerge'], inputs=mkvparts, + output_name=final_novideo_name, concatenate=True) + elif nb_mkv_parts == 1: + copyfile('part-1.mkv', final_novideo_name) + else: + logger.info("Nothing else to do.") + copyfile(mkvfilename, final_with_video_name) + + if nb_mkv_parts >=1 : + try: + final_novideo = open(final_novideo_name, 'rb') + except OSError: + logger.error('Impossible to open file: %s.', final_novideo_name) + exit(-1) + + temporaries.append(final_novideo) + + full_h264_ts.seek(0) + + logger.info('Merging final video track and all other tracks together') + final_with_video = merge_mkvs(mkvmerge_path=options.tools_paths['mkvmerge'], + inputs=[full_h264, final_novideo], + output_name=final_with_video_name, concatenate=False, + timestamps={0: full_h264_ts}) + final_codec_private_data = dump_codec_private_data(main_avc_config) + logger.debug('Final codec private data: %s', hexdump.dump(final_codec_private_data, + sep=':')) + logger.info('Changing codec private data with the new one.') + change_codec_private_data(options.tools_paths['mkvinfo'], final_with_video, + final_codec_private_data) + + if options.subtitles_ocr: + if not options.all_optional_tools: + logger.warning("Missing tools for extracting subtitles.") + move(final_with_video_name, options.output_file) + else: + # Final cut is not any more the final step. + temporaries.append(final_with_video) + duration = get_movie_duration(options.tools_paths['ffprobe'], final_with_video) + supported_langs = get_tesseract_supported_lang(options.tools_paths['tesseract']) + logger.info('Supported lang: %s', supported_langs) + logger.info('Find subtitles tracks and language.') + subtitles = find_subtitles_tracks(options.tools_paths['ffprobe'], final_with_video) + logger.info(subtitles) + sts = {} + for subtitle in subtitles: + index = subtitle['index'] + if 'tags' in subtitle: + if 'language' in subtitle['tags']: + lang = subtitle['tags']['language'] + if lang in sts: + sts[lang].append(index) + else: + sts[lang] = [index] + else: + logger.error("Dropping subtitle: %s because it is missing language\ + indication", subtitle) + else: + logger.error("Dropping subtitle: %s because it is missing language indication", + subtitle) + + logger.info(sts) + if len(sts) > 0: + logger.info('Supported languages: %s', supported_langs) + list_of_subtitles = extract_srt(options.tools_paths['mkvextract'], + final_with_video_name, sts, supported_langs) + logger.info(list_of_subtitles) + for idx_name, sub_name, _, _ in list_of_subtitles: + try: + idx = open(idx_name,'rb') + except OSError: + logger.error("Impossible to open %s.", idx_name) + exit(-1) + try: + sub = open(sub_name,'rb') + except OSError: + logger.error("Impossible to open %s.", sub_name) + exit(-1) + + temporaries.append(idx) + temporaries.append(sub) + + ocr = do_ocr(options.tools_paths['vobsubocr'], list_of_subtitles, duration, + temporaries, options.dump_memory) + logger.info(ocr) + + # Remux SRT subtitles + remux_srt_subtitles(options.tools_paths['mkvmerge'], final_with_video, + options.output_file, ocr) + else: + copyfile(final_with_video_name, options.output_file) + else: + move(final_with_video_name, options.output_file) + + if not options.keep_temporaries: + logger.info("Cleaning temporary files") + for f in temporaries: + path = os.path.realpath(f.name) + logger.info("Removing: %s", path) + f.close() + unlink(path) + + d = datetime(1,1,1) + for c in checks: + logger.info("Please check cut smoothness at %s", (c+d).strftime("%H:%M:%S")) diff --git a/src/tscut/tools/mkvtoolnix.py b/src/tscut/tools/mkvtoolnix.py index 277abf6..ab0be46 100644 --- a/src/tscut/tools/mkvtoolnix.py +++ b/src/tscut/tools/mkvtoolnix.py @@ -17,6 +17,7 @@ from os import ( ftruncate ) import os +from pathlib import Path from subprocess import PIPE, Popen from typing import IO @@ -25,6 +26,7 @@ import hexdump from tqdm import tqdm from tscut.matroska.ebml import change_ebml_element_size +from tscut.exceptions import InvalidMediaError # Found codec private data using mkvinfo @typechecked @@ -411,12 +413,12 @@ def remove_video_tracks_from_mkv(mkvmerge_path:str, input_file: IO[bytes], @typechecked -def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_filename: str, +def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_file: Path, subtitles) -> None: logger = logging.getLogger(__name__) try: - out = open(output_filename, 'w', encoding='utf8') + out = open(output_file, 'w', encoding='utf8') except OSError: logger.error('Impossible to create file: %s', output_filename) return None