Refactoring: extract cutting engine inside a dedicated module.

This commit is contained in:
Frédéric Tronel
2026-08-30 17:19:03 +02:00
parent 0345293664
commit ae1a040590
6 changed files with 406 additions and 303 deletions
+311
View File
@@ -0,0 +1,311 @@
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Copyright (C) 2026 Frédéric Tronel
import logging
from typing import IO, Any
from datetime import timedelta
from shutil import copyfile, move
import hexdump
from tscut.exceptions import InvalidMediaError
from tscut.models import PreparedMedia, ProcessingOptions, CutResult
from tscut.tools.ffprobe import (
find_subtitles_tracks,
get_format,
get_frame_rate,
get_movie_duration,
get_nearest_iframe,
get_streams,
)
from tscut.tools.ffmpeg import extract_all_streams, ffmpeg_convert
from tscut.h264.avc import get_avc_config_from_h264, parse_codec_private
from tscut.tools.mkvtoolnix import (
change_codec_private_data,
extract_mkv_part,
extract_track_from_mkv,
get_codec_private_data_from_mkv,
merge_mkvs,
remove_video_tracks_from_mkv,
remux_srt_subtitles,
)
from tscut.tools.timeframe import get_ts_frame
from tscut.tscut import (
concatenate_h264_parts,
concatenate_h264_ts_parts
)
from tscut.matroska.codec import dump_codec_private_data
logger = logging.getLogger(__name__)
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
)