From 97ef0e972cf035c0a6f09d82c6af9941e3169905 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Tronel?= Date: Fri, 28 Aug 2026 16:20:17 +0200 Subject: [PATCH] Separate main() function in a dedicated module. --- src/tscut/.gitignore | 1 + src/tscut/__main__.py | 10 + src/tscut/cli.py | 559 ++++++++++++++++++++++++++++++++++++++++++ src/tscut/tscut.py | 532 ---------------------------------------- 4 files changed, 570 insertions(+), 532 deletions(-) create mode 100644 src/tscut/.gitignore create mode 100644 src/tscut/__main__.py create mode 100644 src/tscut/cli.py diff --git a/src/tscut/.gitignore b/src/tscut/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/src/tscut/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/src/tscut/__main__.py b/src/tscut/__main__.py new file mode 100644 index 0000000..cf6cdff --- /dev/null +++ b/src/tscut/__main__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: GPL-2.0-or-later +# +# Copyright (C) 2026 Frédéric Tronel + +"""Entry point for ``python -m tscut``.""" + +from tscut.cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/tscut/cli.py b/src/tscut/cli.py new file mode 100644 index 0000000..fb7e715 --- /dev/null +++ b/src/tscut/cli.py @@ -0,0 +1,559 @@ +# SPDX-License-Identifier: GPL-2.0-or-later +# +# Copyright (C) 2026 Frédéric Tronel + +"""Command-line interface for tscut.""" + +import argparse +import logging +from os import unlink +import os.path +from sys import exit +from datetime import datetime,timedelta +from shutil import copyfile, move + +import coloredlogs +import hexdump + +from .tscut import check_required_tools, parse_time_interval, cmp_to_key, get_format,\ + SupportedFormat, compare_time_interval, ffmpeg_convert, parse_codec_private,\ + dump_codec_private_data, get_nearest_iframe, get_movie_duration, extract_all_streams,\ + extract_mkv_part, extract_srt, concatenate_h264_parts, concatenate_h264_ts_parts,\ + do_ocr, get_frame_rate, get_streams, get_ts_frame, remove_video_tracks_from_mkv,\ + merge_mkvs, extract_track_from_mkv, get_avc_config_from_h264,\ + get_codec_private_data_from_mkv, find_subtitles_tracks, change_codec_private_data,\ + get_tesseract_supported_lang, remux_srt_subtitles + +logger = logging.getLogger(__name__) + +def main() -> int: + """Run the tscut command-line interface.""" + + coloredlogs.install() + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input", dest='input_file', type=str, required=True, + help="Input file to process (can be .ts, .mp4 or .mkv).") + parser.add_argument("-o", "--output", dest='output_file', type=str, required=True, + help="Output MKV file to produce.") + parser.add_argument("-p", "--part", dest='parts', nargs='+', required=False, action='append', + metavar="hh:mm:ss[.mmm]-hh:mm:ss[.mmm]", + help="Extract this exact part of the original file.") + parser.add_argument("-k", "--keep", action='store_true', + help="Do not cleanup temporary files after processing.") + parser.add_argument("-t", "--threshold", action='store', type=int, + help="Suppress headers and trailers that are smaller than the threshold.") + parser.add_argument("-c", "--coarse", action='store_true', dest='coarse', + help="Do not take trailers and headers into account.") + parser.add_argument("--dump-memory", action='store_true', dest='dump', + help="For debug purpose, dump all memory mapping of headers (and trailers)\ + before (after) each part. They are kept in memory only otherwise.") + parser.add_argument("-s","--srt", action='store_true', dest='srt', + help="Dump subtitles and make OCR and finally remux them in the movie\ + (as SRT).") + parser.add_argument("-v","--verbose", action='store_true', dest='verbose', help="Debug.") + parser.add_argument("-f","--framerate", action='store', type=int, + help="Override frame rate estimator.") + + args = parser.parse_args() + logger.info('Arguments: %s', args) + + if args.verbose: + logger.info('Setting logging to debug mode') + coloredlogs.set_level(level=logging.DEBUG) + + logger.debug('Arguments: %s', args) + + if args.coarse and args.threshold is not None: + logger.error('--coarse and threshold arguments are exclusive.') + exit(-1) + + if (not args.coarse) and args.threshold is None: + args.threshold = 0 + + all_optional_tools, paths = check_required_tools() + + # Flatten args.parts + intervals = [] + if args.parts is not None: + for part in args.parts: + for subpart in part: + intervals.append(subpart) + + parts=[] + # Parse each interval + for interval in intervals: + ts1, ts2 = parse_time_interval(interval) + if ts1 is None or ts2 is None: + logger.error("Illegal time interval: %s", interval) + exit(-1) + parts.append((ts1,ts2)) + + # Sort intervals + parts.sort(key=cmp_to_key(compare_time_interval)) + + # Check that no intervals are overlapping + prevts = timedelta(0) + for part in parts: + ts1, ts2 = part + if prevts > ts1: + logger.error('Intervals are overlapping') + 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 IOError: + 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 IOError: + logger.error('') + + ffmpeg_convert(paths['ffmpeg'], paths['ffprobe'], mp4, 'mp4', mkv, 'matroska', + duration) + if nb_parts > 0: + temporaries.append(mkv) + except IOError: + logger.error('') + + elif format_of_file == SupportedFormat.MP4: + logger.info("Converting MP4 to MKV") + try: + mkv = open(mkvfilename, 'wb+') + except IOError: + 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 IOError: + logger.error('Impossible to create file: %s', internal_mkv_name) + exit(-1) + + try: + internal_novideo_mkv = open(internal_novideo_mkv_name, 'wb+') + except IOError: + logger.error('Impossible to create file: %s', internal_novideo_mkv_name) + exit(-1) + + try: + internal_h264 = open(internal_h264_name, 'wb+') + except IOError: + 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 IOError: + 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 IOError: + 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 IOError: + 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 IOError: + 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 IOError: + logger.error("Impossible to open %s.", idx_name) + exit(-1) + try: + sub = open(sub_name,'rb') + except IOError: + 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")) + + return 0 diff --git a/src/tscut/tscut.py b/src/tscut/tscut.py index a99c2e3..8aa2351 100755 --- a/src/tscut/tscut.py +++ b/src/tscut/tscut.py @@ -3295,535 +3295,3 @@ def do_coarse_processing(ffmpeg_path:str, ffprobe_path:str, mkvmerge_path:str, begin=begin, end=end) temporaries.append(internal_mkv) - -def main() -> None: - logger = logging.getLogger(__name__) - coloredlogs.install() - parser = argparse.ArgumentParser() - parser.add_argument("-i", "--input", dest='input_file', type=str, required=True, - help="Input file to process (can be .ts, .mp4 or .mkv).") - parser.add_argument("-o", "--output", dest='output_file', type=str, required=True, - help="Output MKV file to produce.") - parser.add_argument("-p", "--part", dest='parts', nargs='+', required=False, action='append', - metavar="hh:mm:ss[.mmm]-hh:mm:ss[.mmm]", - help="Extract this exact part of the original file.") - parser.add_argument("-k", "--keep", action='store_true', - help="Do not cleanup temporary files after processing.") - parser.add_argument("-t", "--threshold", action='store', type=int, - help="Suppress headers and trailers that are smaller than the threshold.") - parser.add_argument("-c", "--coarse", action='store_true', dest='coarse', - help="Do not take trailers and headers into account.") - parser.add_argument("--dump-memory", action='store_true', dest='dump', - help="For debug purpose, dump all memory mapping of headers (and trailers)\ - before (after) each part. They are kept in memory only otherwise.") - parser.add_argument("-s","--srt", action='store_true', dest='srt', - help="Dump subtitles and make OCR and finally remux them in the movie\ - (as SRT).") - parser.add_argument("-v","--verbose", action='store_true', dest='verbose', help="Debug.") - parser.add_argument("-f","--framerate", action='store', type=int, - help="Override frame rate estimator.") - - args = parser.parse_args() - logger.info('Arguments: %s', args) - - if args.verbose: - logger.info('Setting logging to debug mode') - coloredlogs.set_level(level=logging.DEBUG) - - logger.debug('Arguments: %s', args) - - if args.coarse and args.threshold is not None: - logger.error('--coarse and threshold arguments are exclusive.') - exit(-1) - - if (not args.coarse) and args.threshold is None: - args.threshold = 0 - - all_optional_tools, paths = check_required_tools() - - # Flatten args.parts - intervals = [] - if args.parts is not None: - for part in args.parts: - for subpart in part: - intervals.append(subpart) - - parts=[] - # Parse each interval - for interval in intervals: - ts1, ts2 = parse_time_interval(interval) - if ts1 is None or ts2 is None: - logger.error("Illegal time interval: %s", interval) - exit(-1) - parts.append((ts1,ts2)) - - # Sort intervals - parts.sort(key=cmp_to_key(compare_time_interval)) - - # Check that no intervals are overlapping - prevts = timedelta(0) - for part in parts: - ts1, ts2 = part - if prevts > ts1: - logger.error('Intervals are overlapping') - 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 IOError: - 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 IOError: - logger.error('') - - ffmpeg_convert(paths['ffmpeg'], paths['ffprobe'], mp4, 'mp4', mkv, 'matroska', - duration) - if nb_parts > 0: - temporaries.append(mkv) - except IOError: - logger.error('') - - elif format_of_file == SupportedFormat.MP4: - logger.info("Converting MP4 to MKV") - try: - mkv = open(mkvfilename, 'wb+') - except IOError: - 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 IOError: - logger.error('Impossible to create file: %s', internal_mkv_name) - exit(-1) - - try: - internal_novideo_mkv = open(internal_novideo_mkv_name, 'wb+') - except IOError: - logger.error('Impossible to create file: %s', internal_novideo_mkv_name) - exit(-1) - - try: - internal_h264 = open(internal_h264_name, 'wb+') - except IOError: - 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 IOError: - 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 IOError: - 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 IOError: - 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 IOError: - 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") - 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 IOError: - logger.error("Impossible to open %s.", idx_name) - exit(-1) - try: - sub = open(sub_name,'rb') - except IOError: - 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")) - -if __name__ == "__main__": - main()