# SPDX-License-Identifier: GPL-2.0-or-later # # Copyright (C) 2026 Frédéric Tronel import logging from datetime import timedelta from os import ( SEEK_SET, fstat, lseek, read, write, ) from shutil import copyfile from typing import IO, Any, BinaryIO, Sequence, TextIO import hexdump from tqdm import tqdm from typeguard import typechecked from tscut.exceptions import InvalidMediaError from tscut.h264.avc import get_avc_config_from_h264 from tscut.matroska.codec import dump_codec_private_data from tscut.models import CutResult, PreparedMedia, ProcessingOptions from tscut.tools.ffmpeg import extract_all_streams from tscut.tools.ffprobe import ( get_nearest_iframe, ) from tscut.tools.mkvtoolnix import ( change_codec_private_data, extract_mkv_part, extract_track_from_mkv, merge_mkvs, remove_video_tracks_from_mkv, ) from tscut.tools.timeframe import get_ts_frame logger = logging.getLogger(__name__) @typechecked def concatenate_h264_parts(h264parts: Sequence[BinaryIO], output: BinaryIO) -> None: logger = logging.getLogger(__name__) total_length = 0 for h264 in h264parts: fd = h264.fileno() total_length += fstat(fd).st_size logger.info('Total length: %d', total_length) outfd = output.fileno() lseek(outfd, 0, SEEK_SET) pb = tqdm(total=total_length, unit='bytes', desc='Concatenation') for h264 in h264parts: fd = h264.fileno() lseek(fd, 0, SEEK_SET) while True: buf = read(fd, 1000000) if buf is None or len(buf) == 0: break pos = 0 while pos < len(buf): nb_bytes = write(outfd, buf[pos:]) pb.update(nb_bytes) pos += nb_bytes def concatenate_h264_ts_parts(h264_ts_parts: Sequence[TextIO], output: TextIO) -> None: logger = logging.getLogger(__name__) header = '# timestamp format v2\n' output.write(header) last = 0. first = True for part in h264_ts_parts: if first: offset = last else: # TODO: take framerate into account offset = last + 40 logger.debug('Parsing file: %s. Offset=%d', part, offset) isheader = part.readline() if (not isheader) or (isheader != header): logger.error('Impossible to find a valid header: "%s"', isheader) exit(-1) while True: line = part.readline() if not line: break ts = offset + float(line) last = max(last,ts) output.write(f'{ts:f}\n') if first: first = False def cut_recording(media: PreparedMedia, options: ProcessingOptions, temporaries: list[IO[Any]]) -> CutResult: # Pour chaque portion partnum = 0 mkvparts = [] h264parts = [] h264_ts = [] checks = [] pos = timedelta() other_avc_configs = [] mkvfilename = media.basename+'.mkv' 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'], media.movie, ts1, before=False) if head_frames is None: raise InvalidMediaError("Impossible to retrieve first I-frame") # Get the nearest I-frame whose timestamp ... # TODO: wrong here ... tail_frames = get_nearest_iframe(options.tools_paths['ffprobe'], media.movie, ts2, before=True) if tail_frames is None: raise InvalidMediaError("Impossible to retrieve last I-frame") 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: raise InvalidMediaError("Impossible to retrieve timestamp of first i-frame.") tail_iframe_ts = get_ts_frame(tail_iframe) if tail_iframe_ts is None: raise InvalidMediaError("Impossible to retrieve timestamp of final i-frame.") 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=media.movie, begin=ts1, end=head_iframe_ts, nb_frames=nb_head_frames-1, framerate=media.framerate, files_prefix=f'part-{partnum:d}-head', streams=media.streams, width=media.width, height=media.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=media.movie, 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=media.movie, begin=tail_iframe_ts, end=ts2, nb_frames=nb_tail_frames, framerate=media.framerate, files_prefix=f'part-{partnum:d}-tail', streams=media.streams, width=media.width, height=media.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: media.avc_config.merge(avc_config) logger.debug('Merged AVC configuration: %s', media.avc_config) nb_mkv_parts = len(mkvparts) if nb_mkv_parts > 0: try: full_h264 = open(f'{media.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'{media.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'{media.basename}-novideo.mkv' final_with_video_name = f'{media.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(media.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) return CutResult( filename = final_with_video_name, movie = final_with_video, check_positions = checks )