366 lines
15 KiB
Python
366 lines
15 KiB
Python
# 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
|
|
)
|
|
|
|
|
|
|