diff --git a/src/tscut/tools/mkvtoolnix.py b/src/tscut/tools/mkvtoolnix.py new file mode 100644 index 0000000..f81eb20 --- /dev/null +++ b/src/tscut/tools/mkvtoolnix.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: GPL-2.0-or-later +# +# Copyright (C) 2026 Frédéric Tronel + + +import logging +from io import TextIOWrapper +import re +from datetime import timedelta +from os import ( + read, + SEEK_SET, + lseek, + write, + fstat, + set_inheritable, + ftruncate +) +import os +from subprocess import PIPE, Popen +from typing import IO + +from typeguard import typechecked +import hexdump +from tqdm import tqdm + +from tscut.matroska.ebml import change_ebml_element_size + +# Found codec private data using mkvinfo +@typechecked +def get_codec_private_data_from_mkv(mkvinfo_path:str, + input_file: IO[bytes]) -> tuple[int, bytes]|tuple[None,None]: + logger = logging.getLogger(__name__) + + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + set_inheritable(infd, True) + found = False + env = {**os.environ, 'LANG': 'C'} + # Output example + # Codec's private data: size 48 (H.264 profile: High @L4.0) hexdump 01 64 00 28 ff e1 00 1b 67\ + # 64 00 28 ac d9 40 78 04 4f dc d4 04 04 05 00 00 92 ef 00 1d ad a6 1f 16 2d 96 01 00 06 68 fb\ + # a3 cb 22 c0 fd f8 f8 00 at 406 size 51 data size 48 + + with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE, + close_fds=False, env=env) as mkvinfo: + out, _ = mkvinfo.communicate() + lines = out.decode('utf8') + reg_exp = (r"^.*Codec's private data: size ([0-9]+) \(H.264.*\) hexdump " + r"(?P([0-9a-f]{2} )+)at (?P[0-9]+) size (?P[0-9]+).*$") + p = re.compile(reg_exp) + for line in lines.splitlines(): + m = p.match(line) + if m is not None: + size = int(m.group('size')) + position = int(m.group('position')) + logger.debug("Found codec private data at position: %s, size: %d", position, size) + found = True + mkvinfo.wait() + break + + if found: + lseek(infd, position, SEEK_SET) + data = read(infd, size) + return position, data + + logger.error("Impossible to retrieve codec private data from MKV !") + return None, None + +@typechecked +def parse_mkv_tree(mkvinfo_path:str, input_file: IO[bytes]) -> dict[str,tuple[int,int]]: + logger = logging.getLogger(__name__) + + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + set_inheritable(infd, True) + env = {**os.environ, 'LANG': 'C'} + elements = {} + + with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE, + close_fds=False, env=env) as mkvinfo: + out, _ = mkvinfo.communicate() + lines = out.decode('utf8') + prefix = [] + reg_exp = (r"(^(?P\+)|(\|(?P[ ]*\+))).*at (?P[0-9]+)" + r" size (?P[0-9]+).*$") + p = re.compile(reg_exp) + prev_depth = -1 + for line in lines.splitlines(): + m = p.match(line) + if m is None: + logger.error("Impossible to match line: %s", line) + else: + position = int(m.group('position')) + size = int(m.group('size')) + root = m.group('root') is not None + if root: + depth = 0 + else: + depth = len(m.group('depth')) + + if depth > prev_depth: + for _ in range(depth-prev_depth): + prefix.append(1) + elif depth == prev_depth: + subid = prefix[-1] + subid+=1 + prefix.pop() + prefix.append(subid) + else: + for _ in range(prev_depth-depth): + prefix.pop() + subid = prefix[-1] + subid+=1 + prefix.pop() + prefix.append(subid) + + prev_depth = depth + key=".".join(map(str, prefix)) + elements[key] = (position, size) + + mkvinfo.wait() + return elements + +@typechecked +def change_codec_private_data(mkvinfo_path:str, input_file: IO[bytes], codec_data:bytes) -> None: + logger = logging.getLogger(__name__) + + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + + current_length = fstat(infd).st_size + logger.info('Current size of file: %d', current_length) + position, current_data = get_codec_private_data_from_mkv(mkvinfo_path, input_file) + if position is None or current_data is None: + return None + current_data_length = len(current_data) + future_length = current_length - current_data_length + len(codec_data) + logger.info('Expected size of file: %d', future_length) + + logger.info('Current data at position %d: %s', position, hexdump.dump(current_data, sep=":")) + logger.info('Future data: %s', hexdump.dump(codec_data, sep=":")) + + elements = parse_mkv_tree(mkvinfo_path, input_file) + + found = False + for key, (pos,size) in elements.items(): + if pos == position: + logger.info('Codec private data key: %s', key) + found = True + break + + if not found: + logger.error('Impossible to retrieve the key of codec private data') + exit(-1) + + if current_length < future_length: + lseek(infd, position+current_data_length, SEEK_SET) + tail = read(infd, current_length-(position+current_data_length)) + # We extend the file at the end with zeroes + ftruncate(infd, future_length) + lseek(infd, position+len(codec_data), SEEK_SET) + write(infd, tail) + lseek(infd, position, SEEK_SET) + write(infd, codec_data) + elif current_length == future_length: + # Almost nothing to do except overwriting old private codec data with new ones. + lseek(infd, position, SEEK_SET) + write(infd, codec_data) + else: + lseek(infd, position+current_data_length, SEEK_SET) + tail = read(infd, current_length-(position+current_data_length)) + lseek(infd, position+len(codec_data), SEEK_SET) + write(infd, tail) + lseek(infd, position, SEEK_SET) + write(infd, codec_data) + # We reduce the length of file. + ftruncate(infd, future_length) + + # We have to modify the tree elements up to the root that contains the codec private data. + keys = key.split('.') + logger.info(keys) + + delta = future_length-current_length + # if there is no modification of the private codec data, no need to change anything. + if delta != 0: + for _ in range(len(keys)-1): + keys.pop() + key=".".join(map(str, keys)) + pos, size = elements[key] + logger.info('Trying to fix element with key: %s at position: %d with actual size: %d.', + key, pos, size) + # Changing an element can increase its size (in very rare case). + # In that case, we update the new delta that will be larger (because the element has + # been resized). + delta+=change_ebml_element_size(input_file, pos, delta) + +@typechecked +def extract_mkv_part(mkvmerge_path:str, input_file:IO[bytes], output_file:IO[bytes], + begin:timedelta, end:timedelta) -> None: + logger = logging.getLogger(__name__) + + logger.info('Extract video between I-frames at %s and %s', begin,end) + infd = input_file.fileno() + outfd = output_file.fileno() + lseek(infd, 0, SEEK_SET) + lseek(outfd, 0, SEEK_SET) + set_inheritable(infd, True) + set_inheritable(outfd, True) + env = {**os.environ, 'LANG': 'C'} + warnings = [] + + command = [mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '--split', f'parts:{begin}-{end}', + f'/proc/self/fd/{infd:d}'] + logger.debug('Executing: %s', command) + + with Popen(command, stdout=PIPE, close_fds=False, env=env) as mkvmerge: + pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', + desc='Extraction') + for line in pb: + if line.startswith('Progress :'): + p = re.compile('^Progress : (?P[0-9]{1,3})%$') + m = p.match(line) + if m is None: + logger.error('Impossible to parse progress') + pb.update(int(m['progress'])-pb.n) + elif line.startswith('Warning'): + warnings.append(line) + pb.update(100-pb.n) + pb.refresh() + pb.close() + + status = mkvmerge.wait() + if status == 1: + logger.warning('Extraction returns warning') + for w in warnings: + logger.warning(w) + elif status == 2: + logger.error('Extraction returns errors') + + +# Merge a list of mkv files passed as input, and produce a new MKV as output +@typechecked +def merge_mkvs(mkvmerge_path:str, inputs: list[IO[bytes]], output_name:str, + concatenate: bool=True, + timestamps: dict[int, IO[str]] | None = None) -> IO[bytes]|None: + logger = logging.getLogger(__name__) + + if timestamps is None: + timestamps = {} + + fds = [] + try: + out = open(output_name, 'wb+') + except OSError: + logger.error('Impossible to create file: %s', output_name) + return None + + outfd = out.fileno() + lseek(outfd, 0, SEEK_SET) + fds.append(outfd) + set_inheritable(outfd, True) + + # Timestamps of merged tracks are modified by the length of the preceding track. + # The default mode ('file') is using the largest timestamp of the whole file which may create + # desynchronize video and sound. + merge_params = [mkvmerge_path, '--append-mode', 'track'] + + first = True + partnum = 0 + for mkv in inputs: + if mkv is not None: + fd = mkv.fileno() + fds.append(fd) + set_inheritable(fd, True) + # If we pass a timestamps file associated with the considered track, use it. + if partnum in timestamps: + tsfd = timestamps[partnum].fileno() + lseek(tsfd, 0, SEEK_SET) + fds.append(tsfd) + set_inheritable(tsfd, True) + merge_params.extend(['--timestamps', f'{partnum:d}:/proc/self/fd/{tsfd:d}']) + if first: + merge_params.append(f'/proc/self/fd/{fd:d}') + first = False + elif concatenate: + merge_params.append(f'+/proc/self/fd/{fd:d}') + else: + merge_params.append(f'/proc/self/fd/{fd:d}') + partnum+=1 + + merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}']) + + # We merge all files. + warnings = [] + env = {**os.environ, 'LANG': 'C'} + logger.debug('Executing: LANG=C %s', merge_params) + + with Popen(merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge: + pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', + desc='Merging') + for line in pb: + if line.startswith('Progress :'): + p = re.compile('^Progress : (?P[0-9]{1,3})%$') + m = p.match(line) + if m is None: + logger.error('Impossible to parse progress') + pb.n = int(m['progress']) + pb.update() + elif line.startswith('Warning'): + warnings.append(line) + + status = mkvmerge.wait() + if status == 1: + logger.warning('Extraction returns warning') + for w in warnings: + logger.warning(w) + elif status == 2: + logger.error('Extraction returns errors') + + for fd in fds: + set_inheritable(fd, False) + + return out + + + +@typechecked +def extract_track_from_mkv(mkvextract_path: str, input_file: IO[bytes], index, + output_file: IO[bytes], timestamps) -> None: + logger = logging.getLogger(__name__) + + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + set_inheritable(infd, True) + + outfd = output_file.fileno() + lseek(outfd, 0, SEEK_SET) + set_inheritable(outfd, True) + + tsfd = timestamps.fileno() + lseek(tsfd, 0, SEEK_SET) + set_inheritable(tsfd, True) + + params = [ mkvextract_path, f'/proc/self/fd/{infd:d}', 'tracks', + f'{index:d}:/proc/self/fd/{outfd:d}', 'timestamps_v2', + f'{index:d}:/proc/self/fd/{tsfd:d}'] + + env = {**os.environ, 'LANG': 'C'} + logger.debug('Executing: LANG=C %s', params) + + with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract: + pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%', + desc='Extraction of track') + for line in pb: + if line.startswith('Progress :'): + p = re.compile('^Progress : (?P[0-9]{1,3})%$') + m = p.match(line) + if m is None: + logger.error('Impossible to parse progress') + pb.update(int(m['progress'])-pb.n) + pb.update(100-pb.n) + pb.refresh() + pb.close() + + extract.wait() + + if extract.returncode != 0: + logger.error('Mkvextract returns an error code: %d', extract.returncode) + else: + logger.info('Track %d was succesfully extracted.', index) + + + +@typechecked +def remove_video_tracks_from_mkv(mkvmerge_path:str, input_file: IO[bytes], + output_file: IO[bytes]) -> None: + logger = logging.getLogger(__name__) + + outfd = output_file.fileno() + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + lseek(outfd, 0, SEEK_SET) + set_inheritable(infd, True) + set_inheritable(outfd, True) + + params = [ mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '-D', f'/proc/self/fd/{infd:d}'] + logger.debug('Executing: LANG=C %s', params) + + env = {**os.environ, 'LANG': 'C'} + with Popen(params, stdout=PIPE, close_fds=False, env=env) as remove: + pb = tqdm(TextIOWrapper(remove.stdout, encoding="utf-8"), total=100, unit='%', + desc='Removal of video track:') + for line in pb: + if line.startswith('Progress :'): + p = re.compile('^Progress : (?P[0-9]{1,3})%$') + m = p.match(line) + if m is None: + logger.error('Impossible to parse progress') + pb.update(int(m['progress'])-pb.n) + pb.update(100-pb.n) + pb.refresh() + pb.close() + + remove.wait() + + if remove.returncode != 0: + logger.error('Mkvmerge returns an error code: %d', remove.returncode) + else: + logger.info('Video tracks were succesfully extracted.') + + +@typechecked +def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_filename: str, + subtitles) -> None: + logger = logging.getLogger(__name__) + + try: + out = open(output_filename, 'w', encoding='utf8') + except OSError: + logger.error('Impossible to create file: %s', output_filename) + return None + + outfd = out.fileno() + infd = input_file.fileno() + lseek(infd, 0, SEEK_SET) + set_inheritable(infd, True) + set_inheritable(outfd, True) + + mkv_merge_params = [mkvmerge_path, f'/proc/self/fd/{infd:d}'] + for fd, lang in subtitles: + lseek(fd, 0, SEEK_SET) + set_inheritable(fd, True) + mkv_merge_params.extend(['--language', f'0:{lang}', f'/proc/self/fd/{fd:d}']) + + mkv_merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}']) + + warnings = [] + env = {**os.environ, 'LANG': 'C'} + logger.info('Remux subtitles: %s', mkv_merge_params) + with Popen(mkv_merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge: + pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', + desc='Remux subtitles:') + for line in pb: + if line.startswith('Progress :'): + p = re.compile('^Progress : (?P[0-9]{1,3})%$') + m = p.match(line) + if m is None: + logger.error('Impossible to parse progress') + pb.n = int(m['progress']) + pb.update() + elif line.startswith('Warning'): + warnings.append(line) + + status = mkvmerge.wait() + if status == 1: + logger.warning('Remux subtitles returns warning') + for w in warnings: + logger.warning(w) + elif status == 2: + logger.error('Remux subtitles returns errors') + + return None diff --git a/src/tscut/tscut.py b/src/tscut/tscut.py index 9b4a5a0..8805908 100755 --- a/src/tscut/tscut.py +++ b/src/tscut/tscut.py @@ -2,36 +2,24 @@ '''A module to remove parts of video (.e.g advertisements) with single frame precision.''' # Standard modules -import json import logging -import os.path -import re -from datetime import timedelta from enum import IntEnum, unique -from io import BytesIO, TextIOWrapper -from math import ceil, floor, log from os import ( SEEK_SET, - close, fstat, - ftruncate, lseek, - memfd_create, read, - set_inheritable, write, ) -from subprocess import PIPE, Popen from sys import exit from typing import IO # Third party libraries -import hexdump -from iso639 import Lang -from iso639.exceptions import InvalidLanguageValue from tqdm import tqdm from typeguard import typechecked +from tscut.tools.mkvtoolnix import change_codec_private_data + from tscut.h264.avc import ( get_avc_config_from_h264, parse_codec_private @@ -54,16 +42,22 @@ from tscut.tools.ffmpeg import ( ffmpeg_convert ) from tscut.tools.timeframe import ( - compare_time_interval, - parse_time_interval + compare_time_interval, + parse_time_interval ) from tscut.tools.discovery import check_required_tools from tscut.ocr.subtitles import ( - do_ocr, - extract_srt, - get_tesseract_supported_lang + do_ocr, + extract_srt,get_tesseract_supported_lang +) +from tscut.tools.mkvtoolnix import ( + extract_mkv_part, + extract_track_from_mkv, + get_codec_private_data_from_mkv, + merge_mkvs, + remove_video_tracks_from_mkv, + remux_srt_subtitles ) - # Useful SPS/PPS discussion. # https://copyprogramming.com/howto/including-sps-and-pps-in-a-raw-h264-track @@ -116,441 +110,8 @@ class SupportedFormat(IntEnum): # ffmpeg -i -c:v copy -an -sn -bsf:v trace_headers -t 0.01\ # -report -loglevel 0 -f null - -# Found codec private data using mkvinfo -@typechecked -def get_codec_private_data_from_mkv(mkvinfo_path:str, - input_file: IO[bytes]) -> tuple[int, bytes]|tuple[None,None]: - logger = logging.getLogger(__name__) - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - set_inheritable(infd, True) - found = False - env = {**os.environ, 'LANG': 'C'} - # Output example - # Codec's private data: size 48 (H.264 profile: High @L4.0) hexdump 01 64 00 28 ff e1 00 1b 67\ - # 64 00 28 ac d9 40 78 04 4f dc d4 04 04 05 00 00 92 ef 00 1d ad a6 1f 16 2d 96 01 00 06 68 fb\ - # a3 cb 22 c0 fd f8 f8 00 at 406 size 51 data size 48 - with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE, - close_fds=False, env=env) as mkvinfo: - out, _ = mkvinfo.communicate() - out = out.decode('utf8') - reg_exp = (r"^.*Codec's private data: size ([0-9]+) \(H.264.*\) hexdump " - r"(?P([0-9a-f]{2} )+)at (?P[0-9]+) size (?P[0-9]+).*$") - p = re.compile(reg_exp) - for line in out.splitlines(): - m = p.match(line) - if m is not None: - size = int(m.group('size')) - position = int(m.group('position')) - logger.debug("Found codec private data at position: %s, size: %d", position, size) - found = True - mkvinfo.wait() - break - - if found: - lseek(infd, position, SEEK_SET) - data = read(infd, size) - return position, data - - return None, None - -@typechecked -def parse_mkv_tree(mkvinfo_path:str, input_file: IO[bytes]) -> dict[str,tuple[int,int]]: - logger = logging.getLogger(__name__) - - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - set_inheritable(infd, True) - env = {**os.environ, 'LANG': 'C'} - elements = {} - - with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE, - close_fds=False, env=env) as mkvinfo: - out, _ = mkvinfo.communicate() - out = out.decode('utf8') - prefix = [] - reg_exp = (r"(^(?P\+)|(\|(?P[ ]*\+))).*at (?P[0-9]+)" - r" size (?P[0-9]+).*$") - p = re.compile(reg_exp) - prev_depth = -1 - for line in out.splitlines(): - m = p.match(line) - if m is None: - logger.error("Impossible to match line: %s", line) - else: - position = int(m.group('position')) - size = int(m.group('size')) - root = m.group('root') is not None - if root: - depth = 0 - else: - depth = len(m.group('depth')) - - if depth > prev_depth: - for _ in range(depth-prev_depth): - prefix.append(1) - elif depth == prev_depth: - subid = prefix[-1] - subid+=1 - prefix.pop() - prefix.append(subid) - else: - for _ in range(prev_depth-depth): - prefix.pop() - subid = prefix[-1] - subid+=1 - prefix.pop() - prefix.append(subid) - - prev_depth = depth - key=".".join(map(str, prefix)) - elements[key] = (position, size) - - mkvinfo.wait() - return elements - - - -@typechecked -def change_codec_private_data(mkvinfo_path:str, input_file: IO[bytes], codec_data:bytes) -> None: - logger = logging.getLogger(__name__) - - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - - current_length = fstat(infd).st_size - logger.info('Current size of file: %d', current_length) - position, current_data = get_codec_private_data_from_mkv(mkvinfo_path, input_file) - current_data_length = len(current_data) - future_length = current_length - current_data_length + len(codec_data) - logger.info('Expected size of file: %d', future_length) - - logger.info('Current data at position %d: %s', position, hexdump.dump(current_data, sep=":")) - logger.info('Future data: %s', hexdump.dump(codec_data, sep=":")) - - elements = parse_mkv_tree(mkvinfo_path, input_file) - - found = False - for key, (pos,size) in elements.items(): - if pos == position: - logger.info('Codec private data key: %s', key) - found = True - break - - if not found: - logger.error('Impossible to retrieve the key of codec private data') - exit(-1) - - if current_length < future_length: - lseek(infd, position+current_data_length, SEEK_SET) - tail = read(infd, current_length-(position+current_data_length)) - # We extend the file at the end with zeroes - ftruncate(infd, future_length) - lseek(infd, position+len(codec_data), SEEK_SET) - write(infd, tail) - lseek(infd, position, SEEK_SET) - write(infd, codec_data) - elif current_length == future_length: - # Almost nothing to do except overwriting old private codec data with new ones. - lseek(infd, position, SEEK_SET) - write(infd, codec_data) - else: - lseek(infd, position+current_data_length, SEEK_SET) - tail = read(infd, current_length-(position+current_data_length)) - lseek(infd, position+len(codec_data), SEEK_SET) - write(infd, tail) - lseek(infd, position, SEEK_SET) - write(infd, codec_data) - # We reduce the length of file. - ftruncate(infd, future_length) - - # We have to modify the tree elements up to the root that contains the codec private data. - keys = key.split('.') - logger.info(keys) - - delta = future_length-current_length - # if there is no modification of the private codec data, no need to change anything. - if delta != 0: - for _ in range(len(keys)-1): - keys.pop() - key=".".join(map(str, keys)) - pos, size = elements[key] - logger.info('Trying to fix element with key: %s at position: %d with actual size: %d.', - key, pos, size) - # Changing an element can increase its size (in very rare case). - # In that case, we update the new delta that will be larger (because the element has - # been resized). - delta+=change_ebml_element_size(input_file, pos, delta) - - - - - -@typechecked -def extract_mkv_part(mkvmerge_path:str, input_file:IO[bytes], output_file:IO[bytes], - begin:timedelta, end:timedelta) -> None: - logger = logging.getLogger(__name__) - - logger.info('Extract video between I-frames at %s and %s', begin,end) - infd = input_file.fileno() - outfd = output_file.fileno() - lseek(infd, 0, SEEK_SET) - lseek(outfd, 0, SEEK_SET) - set_inheritable(infd, True) - set_inheritable(outfd, True) - env = {**os.environ, 'LANG': 'C'} - warnings = [] - - command = [mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '--split', f'parts:{begin}-{end}', - f'/proc/self/fd/{infd:d}'] - logger.debug('Executing: %s', command) - - with Popen(command, stdout=PIPE, close_fds=False, env=env) as mkvmerge: - pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', - desc='Extraction') - for line in pb: - if line.startswith('Progress :'): - p = re.compile('^Progress : (?P[0-9]{1,3})%$') - m = p.match(line) - if m is None: - logger.error('Impossible to parse progress') - pb.update(int(m['progress'])-pb.n) - elif line.startswith('Warning'): - warnings.append(line) - pb.update(100-pb.n) - pb.refresh() - pb.close() - - status = mkvmerge.wait() - if status == 1: - logger.warning('Extraction returns warning') - for w in warnings: - logger.warning(w) - elif status == 2: - logger.error('Extraction returns errors') - - - -# Merge a list of mkv files passed as input, and produce a new MKV as output -@typechecked -def merge_mkvs(mkvmerge_path:str, inputs: list[IO[bytes]], output_name:str, - concatenate: bool=True, timestamps: dict[int, IO[str]] | None = None) -> IO[bytes]|None: - logger = logging.getLogger(__name__) - - if timestamps is None: - timestamps = {} - - fds = [] - try: - out = open(output_name, 'wb+') - except OSError: - logger.error('Impossible to create file: %s', output_name) - return None - - outfd = out.fileno() - lseek(outfd, 0, SEEK_SET) - fds.append(outfd) - set_inheritable(outfd, True) - - # Timestamps of merged tracks are modified by the length of the preceding track. - # The default mode ('file') is using the largest timestamp of the whole file which may create - # desynchronize video and sound. - merge_params = [mkvmerge_path, '--append-mode', 'track'] - - first = True - partnum = 0 - for mkv in inputs: - if mkv is not None: - fd = mkv.fileno() - fds.append(fd) - set_inheritable(fd, True) - # If we pass a timestamps file associated with the considered track, use it. - if partnum in timestamps: - tsfd = timestamps[partnum].fileno() - lseek(tsfd, 0, SEEK_SET) - fds.append(tsfd) - set_inheritable(tsfd, True) - merge_params.extend(['--timestamps', f'{partnum:d}:/proc/self/fd/{tsfd:d}']) - if first: - merge_params.append(f'/proc/self/fd/{fd:d}') - first = False - elif concatenate: - merge_params.append(f'+/proc/self/fd/{fd:d}') - else: - merge_params.append(f'/proc/self/fd/{fd:d}') - partnum+=1 - - merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}']) - - # We merge all files. - warnings = [] - env = {**os.environ, 'LANG': 'C'} - logger.debug('Executing: LANG=C %s', merge_params) - - with Popen(merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge: - pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', - desc='Merging') - for line in pb: - if line.startswith('Progress :'): - p = re.compile('^Progress : (?P[0-9]{1,3})%$') - m = p.match(line) - if m is None: - logger.error('Impossible to parse progress') - pb.n = int(m['progress']) - pb.update() - elif line.startswith('Warning'): - warnings.append(line) - - status = mkvmerge.wait() - if status == 1: - logger.warning('Extraction returns warning') - for w in warnings: - logger.warning(w) - elif status == 2: - logger.error('Extraction returns errors') - - for fd in fds: - set_inheritable(fd, False) - - return out - - - -@typechecked -def extract_track_from_mkv(mkvextract_path: str, input_file: IO[bytes], index, - output_file: IO[bytes], timestamps) -> None: - logger = logging.getLogger(__name__) - - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - set_inheritable(infd, True) - - outfd = output_file.fileno() - lseek(outfd, 0, SEEK_SET) - set_inheritable(outfd, True) - - tsfd = timestamps.fileno() - lseek(tsfd, 0, SEEK_SET) - set_inheritable(tsfd, True) - - params = [ mkvextract_path, f'/proc/self/fd/{infd:d}', 'tracks', - f'{index:d}:/proc/self/fd/{outfd:d}', 'timestamps_v2', - f'{index:d}:/proc/self/fd/{tsfd:d}'] - - env = {**os.environ, 'LANG': 'C'} - logger.debug('Executing: LANG=C %s', params) - - with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract: - pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%', - desc='Extraction of track') - for line in pb: - if line.startswith('Progress :'): - p = re.compile('^Progress : (?P[0-9]{1,3})%$') - m = p.match(line) - if m is None: - logger.error('Impossible to parse progress') - pb.update(int(m['progress'])-pb.n) - pb.update(100-pb.n) - pb.refresh() - pb.close() - - extract.wait() - - if extract.returncode != 0: - logger.error('Mkvextract returns an error code: %d', extract.returncode) - else: - logger.info('Track %d was succesfully extracted.', index) - -@typechecked -def remove_video_tracks_from_mkv(mkvmerge_path:str, input_file: IO[bytes], - output_file: IO[bytes]) -> None: - logger = logging.getLogger(__name__) - - outfd = output_file.fileno() - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - lseek(outfd, 0, SEEK_SET) - set_inheritable(infd, True) - set_inheritable(outfd, True) - - params = [ mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '-D', f'/proc/self/fd/{infd:d}'] - logger.debug('Executing: LANG=C %s', params) - - env = {**os.environ, 'LANG': 'C'} - with Popen(params, stdout=PIPE, close_fds=False, env=env) as remove: - pb = tqdm(TextIOWrapper(remove.stdout, encoding="utf-8"), total=100, unit='%', - desc='Removal of video track:') - for line in pb: - if line.startswith('Progress :'): - p = re.compile('^Progress : (?P[0-9]{1,3})%$') - m = p.match(line) - if m is None: - logger.error('Impossible to parse progress') - pb.update(int(m['progress'])-pb.n) - pb.update(100-pb.n) - pb.refresh() - pb.close() - - remove.wait() - - if remove.returncode != 0: - logger.error('Mkvmerge returns an error code: %d', remove.returncode) - else: - logger.info('Video tracks were succesfully extracted.') - -@typechecked -def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_filename: str, - subtitles) -> None: - logger = logging.getLogger(__name__) - - try: - out = open(output_filename, 'w', encoding='utf8') - except OSError: - logger.error('Impossible to create file: %s', output_filename) - return None - - outfd = out.fileno() - infd = input_file.fileno() - lseek(infd, 0, SEEK_SET) - set_inheritable(infd, True) - set_inheritable(outfd, True) - - mkv_merge_params = [mkvmerge_path, f'/proc/self/fd/{infd:d}'] - for fd, lang in subtitles: - lseek(fd, 0, SEEK_SET) - set_inheritable(fd, True) - mkv_merge_params.extend(['--language', f'0:{lang}', f'/proc/self/fd/{fd:d}']) - - mkv_merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}']) - - warnings = [] - env = {**os.environ, 'LANG': 'C'} - logger.info('Remux subtitles: %s', mkv_merge_params) - with Popen(mkv_merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge: - pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', - desc='Remux subtitles:') - for line in pb: - if line.startswith('Progress :'): - p = re.compile('^Progress : (?P[0-9]{1,3})%$') - m = p.match(line) - if m is None: - logger.error('Impossible to parse progress') - pb.n = int(m['progress']) - pb.update() - elif line.startswith('Warning'): - warnings.append(line) - - status = mkvmerge.wait() - if status == 1: - logger.warning('Remux subtitles returns warning') - for w in warnings: - logger.warning(w) - elif status == 2: - logger.error('Remux subtitles returns errors') - - return None @typechecked def concatenate_h264_parts(h264parts: list[IO[bytes]], output: IO[bytes]) -> None: @@ -582,7 +143,7 @@ def concatenate_h264_parts(h264parts: list[IO[bytes]], output: IO[bytes]) -> Non def concatenate_h264_ts_parts(h264_ts_parts: list[IO[bytes]], output: IO[bytes]) -> None: logger = logging.getLogger(__name__) - header = '# timestamp format v2\n' + header = '# timestamp format v2\n'.encode('ascii') output.write(header) @@ -605,7 +166,7 @@ def concatenate_h264_ts_parts(h264_ts_parts: list[IO[bytes]], output: IO[bytes]) break ts = offset + float(line) last = max(last,ts) - output.write(f'{ts:f}\n') + output.write(f'{ts:f}\n'.encode('ascii')) if first: first = False