Compare commits
52
Commits
97ef0e972c
...
refactor
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
814c64cc52 | ||
|
|
c17ccb6563 | ||
|
|
e1d7474e00 | ||
|
|
51bab76c00 | ||
|
|
153873637e | ||
|
|
b0bc0a338e | ||
|
|
4597c459a8 | ||
|
|
5b2afd29b3 | ||
|
|
d1df44c82c | ||
|
|
1a616f6cbd | ||
|
|
d756fede54 | ||
|
|
34d1bc26f3 | ||
|
|
04781021eb | ||
|
|
e1ac3e48bb | ||
|
|
ae1a040590 | ||
|
|
0345293664 | ||
|
|
63d2deafb1 | ||
|
|
6f5dd3a9e4 | ||
|
|
093ad733d3 | ||
|
|
a6a41c112e | ||
|
|
447d04d13d | ||
|
|
2590faf20f | ||
|
|
bea482487e | ||
|
|
c48b77f294 | ||
|
|
0ff2e3a6dd | ||
|
|
294d857206 | ||
|
|
280eb1facb | ||
|
|
3769244bdb | ||
|
|
b241c45efc | ||
|
|
edcbf91101 | ||
|
|
4658cc16e0 | ||
|
|
5bbdfe9929 | ||
|
|
87e326d976 | ||
|
|
e2936d565f | ||
|
|
fa18247fd6 | ||
|
|
87b557bc38 | ||
|
|
a12a4e1b2b | ||
|
|
795d314f13 | ||
|
|
700d1fedc5 | ||
|
|
92df1d24d6 | ||
|
|
ff20f31e74 | ||
|
|
c32018caae | ||
|
|
f6491be5b7 | ||
|
|
a1a1ae5399 | ||
|
|
8dc2af580d | ||
|
|
413df48a77 | ||
|
|
7c539382e6 | ||
|
|
451fd50c68 | ||
|
|
edd034b3b6 | ||
|
|
66e20e126a | ||
|
|
f18ebfadc2 | ||
|
|
e6814f483a |
@@ -4,3 +4,5 @@
|
||||
*.mkv
|
||||
part*
|
||||
venv/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
Generated
+18
-1
@@ -872,6 +872,23 @@ notebook = ["ipywidgets (>=6)"]
|
||||
slack = ["envwrap", "slack-sdk"]
|
||||
telegram = ["envwrap", "requests"]
|
||||
|
||||
[[package]]
|
||||
name = "types-tqdm"
|
||||
version = "4.70.0.20260827"
|
||||
description = "Typing stubs for tqdm"
|
||||
optional = false
|
||||
python-versions = ">=3.10"
|
||||
groups = ["dev"]
|
||||
files = [
|
||||
{file = "types_tqdm-4.70.0.20260827-py3-none-any.whl", hash = "sha256:d6e02da8ec2bc11124e02d4bc84bd8b67a444f4c5decf544846897c8272c7023"},
|
||||
{file = "types_tqdm-4.70.0.20260827.tar.gz", hash = "sha256:d3efd69243ac9093c552f1c0d081b77286d61ace96b6899a142542555bdad2dc"},
|
||||
]
|
||||
|
||||
[package.extras]
|
||||
all = ["types-requests", "types-tensorflow"]
|
||||
requests = ["types-requests"]
|
||||
tensorflow = ["types-tensorflow"]
|
||||
|
||||
[[package]]
|
||||
name = "typing-extensions"
|
||||
version = "4.16.0"
|
||||
@@ -920,4 +937,4 @@ test = ["pytest", "pytest-cov"]
|
||||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = ">=3.11,<4"
|
||||
content-hash = "bb1383a0bd739449d830d9146085c9dcd7f67d7ddaf93bf3e371b90f2d6c069c"
|
||||
content-hash = "76bcb7a8ee79106ddac32401277b69dd4895e0d8067d71277e344c3a4e083be2"
|
||||
|
||||
+34
-1
@@ -15,7 +15,7 @@ packages = [
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
tscut = "tscut:main"
|
||||
tscut = "tscut.cli:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core>=2.0.0,<3.0.0"]
|
||||
@@ -26,4 +26,37 @@ ruff = "^0.16.2"
|
||||
mypy = "^2.3.0"
|
||||
pytest = "^9.1.1"
|
||||
pylint = "^4.0.7"
|
||||
types-tqdm = "^4.70.0.20260827"
|
||||
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py311"
|
||||
line-length = 100
|
||||
src = ["src"]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"F", # Pyflakes
|
||||
"I", # isort
|
||||
"UP", # pyupgrade
|
||||
"B", # flake8-bugbear
|
||||
"SIM", # flake8-simplify
|
||||
"RET", # flake8-return
|
||||
"PERF", # Perflint
|
||||
]
|
||||
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.11"
|
||||
files = ["src"]
|
||||
warn_unused_configs = true
|
||||
show_error_codes = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = [
|
||||
"coloredlogs",
|
||||
"hexdump",
|
||||
]
|
||||
ignore_missing_imports = true
|
||||
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
__pycache__
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
"""Entry point for ``python -m tscut``."""
|
||||
|
||||
from tscut.cli import main
|
||||
from .cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
+33
-469
@@ -6,23 +6,18 @@
|
||||
|
||||
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
|
||||
from datetime import timedelta
|
||||
from functools import cmp_to_key
|
||||
from pathlib import Path
|
||||
|
||||
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
|
||||
from tscut.exceptions import TSCutError
|
||||
from tscut.models import ProcessingOptions
|
||||
from tscut.pipeline import process_recording
|
||||
from tscut.temporaries import TemporaryFiles
|
||||
from tscut.tools.discovery import check_required_tools
|
||||
from tscut.tools.timeframe import compare_time_interval, parse_time_interval
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -65,7 +60,7 @@ def main() -> int:
|
||||
|
||||
if args.coarse and args.threshold is not None:
|
||||
logger.error('--coarse and threshold arguments are exclusive.')
|
||||
exit(-1)
|
||||
return 2
|
||||
|
||||
if (not args.coarse) and args.threshold is None:
|
||||
args.threshold = 0
|
||||
@@ -85,7 +80,7 @@ def main() -> int:
|
||||
ts1, ts2 = parse_time_interval(interval)
|
||||
if ts1 is None or ts2 is None:
|
||||
logger.error("Illegal time interval: %s", interval)
|
||||
exit(-1)
|
||||
return 2
|
||||
parts.append((ts1,ts2))
|
||||
|
||||
# Sort intervals
|
||||
@@ -97,463 +92,32 @@ def main() -> int:
|
||||
ts1, ts2 = part
|
||||
if prevts > ts1:
|
||||
logger.error('Intervals are overlapping')
|
||||
exit(-1)
|
||||
return 2
|
||||
prevts = ts2
|
||||
|
||||
nb_parts = len(parts)
|
||||
temporaries = []
|
||||
|
||||
basename = os.path.splitext(os.path.basename(args.input_file))[0]
|
||||
mp4filename = basename+'.mp4'
|
||||
mkvfilename = basename+'.mkv'
|
||||
config = ProcessingOptions(
|
||||
input_file=Path(args.input_file),
|
||||
output_file=Path(args.output_file),
|
||||
parts=parts,
|
||||
framerate = args.framerate,
|
||||
verbose=args.verbose,
|
||||
coarse=args.coarse,
|
||||
threshold=args.threshold,
|
||||
keep_temporaries=args.keep,
|
||||
subtitles_ocr=args.srt,
|
||||
dump_memory=args.dump,
|
||||
tools_paths = paths,
|
||||
all_optional_tools = all_optional_tools)
|
||||
|
||||
logger.debug("Configuration: %s", config)
|
||||
temporaries = TemporaryFiles(config.keep_temporaries)
|
||||
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:
|
||||
process_recording(config, temporaries)
|
||||
except TSCutError as exc:
|
||||
logger.error("%s", exc)
|
||||
return 1
|
||||
finally:
|
||||
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"))
|
||||
temporaries.cleanup()
|
||||
|
||||
return 0
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import logging
|
||||
from typing import IO
|
||||
|
||||
from tscut.exceptions import TemporaryFileError
|
||||
from tscut.tools.mkvtoolnix import extract_mkv_part
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TODO: finish this procedure
|
||||
def do_coarse_processing(ffmpeg_path:str, ffprobe_path:str, mkvmerge_path:str,
|
||||
input_file: IO[bytes], begin, end, nb_frames, framerate,
|
||||
files_prefix, streams, width, height, temporaries, dump_mem_fd) -> None:
|
||||
# Internal video with all streams (video, audio and subtitles)
|
||||
internal_mkv_name = f'{files_prefix}.mkv'
|
||||
|
||||
try:
|
||||
internal_mkv = open(internal_mkv_name, 'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create {internal_mkv_name}") from e
|
||||
|
||||
# Extract internal part of MKV
|
||||
extract_mkv_part(mkvmerge_path=mkvmerge_path, input_file=input_file, output_file=internal_mkv,
|
||||
begin=begin, end=end)
|
||||
|
||||
temporaries.add(internal_mkv)
|
||||
@@ -0,0 +1,365 @@
|
||||
# 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 BinaryIO, Sequence, TextIO
|
||||
|
||||
import hexdump
|
||||
from tqdm import tqdm
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import InvalidMediaError, TemporaryFileError
|
||||
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.temporaries import TemporaryFiles
|
||||
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):
|
||||
raise InvalidMediaError(f"Impossible to find a valid header: {isheader}")
|
||||
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: TemporaryFiles) -> 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)
|
||||
try:
|
||||
head_iframe_ts = get_ts_frame(head_iframe)
|
||||
except InvalidMediaError:
|
||||
raise InvalidMediaError("Impossible to retrieve timestamp of first i-frame.")
|
||||
try:
|
||||
tail_iframe_ts = get_ts_frame(tail_iframe)
|
||||
except InvalidMediaError:
|
||||
raise InvalidMediaError("Impossible to retrieve timestamp of final i-frame.")
|
||||
|
||||
assert head_iframe_ts <= tail_iframe_ts
|
||||
|
||||
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)
|
||||
|
||||
if head_iframe_ts < tail_iframe_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 as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {internal_mkv_name}") from e
|
||||
|
||||
try:
|
||||
internal_novideo_mkv = open(internal_novideo_mkv_name, 'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: \
|
||||
{internal_novideo_mkv_name}") from e
|
||||
|
||||
try:
|
||||
internal_h264 = open(internal_h264_name, 'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {internal_h264_name}") from e
|
||||
|
||||
try:
|
||||
internal_h264_ts = open(internal_h264_ts_name, 'w+', encoding='utf8')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {internal_h264_ts_name}") from e
|
||||
|
||||
# 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.add(internal_mkv)
|
||||
temporaries.add(internal_h264)
|
||||
temporaries.add(internal_h264_ts)
|
||||
temporaries.add(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.add(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 as e:
|
||||
raise TemporaryFileError("Impossible to create file full H264 stream.") from e
|
||||
|
||||
logger.info('Merging all H264 tracks')
|
||||
concatenate_h264_parts(h264parts=h264parts, output=full_h264)
|
||||
temporaries.add(full_h264)
|
||||
|
||||
try:
|
||||
full_h264_ts = open(f'{media.basename}-ts.txt', 'w+', encoding='utf8')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError("Impossible to create file containing all video timestamps.")\
|
||||
from e
|
||||
|
||||
logger.info('Merging H264 timestamps')
|
||||
concatenate_h264_ts_parts(h264_ts_parts=h264_ts, output=full_h264_ts)
|
||||
temporaries.add(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 as e:
|
||||
raise TemporaryFileError(f"Impossible to open file: {final_novideo_name}") from e
|
||||
|
||||
temporaries.add(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
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
Useful SPS/PPS discussion:
|
||||
1.https://copyprogramming.com/howto/including-sps-and-pps-in-a-raw-h264-track
|
||||
2. https://gitlab.com/mbunkus/mkvtoolnix/-/issues/2390
|
||||
|
||||
Strategy: a possible way of handling multiple SPS/PPS gracefully.
|
||||
Encode each head and trailer with FFMPEG using only I-frame (to be sure the NAL unit will never refer to another image).
|
||||
Encode using an different SPS-ID all of them (using sps-id parameter of libx264 library, e.g 1 instead of 0).
|
||||
For the video track produce only a raw H264 file and a file containing timestamps of the different frames.
|
||||
For the rest of the tracks (audio, subtitles) produce directly a MKV (this is already done).
|
||||
Concatenate all raw H264 in a giant one (like cat), and the same for timestamps of video frames (to keep sound and video synchronized).
|
||||
Then use mkvmerge to remux the H264 track and the rest of tracks.
|
||||
MKVmerge "concatenate" subcommand is able to concatenate different SPS/PPS data into a bigger Private Codec Data.
|
||||
However, this is proved to be not reliable.
|
||||
Sometimes it results in a AVC context containing a single SPS/PPS.
|
||||
So we have to rely on a manual parsing of the H264 AVC context of original movie and the ones produced for headers and trailers, and then merging them into a bigger AVC context.
|
||||
Then finally, change the Private Codec Data in the final MKV.
|
||||
|
||||
|
||||
Extract SPS/PPS:
|
||||
1. https://gitlab.com/mbunkus/mkvtoolnix/-/issues/2390
|
||||
```bash
|
||||
ffmpeg -i <InputFile (before concatenation)> -c:v copy -an -sn -bsf:v trace_headers -t 0.01 -report -loglevel 0 -f null -
|
||||
``
|
||||
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
class TSCutError(Exception):
|
||||
"""Base exception for tscut."""
|
||||
|
||||
|
||||
class MissingToolError(TSCutError):
|
||||
"""Raised when a required external tool is unavailable."""
|
||||
|
||||
|
||||
class ExternalToolError(TSCutError):
|
||||
"""Raised when an external tool fails."""
|
||||
|
||||
class InvalidMediaError(TSCutError):
|
||||
"""Raised when input media cannot be processed."""
|
||||
|
||||
class TemporaryFileError(TSCutError):
|
||||
"""Raised when the manipulation of temporary file goes wrong"""
|
||||
|
||||
class UnimplementedFeatureError(TSCutError):
|
||||
"""Raised when a rare feature is encountered and not yet implemented"""
|
||||
@@ -0,0 +1,303 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from math import floor
|
||||
from typing import IO
|
||||
|
||||
import hexdump
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.h264.bitstream import (
|
||||
rbsp_to_sodb,
|
||||
read_bit,
|
||||
read_bits,
|
||||
read_byte,
|
||||
read_long,
|
||||
read_word,
|
||||
sodb_to_rbsp,
|
||||
write_bits,
|
||||
write_byte,
|
||||
write_word,
|
||||
)
|
||||
from tscut.h264.parameters import (
|
||||
PPS,
|
||||
SPS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AVCDecoderConfiguration:
|
||||
configuration_version:int=1 # u(8)
|
||||
avc_profile_indication:int=0 # u(8)
|
||||
profile_compatibility:int=0 # u(8)
|
||||
avc_level_indication:int=0 # u(8)
|
||||
length_size_minus_one:int=0 # u(2) (0,1 or 3)
|
||||
num_of_sequence_parameter_sets:int=0 # u(5)
|
||||
sps:dict = field(default_factory=dict)
|
||||
num_of_picture_parameter_sets:int=0 #u(8)
|
||||
pps:dict = field(default_factory=dict)
|
||||
chroma_format:int=0 # u(2)
|
||||
bit_depth_luma_minus8:int=0 # u(3)
|
||||
bit_depth_chroma_minus8:int=0 # u(3)
|
||||
num_of_sequence_parameter_set_ext:int=0 # u(8)
|
||||
spsext:dict = field(default_factory=dict)
|
||||
|
||||
def __init__(self):
|
||||
self.sps = {}
|
||||
self.spsext = {}
|
||||
self.pps = {}
|
||||
|
||||
def from_bytes(self, buf):
|
||||
logger.debug('Parsing: %s', (hexdump.dump(buf,sep=':')))
|
||||
bit_position = 0
|
||||
bit_position, self.configuration_version = read_byte(buf, bit_position)
|
||||
bit_position, self.avc_profile_indication = read_byte(buf, bit_position)
|
||||
bit_position, self.profile_compatibility = read_byte(buf, bit_position)
|
||||
bit_position, self.avc_level_indication = read_byte(buf, bit_position)
|
||||
bit_position, v = read_bits(buf, bit_position, 6)
|
||||
if v != 0b111111:
|
||||
raise ValueError(f'Reserved bits are not equal to 0b111111: {v:x}')
|
||||
bit_position, self.length_size_minus_one = read_bits(buf, bit_position, 2)
|
||||
bit_position, v = read_bits(buf, bit_position, 3)
|
||||
if v != 0b111:
|
||||
raise ValueError(f'Reserved bits are not equal to 0b111: {v:x}')
|
||||
bit_position, self.num_of_sequence_parameter_sets= read_bits(buf, bit_position, 5)
|
||||
logger.debug('Number of SPS: %d', self.num_of_sequence_parameter_sets)
|
||||
for _ in range(self.num_of_sequence_parameter_sets):
|
||||
bit_position, length = read_word(buf, bit_position)
|
||||
if bit_position % 8 != 0:
|
||||
raise ValueError(f'SPS is not located at a byte boundary: {bit_position:d}')
|
||||
|
||||
sps = SPS()
|
||||
sodb = rbsp_to_sodb(buf[floor(bit_position/8):])
|
||||
bit_length = sps.from_bytes(sodb)
|
||||
spsid = sps.seq_parameter_set_id
|
||||
self.sps[spsid] = sps
|
||||
|
||||
parsed_length = floor(bit_length/8)
|
||||
logger.debug('Expected length of SPS: %d bytes. Parsed: %d bytes', length,
|
||||
parsed_length)
|
||||
# Parse length can be shorter than length because of rewriting from RBSP to SODB
|
||||
# (that is shorter).
|
||||
# So we advance of indicated length.
|
||||
bit_position+=length*8
|
||||
|
||||
logger.debug('Bit position:%d. Reading one byte of: %s', bit_position,
|
||||
hexdump.dump(buf[floor(bit_position/8):], sep=':'))
|
||||
bit_position, self.num_of_picture_parameter_sets = read_byte(buf, bit_position)
|
||||
logger.debug('Number of PPS: %d', self.num_of_picture_parameter_sets)
|
||||
for _ in range(self.num_of_picture_parameter_sets):
|
||||
bit_position, length = read_word(buf, bit_position)
|
||||
if bit_position % 8 != 0:
|
||||
raise ValueError('PPS is not located at a byte boundary: {bit_position:d}')
|
||||
|
||||
pps = PPS()
|
||||
sodb = rbsp_to_sodb(buf[floor(bit_position/8):])
|
||||
bit_length = pps.from_bytes(sodb, self.chroma_format)
|
||||
ppsid = pps.pic_parameter_set_id
|
||||
self.pps[ppsid] = pps
|
||||
|
||||
parsed_length = floor(bit_length/8)
|
||||
logger.debug('Expected length of PPS: %d bytes. Parsed: %d bytes', length,
|
||||
parsed_length)
|
||||
# Parse length can be shorter than length because of rewriting from RBSP to SODB
|
||||
# (that is shorter).
|
||||
# So we advance of indicated length.
|
||||
bit_position+=length*8
|
||||
|
||||
logger.debug('Remaining bits: %s', hexdump.dump(buf[floor(bit_position/8):]))
|
||||
|
||||
if self.avc_profile_indication in [100, 110, 122, 144]:
|
||||
bit_position, reserved = read_bits(buf, bit_position, 6)
|
||||
if reserved != 0b111111:
|
||||
raise ValueError(f'Reserved bits are different from 111111: {reserved:x}')
|
||||
bit_position, self.chroma_format = read_bits(buf, bit_position, 2)
|
||||
bit_position, reserved = read_bits(buf, bit_position, 5)
|
||||
if reserved != 0b11111:
|
||||
raise ValueError(f'Reserved bits are different from 11111: {reserved:x}')
|
||||
bit_position, self.bit_depth_luma_minus8 = read_bits(buf, bit_position, 3)
|
||||
bit_position, reserved = read_bits(buf, bit_position, 5)
|
||||
if reserved != 0b11111:
|
||||
raise ValueError(f'Reserved bits are different from 11111: {reserved:x}')
|
||||
bit_position, self.bit_depth_chroma_minus8 = read_bits(buf, bit_position, 3)
|
||||
bit_position, self.num_of_sequence_parameter_set_ext = read_byte(buf, bit_position)
|
||||
for _ in range(self.num_of_sequence_parameter_set_ext):
|
||||
# TODO: parse SPSextended
|
||||
logger.error('Parsing of SPS extended not yet implemented !')
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def to_bytes(self):
|
||||
buf = bytearray()
|
||||
bit_position = 0
|
||||
bit_position = write_byte(buf, bit_position, self.configuration_version)
|
||||
bit_position = write_byte(buf, bit_position, self.avc_profile_indication)
|
||||
bit_position = write_byte(buf, bit_position, self.profile_compatibility)
|
||||
bit_position = write_byte(buf, bit_position, self.avc_level_indication)
|
||||
bit_position = write_bits(buf, bit_position, 0b111111, 6)
|
||||
bit_position = write_bits(buf, bit_position, self.length_size_minus_one, 2)
|
||||
bit_position = write_bits(buf, bit_position, 0b111, 3)
|
||||
bit_position = write_bits(buf, bit_position, self.num_of_sequence_parameter_sets, 5)
|
||||
for spsid, sps in self.sps.items():
|
||||
sodb = sps.to_bytes()
|
||||
sodb_length = len(sodb)
|
||||
rbsp = sodb_to_rbsp(sodb)
|
||||
rbsp_length = len(rbsp)
|
||||
|
||||
logger.debug('SODB length: %d RBSP length:%d', sodb_length, rbsp_length)
|
||||
|
||||
bit_position = write_word(buf, bit_position, rbsp_length)
|
||||
buf.extend(rbsp)
|
||||
bit_position+=rbsp_length*8
|
||||
|
||||
logger.debug('2. Buffer: %s', hexdump.dump(buf, sep=':'))
|
||||
|
||||
bit_position = write_byte(buf, bit_position, self.num_of_picture_parameter_sets)
|
||||
for ppsid, lpps in self.pps.items():
|
||||
logger.debug('Writing PPS: %d', ppsid)
|
||||
# TODO: does chroma_format should come from self ?
|
||||
sodb = lpps.to_bytes(self.chroma_format)
|
||||
sodb_length = len(sodb)
|
||||
rbsp = sodb_to_rbsp(sodb)
|
||||
rbsp_length = len(rbsp)
|
||||
|
||||
logger.debug('SODB length: %d RBSP length:%d', sodb_length, rbsp_length)
|
||||
|
||||
bit_position = write_word(buf, bit_position, rbsp_length)
|
||||
buf.extend(rbsp)
|
||||
bit_position+=rbsp_length*8
|
||||
|
||||
if self.avc_profile_indication in [ 100, 110, 122, 144]:
|
||||
bit_position = write_bits(buf, bit_position, 0b111111, 6)
|
||||
bit_position = write_bits(buf, bit_position, self.chroma_format, 2)
|
||||
bit_position = write_bits(buf, bit_position, 0b11111, 5)
|
||||
bit_position = write_bits(buf, bit_position, self.bit_depth_luma_minus8, 3)
|
||||
bit_position = write_bits(buf, bit_position, 0b11111, 5)
|
||||
bit_position = write_bits(buf, bit_position, self.bit_depth_chroma_minus8, 3)
|
||||
bit_position = write_byte(buf, bit_position, self.num_of_sequence_parameter_set_ext)
|
||||
for _ in range(self.num_of_sequence_parameter_set_ext):
|
||||
# TODO: dump SPSextended
|
||||
logger.error('Dumping SPS extended not yet implemented')
|
||||
pass
|
||||
|
||||
return buf
|
||||
|
||||
def merge(self, config):
|
||||
# Check config compatibility
|
||||
if self.configuration_version != config.configuration_version:
|
||||
raise ValueError(f'Configuration versions are different: {self.configuration_version:d}\
|
||||
vs {config.configuration_version:d}')
|
||||
if self.avc_profile_indication != config.avc_profile_indication:
|
||||
raise ValueError(f'AVC profiles are different: {self.avc_profile_indication:d} vs \
|
||||
{config.avc_profile_indication:d}')
|
||||
if self.profile_compatibility != config.profile_compatibility:
|
||||
raise ValueError(f'Profile compatilities are different: {self.profile_compatibility:d} \
|
||||
vs {config.profile_compatibility:d}')
|
||||
if self.avc_level_indication != config.avc_level_indication:
|
||||
raise ValueError(f'Level indications are different: {self.avc_level_indication:d} vs \
|
||||
{config.avc_level_indication:d}')
|
||||
if self.length_size_minus_one != config.length_size_minus_one:
|
||||
raise ValueError(f'Length units are different: {self.length_size_minus_one:d} vs \
|
||||
{config.length_size_minus_one:d}')
|
||||
if self.chroma_format != config.chroma_format:
|
||||
raise ValueError(f'Colour format are different: {self.chroma_format:d} vs \
|
||||
{config.chroma_format:d}')
|
||||
if self.bit_depth_luma_minus8 != config.bit_depth_luma_minus8:
|
||||
raise ValueError(f'Depth of luminance are different: {self.bit_depth_luma_minus8:d} vs \
|
||||
{config.bit_depth_luma_minus8:d}')
|
||||
if self.bit_depth_chroma_minus8 != config.bit_depth_chroma_minus8:
|
||||
raise ValueError(f'Depth of chromaticity are different: \
|
||||
{self.bit_depth_chroma_minus8:d} vs {config.bit_depth_luma_minus8:d}')
|
||||
|
||||
for spsid in config.sps:
|
||||
sps = config.sps[spsid]
|
||||
if spsid in self.sps:
|
||||
localsps = self.sps[spsid]
|
||||
if sps!=localsps:
|
||||
raise ValueError(f'Profile are not compatible. They contain two different SPS\
|
||||
with the same identifier ({spsid:d}): {localsps}\n{sps}\n')
|
||||
self.sps[spsid] = sps
|
||||
|
||||
self.num_of_sequence_parameter_sets = len(self.sps)
|
||||
|
||||
for ppsid in config.pps:
|
||||
pps = config.pps[ppsid]
|
||||
if ppsid in self.pps:
|
||||
localpps = self.pps[ppsid]
|
||||
if pps!=localpps:
|
||||
raise ValueError(f'Profile are not compatible. They contain two different PPS\
|
||||
with the same identifier ({ppsid:d}): {localpps}\n{pps}\n')
|
||||
self.pps[ppsid] = pps
|
||||
|
||||
self.num_of_picture_parameter_sets = len(self.pps)
|
||||
|
||||
# TODO: do the same with extended SPS !
|
||||
|
||||
@typechecked
|
||||
def parse_codec_private(codec_private_data: bytes | bytearray) -> AVCDecoderConfiguration:
|
||||
if codec_private_data[0] != 0x63:
|
||||
raise ValueError(f'Matroska header is wrong: {codec_private_data[0]:x}')
|
||||
if codec_private_data[1] != 0xA2:
|
||||
raise ValueError(f'Matroska header is wrong: {codec_private_data[1]:x}')
|
||||
length = codec_private_data[2]
|
||||
if length == 0:
|
||||
raise ValueError('Matroska length cannot start with zero byte.')
|
||||
for nb_zeroes in range(8):
|
||||
b = read_bit(codec_private_data[2:], nb_zeroes)
|
||||
if b != 0:
|
||||
break
|
||||
mask = 2^(7-nb_zeroes)-1
|
||||
length = codec_private_data[2] and mask
|
||||
for i in range(nb_zeroes):
|
||||
length*=256
|
||||
length+=(codec_private_data[3+i])
|
||||
byte_position = 3+nb_zeroes
|
||||
avcconfig = AVCDecoderConfiguration()
|
||||
avcconfig.from_bytes(codec_private_data[byte_position:])
|
||||
|
||||
return avcconfig
|
||||
|
||||
@typechecked
|
||||
def get_avc_config_from_h264(input_file: IO[bytes]) -> AVCDecoderConfiguration:
|
||||
# TODO: improve this ...
|
||||
rbsp = input_file.read(1000)
|
||||
sodb = rbsp_to_sodb(rbsp)
|
||||
|
||||
bit_position = 0
|
||||
bit_position, start_code = read_long(sodb, bit_position)
|
||||
if start_code != 1:
|
||||
raise ValueError(f'Starting code not detected: {start_code:x}')
|
||||
sps = SPS()
|
||||
bit_length = sps.from_bytes(sodb[4:])
|
||||
bit_position+=bit_length
|
||||
|
||||
bit_position, start_code = read_long(sodb, bit_position)
|
||||
if start_code != 1:
|
||||
raise ValueError(f'Starting code not detected: {start_code:x}')
|
||||
pps = PPS()
|
||||
bit_length = pps.from_bytes(sodb[floor(bit_position/8):], sps.chroma_format_idc)
|
||||
logger.debug(pps)
|
||||
|
||||
avcconfig = AVCDecoderConfiguration()
|
||||
avcconfig.configuration_version = 1
|
||||
avcconfig.avc_profile_indication = sps.profile_idc
|
||||
avcconfig.profile_compatibility = 0
|
||||
avcconfig.avc_level_indication = sps.level_idc
|
||||
avcconfig.length_size_minus_one = 3
|
||||
avcconfig.num_of_sequence_parameter_sets = 1
|
||||
avcconfig.num_of_picture_parameter_sets = 1
|
||||
avcconfig.num_of_sequence_parameter_set_ext = 0
|
||||
avcconfig.chroma_format = sps.chroma_format_idc
|
||||
avcconfig.bit_depth_chroma_minus8 = sps.bit_depth_chroma_minus8
|
||||
avcconfig.bit_depth_luma_minus8 = sps.bit_depth_luma_minus8
|
||||
avcconfig.sps[sps.seq_parameter_set_id] = sps
|
||||
avcconfig.pps[pps.pic_parameter_set_id] = pps
|
||||
|
||||
return avcconfig
|
||||
@@ -0,0 +1,303 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
""" All the following code is a transposition of documents:
|
||||
ISO/IEC H.264-201602
|
||||
ISO/IEC 14496-15
|
||||
"""
|
||||
|
||||
import logging
|
||||
from math import floor, log
|
||||
|
||||
import hexdump
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import UnimplementedFeatureError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@typechecked
|
||||
def read_bit(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
"""
|
||||
Read a single bit from a byte buffer.
|
||||
|
||||
This function is part of the implementation of the H.264/AVC video compression standard,
|
||||
as specified in ISO/IEC H.264-201602 and ISO/IEC 14496-15.
|
||||
It takes a byte buffer and a bit position as input, and returns a tuple containing
|
||||
the updated bit position and the value of the bit at the specified position.
|
||||
|
||||
Args:
|
||||
buf (bytes): The byte buffer to read from.
|
||||
bit_position (int): The position of the bit to read, starting from 0.
|
||||
|
||||
Returns:
|
||||
tuple[int, int]: A tuple containing the updated bit position (bit_position + 1) and the
|
||||
value of the bit (0 or 1).
|
||||
"""
|
||||
byte_position = floor(floor(bit_position/8))
|
||||
byte = buf[byte_position]
|
||||
bit = (byte >> (7-(bit_position % 8))) & 1
|
||||
return bit_position+1, bit
|
||||
|
||||
@typechecked
|
||||
def read_boolean(buf:bytes|bytearray, bit_position: int) -> tuple[int, bool]:
|
||||
"""
|
||||
Read a boolean value from a byte buffer.
|
||||
|
||||
This function reads a single bit from the byte buffer at the specified position and interprets
|
||||
it as a boolean value.
|
||||
It returns a tuple containing the updated bit position and the boolean value.
|
||||
|
||||
Args:
|
||||
buf (bytes): The byte buffer to read from.
|
||||
bit_position (int): The position of the bit to read, starting from 0.
|
||||
|
||||
Returns:
|
||||
tuple[int, bool]: A tuple containing the updated bit position and the boolean value
|
||||
(True if the bit is 1, False if the bit is 0).
|
||||
"""
|
||||
bit_position, b = read_bit(buf, bit_position)
|
||||
return bit_position, b==1
|
||||
|
||||
@typechecked
|
||||
def read_bits(buf:bytes|bytearray, bit_position: int, nb_bits: int) -> tuple[int, int]:
|
||||
v = 0
|
||||
for _ in range(nb_bits):
|
||||
bit_position, bit = read_bit(buf, bit_position)
|
||||
v = v*2+bit
|
||||
return bit_position, v
|
||||
|
||||
@typechecked
|
||||
def read_byte(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
bit_position, b = read_bits(buf, bit_position, 8)
|
||||
return bit_position, b
|
||||
|
||||
@typechecked
|
||||
def read_word(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
bit_position, w = read_bits(buf, bit_position, 16)
|
||||
return bit_position, w
|
||||
|
||||
@typechecked
|
||||
def read_long(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
bit_position, value = read_bits(buf, bit_position, 32)
|
||||
return bit_position, value
|
||||
|
||||
@typechecked
|
||||
def read_unsigned_exp_golomb(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
nb_zeroes=0
|
||||
while True:
|
||||
bit_position, b = read_bit(buf, bit_position)
|
||||
if b!=0:
|
||||
break
|
||||
nb_zeroes+=1
|
||||
|
||||
v1 = 1
|
||||
bit_position, v2 = read_bits(buf, bit_position, nb_zeroes)
|
||||
v = (v1<<nb_zeroes)+v2
|
||||
return bit_position, v-1
|
||||
|
||||
@typechecked
|
||||
def read_signed_exp_golomb(buf:bytes|bytearray, bit_position: int) -> tuple[int, int]:
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
match v%2:
|
||||
case 0:
|
||||
return bit_position, -(v>>1)
|
||||
case 1:
|
||||
return bit_position, (v+1)>>1
|
||||
case _:
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
@typechecked
|
||||
def write_bit(buf:bytearray, bit_position: int, b) -> int:
|
||||
buf_length = len(buf)
|
||||
byte_position = floor(bit_position/8)
|
||||
|
||||
if byte_position >= buf_length:
|
||||
extension = bytearray(byte_position+1-buf_length)
|
||||
buf.extend(extension)
|
||||
|
||||
buf[byte_position] |= (b<<(7-(bit_position % 8)))
|
||||
bit_position+=1
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_boolean(buf:bytearray, bit_position: int, b: bool) -> int:
|
||||
if b:
|
||||
bit_position = write_bit(buf, bit_position, 1)
|
||||
else:
|
||||
bit_position = write_bit(buf, bit_position, 0)
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_bits(buf:bytearray, bit_position: int, v, size) -> int:
|
||||
for i in range(size-1,-1,-1):
|
||||
b = (v>>i)&1
|
||||
bit_position = write_bit(buf, bit_position, b)
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_byte(buf:bytearray, bit_position: int, v) -> int:
|
||||
bit_position = write_bits(buf, bit_position, v, 8)
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_word(buf:bytearray, bit_position: int, v) -> int:
|
||||
bit_position = write_bits(buf, bit_position, v, 16)
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_long(buf:bytearray, bit_position: int, v) -> int:
|
||||
bit_position = write_bits(buf, bit_position, v, 32)
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_unsigned_exp_golomb(buf:bytearray, bit_position: int, v) -> int:
|
||||
n = floor(log(v+1)/log(2))+1
|
||||
# Write zeroes
|
||||
bit_position = write_bits(buf, bit_position, 0, n-1)
|
||||
bit_position = write_bit(buf, bit_position, 1)
|
||||
bit_position = write_bits(buf, bit_position, v+1, n-1)
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_signed_exp_golomb(buf:bytearray, bit_position: int, v) -> int:
|
||||
if v <= 0:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, -v*2)
|
||||
else:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v*2-1)
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def parse_rbsp_trailing_bits(buf:bytes|bytearray, bit_position: int) -> int:
|
||||
bit_position, one = read_bit(buf, bit_position)
|
||||
if one==0:
|
||||
raise ValueError(f'Stop bit should be equal to one. Read: {one:d}')
|
||||
while bit_position%8 != 0:
|
||||
bit_position, zero = read_bit(buf, bit_position)
|
||||
if zero==1:
|
||||
raise ValueError('Trailing bit should be equal to zero')
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def write_rbsp_trailing_bits(buf:bytearray, bit_position: int) -> int:
|
||||
bit_position = write_bit(buf, bit_position, 1)
|
||||
while bit_position%8 != 0:
|
||||
bit_position = write_bit(buf, bit_position, 0)
|
||||
|
||||
return bit_position
|
||||
|
||||
@typechecked
|
||||
def more_rbsp_data(buf:bytes|bytearray, bit_position: int) -> bool:
|
||||
logger.debug('Is there more data in buffer of length: %d at bit position: %d',
|
||||
len(buf), bit_position)
|
||||
|
||||
byte_length = len(buf)
|
||||
bit_length = byte_length*8
|
||||
|
||||
# We are at the end of buffer
|
||||
if bit_position == bit_length:
|
||||
return False
|
||||
|
||||
found = False
|
||||
for i in range(bit_length-1,-1,-1):
|
||||
pos, b = read_bit(buf, i)
|
||||
if b == 1:
|
||||
found = True
|
||||
break
|
||||
|
||||
if not found:
|
||||
raise ValueError('Impossible to find trailing stop bit !')
|
||||
|
||||
# No more data
|
||||
if bit_position == pos:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
# Convert from RBSP (Raw Byte Sequence Payload) to SODB (String Of Data Bits)
|
||||
@typechecked
|
||||
def rbsp_to_sodb(buf:bytes) -> bytes:
|
||||
logger.debug('RBSP: %s', hexdump.dump(buf, sep=':'))
|
||||
|
||||
res = buf
|
||||
for b in [ b'\x00', b'\x01', b'\x02', b'\x03']:
|
||||
pattern = b'\x00\x00\x03'+b
|
||||
replacement = b'\x00\x00' + b
|
||||
res = res.replace(pattern, replacement)
|
||||
|
||||
logger.debug('SODB: %s', hexdump.dump(res, sep=':'))
|
||||
return res
|
||||
|
||||
# Reverse operation SODB to RBSP.
|
||||
@typechecked
|
||||
def sodb_to_rbsp(buf:bytes) -> bytes:
|
||||
logger.debug('SODB: %s', hexdump.dump(buf, sep=':'))
|
||||
|
||||
res = buf
|
||||
for b in [ b'\x03', b'\x00', b'\x01', b'\x02']:
|
||||
pattern = b'\x00\x00'+b
|
||||
replacement = b'\x00\x00\x03' + b
|
||||
res = res.replace(pattern, replacement)
|
||||
|
||||
logger.debug('RBSP: %s', hexdump.dump(res, sep=':'))
|
||||
return res
|
||||
|
||||
# Useful for SPS and PPS
|
||||
@typechecked
|
||||
def parse_scaling_list(buf:bytes|bytearray, bit_position: int, size) -> tuple[int,list[int]]:
|
||||
res = []
|
||||
last_scale = 8
|
||||
next_scale = 8
|
||||
for _ in range(size):
|
||||
if next_scale != 0:
|
||||
bit_position, delta_scale = read_signed_exp_golomb(buf, bit_position)
|
||||
next_scale = (last_scale+delta_scale+256) % 256
|
||||
v = last_scale if next_scale==0 else next_scale
|
||||
res.append(v)
|
||||
last_scale = v
|
||||
|
||||
return bit_position,res
|
||||
|
||||
# TODO: test optimized version.
|
||||
# The ISO/IEC H.264-201602 seems to take into account the case where the end of the deltas list
|
||||
# is full of zeroes.
|
||||
@typechecked
|
||||
def write_scaling_list(buf:bytearray, bit_position: int, size, matrix:list[int],
|
||||
optimized: bool = False) -> int:
|
||||
logger.debug('Dumping matrix: %s of size: %d, size parameter: %d.', matrix, len(matrix), size)
|
||||
|
||||
prev = 8
|
||||
deltas = []
|
||||
for i in range(size):
|
||||
v = matrix[i]
|
||||
delta = v - prev
|
||||
deltas.append(delta)
|
||||
prev = v
|
||||
|
||||
if not optimized:
|
||||
for delta in deltas:
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position, delta)
|
||||
else:
|
||||
logger.error('Not yet implemented')
|
||||
raise UnimplementedFeatureError("Optimized scaling list of H264 header is not implemented")
|
||||
# reverse = deltas.reverse()
|
||||
# compressed = False
|
||||
# while len(reverse)>0:
|
||||
# if reverse[0] == 0:
|
||||
# compressed = True
|
||||
# reverse.pop()
|
||||
# else:
|
||||
# break
|
||||
# deltas = reverse.reverse()
|
||||
# if compressed:
|
||||
# deltas.append(0)
|
||||
# for delta in deltas:
|
||||
# bit_position = write_signed_exp_golomb(buf, bit_position, delta)
|
||||
|
||||
return bit_position
|
||||
@@ -0,0 +1,716 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from math import ceil, floor, log
|
||||
|
||||
import hexdump
|
||||
|
||||
from tscut.h264.bitstream import (
|
||||
more_rbsp_data,
|
||||
parse_rbsp_trailing_bits,
|
||||
parse_scaling_list,
|
||||
read_bit,
|
||||
read_bits,
|
||||
read_boolean,
|
||||
read_byte,
|
||||
read_long,
|
||||
read_signed_exp_golomb,
|
||||
read_unsigned_exp_golomb,
|
||||
read_word,
|
||||
write_bit,
|
||||
write_bits,
|
||||
write_boolean,
|
||||
write_byte,
|
||||
write_long,
|
||||
write_rbsp_trailing_bits,
|
||||
write_scaling_list,
|
||||
write_signed_exp_golomb,
|
||||
write_unsigned_exp_golomb,
|
||||
write_word,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class HRD:
|
||||
cpb_cnt_minus1: int=0
|
||||
bit_rate_scale: int=0
|
||||
cpb_size_scale: int=0
|
||||
bit_rate_value_minus1: dict = field(default_factory=dict)
|
||||
cpb_size_value_minus1: dict = field(default_factory=dict)
|
||||
cbr_flag: dict = field(default_factory=dict)
|
||||
initial_cpb_removal_delay_length_minus1: int=0
|
||||
cpb_removal_delay_length_minus1: int=0
|
||||
dpb_output_delay_length_minus1: int=0
|
||||
time_offset_length: int=0
|
||||
|
||||
def __init__(self):
|
||||
self.bit_rate_value_minus1 = {}
|
||||
self.cpb_size_value_minus1 = {}
|
||||
self.cbr_flag = {}
|
||||
|
||||
def from_bytes(self, buf, bit_position):
|
||||
bit_position, self.cpb_cnt_minus1 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.bit_rate_scale = read_bits(buf, bit_position, 4)
|
||||
bit_position, self.cpb_size_scale = read_bits(buf, bit_position, 4)
|
||||
for i in range(self.cpb_cnt_minus1+1):
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.bit_rate_value_minus1[i] = v
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.cpb_size_value_minus1[i] = v
|
||||
bit_position, b = read_boolean(buf, bit_position)
|
||||
self.cbr_flag[i] = b
|
||||
bit_position, self.initial_cpb_removal_delay_length_minus1 = read_bits(buf, bit_position, 5)
|
||||
bit_position, self.cpb_removal_delay_length_minus1 = read_bits(buf, bit_position, 5)
|
||||
bit_position, self.dpb_output_delay_length_minus1 = read_bits(buf, bit_position, 5)
|
||||
bit_position, self.time_offset_length = read_bits(buf, bit_position, 5)
|
||||
|
||||
return bit_position
|
||||
|
||||
def to_bytes(self, buf, bit_position):
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.cpb_cnt_minus1)
|
||||
bit_position = write_bits(buf, bit_position, self.bit_rate_scale, 4)
|
||||
bit_position = write_bits(buf, bit_position, self.cpb_size_scale, 4)
|
||||
for i in range(self.cpb_cnt_minus1+1):
|
||||
v = self.bit_rate_value_minus1[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
v = self.cpb_size_value_minus1[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
b = self.cbr_flag[i]
|
||||
bit_position = write_boolean(buf, bit_position, b)
|
||||
bit_position = write_bits(buf, bit_position, self.initial_cpb_removal_delay_length_minus1,
|
||||
5)
|
||||
bit_position = write_bits(buf, bit_position, self.cpb_removal_delay_length_minus1, 5)
|
||||
bit_position = write_bits(buf, bit_position, self.dpb_output_delay_length_minus1, 5)
|
||||
bit_position = write_bits(buf, bit_position, self.time_offset_length, 5)
|
||||
|
||||
return bit_position
|
||||
|
||||
@dataclass
|
||||
class VUI:
|
||||
aspect_ratio_info_present_flag:bool=False
|
||||
aspect_ratio_idc:int=0
|
||||
sar_width:int=0
|
||||
sar_height:int=0
|
||||
overscan_info_present_flag:bool=False
|
||||
overscan_appropriate_flag:bool=False
|
||||
video_signal_type_present_flag:bool=False
|
||||
video_format:int=0
|
||||
video_full_range_flag:bool=False
|
||||
colour_description_present_flag:bool=False
|
||||
colour_primaries:int=0
|
||||
transfer_characteristics:int=0
|
||||
matrix_coefficients:int=0
|
||||
chroma_loc_info_present_flag:bool=False
|
||||
chroma_sample_loc_type_top_field:int=0
|
||||
chroma_sample_loc_type_bottom_field:int=0
|
||||
timing_info_present_flag:bool=False
|
||||
num_units_in_tick:int=0
|
||||
time_scale:int=0
|
||||
fixed_frame_rate_flag:bool=False
|
||||
nal_hrd_parameters_present_flag:bool=False
|
||||
hrd_parameters:HRD|None=None
|
||||
vcl_hrd_parameters_present_flag:bool=False
|
||||
vcl_hrd_parameters:HRD|None=None
|
||||
low_delay_hrd_flag:bool=False
|
||||
pic_struct_present_flag:bool=False
|
||||
bitstream_restriction_flag:bool=False
|
||||
motion_vectors_over_pic_boundaries_flag:bool=False
|
||||
max_bytes_per_pic_denom:int=0
|
||||
max_bits_per_mb_denom:int=0
|
||||
log2_max_mv_length_horizontal:int=0
|
||||
log2_max_mv_length_vertical:int=0
|
||||
max_num_reorder_frames:int=0
|
||||
max_dec_frame_buffering:int=0
|
||||
|
||||
# This structure is not guaranteed to be located at a byte boundary.
|
||||
# We must explicitely indicate bit offset.
|
||||
def from_bytes(self, buf, bit_position):
|
||||
bit_position, self.aspect_ratio_info_present_flag = read_boolean(buf, bit_position)
|
||||
if self.aspect_ratio_info_present_flag:
|
||||
bit_position, self.aspect_ratio_idc = read_byte(buf, bit_position)
|
||||
if self.aspect_ratio_idc == 255: # Extended_SAR
|
||||
bit_position, self.sar_width = read_word(buf, bit_position)
|
||||
bit_position, self.sar_height = read_word(buf, bit_position)
|
||||
bit_position, self.overscan_info_present_flag = read_boolean(buf, bit_position)
|
||||
if self.overscan_info_present_flag:
|
||||
bit_position, self.overscan_appropriate_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.video_signal_type_present_flag = read_boolean(buf, bit_position)
|
||||
if self.video_signal_type_present_flag:
|
||||
bit_position, self.video_format = read_bits(buf, bit_position, 3)
|
||||
bit_position, self.video_full_range_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.colour_description_present_flag = read_boolean(buf, bit_position)
|
||||
if self.colour_description_present_flag:
|
||||
bit_position, self.colour_primaries = read_byte(buf, bit_position)
|
||||
bit_position, self.transfer_characteristics = read_byte(buf, bit_position)
|
||||
bit_position, self.matrix_coefficients = read_byte(buf, bit_position)
|
||||
bit_position, self.chroma_loc_info_present_flag = read_boolean(buf, bit_position)
|
||||
if self.chroma_loc_info_present_flag:
|
||||
bit_position, self.chroma_sample_loc_type_top_field =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.chroma_sample_loc_type_bottom_field =\
|
||||
read_unsigned_exp_golomb(buf,bit_position)
|
||||
bit_position, self.timing_info_present_flag = read_boolean(buf, bit_position)
|
||||
if self.timing_info_present_flag:
|
||||
bit_position, self.num_units_in_tick = read_long(buf, bit_position)
|
||||
bit_position, self.time_scale = read_long(buf, bit_position)
|
||||
bit_position, self.fixed_frame_rate_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.nal_hrd_parameters_present_flag = read_boolean(buf, bit_position)
|
||||
if self.nal_hrd_parameters_present_flag:
|
||||
hrd = HRD()
|
||||
bit_position = hrd.from_bytes(buf, bit_position)
|
||||
self.hrd_parameters = hrd
|
||||
bit_position, self.vcl_hrd_parameters_present_flag = read_boolean(buf, bit_position)
|
||||
if self.vcl_hrd_parameters_present_flag:
|
||||
hrd = HRD()
|
||||
bit_position = hrd.from_bytes(buf, bit_position)
|
||||
self.vcl_hrd_parameters = hrd
|
||||
if self.nal_hrd_parameters_present_flag or self.vcl_hrd_parameters_present_flag:
|
||||
bit_position, self.low_delay_hrd_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.pic_struct_present_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.bitstream_restriction_flag = read_boolean(buf, bit_position)
|
||||
if self.bitstream_restriction_flag:
|
||||
bit_position, self.motion_vectors_over_pic_boundaries_flag =\
|
||||
read_boolean(buf, bit_position)
|
||||
bit_position, self.max_bytes_per_pic_denom = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.max_bits_per_mb_denom = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.log2_max_mv_length_horizontal = read_unsigned_exp_golomb(buf,
|
||||
bit_position)
|
||||
bit_position, self.log2_max_mv_length_vertical = read_unsigned_exp_golomb(buf,
|
||||
bit_position)
|
||||
bit_position, self.max_num_reorder_frames = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.max_dec_frame_buffering = read_unsigned_exp_golomb(buf, bit_position)
|
||||
|
||||
return bit_position
|
||||
|
||||
def to_bytes(self, buf, bit_position):
|
||||
bit_position = write_boolean(buf, bit_position, self.aspect_ratio_info_present_flag)
|
||||
if self.aspect_ratio_info_present_flag:
|
||||
bit_position = write_byte(buf, bit_position, self.aspect_ratio_idc)
|
||||
if self.aspect_ratio_idc == 255: # Extended_SAR
|
||||
bit_position = write_word(buf, bit_position, self.sar_width)
|
||||
bit_position = write_word(buf, bit_position, self.sar_height)
|
||||
bit_position = write_boolean(buf, bit_position, self.overscan_info_present_flag)
|
||||
if self.overscan_info_present_flag:
|
||||
bit_position = write_boolean(buf, bit_position, self.overscan_appropriate_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.video_signal_type_present_flag)
|
||||
if self.video_signal_type_present_flag:
|
||||
bit_position = write_bits(buf, bit_position, self.video_format, 3)
|
||||
bit_position = write_boolean(buf, bit_position, self.video_full_range_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.colour_description_present_flag)
|
||||
if self.colour_description_present_flag:
|
||||
bit_position = write_byte(buf, bit_position, self.colour_primaries)
|
||||
bit_position = write_byte(buf, bit_position, self.transfer_characteristics)
|
||||
bit_position = write_byte(buf, bit_position, self.matrix_coefficients)
|
||||
bit_position = write_boolean(buf, bit_position, self.chroma_loc_info_present_flag)
|
||||
if self.chroma_loc_info_present_flag:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.chroma_sample_loc_type_top_field)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.chroma_sample_loc_type_bottom_field)
|
||||
bit_position = write_boolean(buf, bit_position, self.timing_info_present_flag )
|
||||
if self.timing_info_present_flag:
|
||||
bit_position = write_long(buf, bit_position, self.num_units_in_tick )
|
||||
bit_position = write_long(buf, bit_position, self.time_scale)
|
||||
bit_position = write_boolean(buf, bit_position, self.fixed_frame_rate_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.nal_hrd_parameters_present_flag)
|
||||
if self.nal_hrd_parameters_present_flag:
|
||||
bit_position = self.hrd_parameters.to_bytes(buf, bit_position)
|
||||
bit_position = write_boolean(buf, bit_position, self.vcl_hrd_parameters_present_flag)
|
||||
if self.vcl_hrd_parameters_present_flag:
|
||||
bit_position = self.vcl_hrd_parameters.to_bytes(buf, bit_position)
|
||||
if self.nal_hrd_parameters_present_flag or self.vcl_hrd_parameters_present_flag:
|
||||
bit_position = write_boolean(buf, bit_position, self.low_delay_hrd_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.pic_struct_present_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.bitstream_restriction_flag)
|
||||
if self.bitstream_restriction_flag:
|
||||
bit_position = write_boolean(buf, bit_position,
|
||||
self.motion_vectors_over_pic_boundaries_flag)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.max_bytes_per_pic_denom)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_bits_per_mb_denom)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.log2_max_mv_length_horizontal)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.log2_max_mv_length_vertical)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_num_reorder_frames)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.max_dec_frame_buffering)
|
||||
|
||||
return bit_position
|
||||
|
||||
@dataclass
|
||||
class SPS:
|
||||
profile_idc:int=0 # u(8)
|
||||
constraint_set0_flag:bool=False # u(1)
|
||||
constraint_set1_flag:bool=False # u(1)
|
||||
constraint_set2_flag:bool=False # u(1)
|
||||
constraint_set3_flag:bool=False # u(1)
|
||||
constraint_set4_flag:bool=False # u(1)
|
||||
constraint_set5_flag:bool=False # u(1)
|
||||
level_idc:int=0 # u(8)
|
||||
seq_parameter_set_id:int=0 # ue(v)
|
||||
chroma_format_idc:int=0 # ue(v)
|
||||
separate_colour_plane_flag:bool=False # u(1)
|
||||
bit_depth_luma_minus8:int=0 # ue(v)
|
||||
bit_depth_chroma_minus8:int=0 # ue(v)
|
||||
qpprime_y_zero_transform_bypass_flag:bool=False # u(1)
|
||||
seq_scaling_matrix_present_flag:bool=False # u(1)
|
||||
scaling_list: dict = field(default_factory=dict)
|
||||
log2_max_frame_num_minus4:int=0 # ue(v)
|
||||
pic_order_cnt_type:int=0 # ue(v)
|
||||
log2_max_pic_order_cnt_lsb_minus4:int=0 # ue(v)
|
||||
delta_pic_order_always_zero_flag:bool=False # ue(1)
|
||||
offset_for_non_ref_pic:int=0 # se(v)
|
||||
offset_for_top_to_bottom_field:int=0 # se(v)
|
||||
num_ref_frames_in_pic_order_cnt_cycle:int=0 # ue(v)
|
||||
offset_for_ref_frame:dict[int,int] = field(default_factory=dict)
|
||||
max_num_ref_frames:int=9 # ue(v)
|
||||
gaps_in_frame_num_value_allowed_flag:bool=False # u(1)
|
||||
pic_width_in_mbs_minus1:int=0 # ue(v)
|
||||
pic_height_in_map_units_minus1:int=0 # ue(v)
|
||||
frame_mbs_only_flag:bool=False # u(1)
|
||||
mb_adaptive_frame_field_flag:bool=False # u(1)
|
||||
direct_8x8_inference_flag:bool=False # u(1)
|
||||
frame_cropping_flag:bool=False # u(1)
|
||||
frame_crop_left_offset:int=0 # ue(v)
|
||||
frame_crop_right_offset:int=0 # ue(v)
|
||||
frame_crop_top_offset:int=0 # ue(v)
|
||||
frame_crop_bottom_offset:int=0 # ue(v)
|
||||
vui_parameters_present_flag:bool=False # u(1)
|
||||
vui:VUI|None=None # VUI object
|
||||
|
||||
def __init__(self):
|
||||
self.scaling_list={}
|
||||
self.offset_for_ref_frame={}
|
||||
|
||||
# TODO: ...
|
||||
# Compute options to pass to ffmpeg so as to reproduce the same SPS.
|
||||
# Very complex since some codec configuration are not provided by ffmpeg and/or libx264.
|
||||
# This is only an attempt for now and it is almost impossible to mimic any profile without
|
||||
# patching ffmpeg and/or libx264 to add the support for corner cases.
|
||||
def ffmpeg_options(self, video_id=0):
|
||||
x264opts = []
|
||||
|
||||
try:
|
||||
profile = {0x42:'baseline', 0x4D:'main', 0x64:'high', 0x6E:'high10', 0x7A:'high422',
|
||||
0xF4:'high444'}[self.profile_idc]
|
||||
except KeyError:
|
||||
logger.error('Unknow profile: %x', self.profile_idc)
|
||||
return []
|
||||
|
||||
level = f'{floor(self.level_idc/10):d}.{self.level_idc % 10:d}'
|
||||
x264opts.extend([f'sps-id={self.seq_parameter_set_id:d}'] )
|
||||
|
||||
if self.bit_depth_chroma_minus8 not in [0,1,2,4,6,8]:
|
||||
logger.error('Bit depth of chrominance is not supported: %d',
|
||||
self.bit_depth_chroma_minus8+8)
|
||||
return []
|
||||
|
||||
if self.chroma_format_idc in range(4):
|
||||
if self.chroma_format_idc == 0:
|
||||
# Monochrome
|
||||
pass
|
||||
elif self.chroma_format_idc == 1:
|
||||
# YUV:4:2:0
|
||||
pass
|
||||
elif self.chroma_format_idc == 2:
|
||||
# YUV:4:2:2
|
||||
pass
|
||||
elif self.chroma_format_idc == 3:
|
||||
# YUV:4:4:4
|
||||
pass
|
||||
else:
|
||||
logger.error('Unknow chrominance format: %x', self.chroma_format_idc)
|
||||
return []
|
||||
|
||||
res = [f'-profile:v:{video_id:d}', self.profile_idc, f'-level:v:{video_id:d}', level]
|
||||
return res
|
||||
|
||||
def from_bytes(self, buf):
|
||||
logger.debug('Parsing: %s', hexdump.dump(buf,sep=':'))
|
||||
|
||||
bit_position=0
|
||||
|
||||
# NAL Unit SPS
|
||||
bit_position, zero = read_bit(buf, bit_position)
|
||||
if zero != 0:
|
||||
raise ValueError(f'Reserved bit is not equal to 0: {zero:d}')
|
||||
bit_position, nal_ref_idc = read_bits(buf, bit_position,2)
|
||||
if nal_ref_idc != 3:
|
||||
raise ValueError(f'NAL ref idc is not equal to 3: {nal_ref_idc:d}')
|
||||
bit_position, nal_unit_type = read_bits(buf, bit_position,5)
|
||||
if nal_unit_type != 7:
|
||||
raise ValueError(f'NAL unit type is not a SPS: {nal_unit_type:d}')
|
||||
|
||||
bit_position, self.profile_idc = read_byte(buf, bit_position)
|
||||
bit_position, self.constraint_set0_flag = read_bit(buf,bit_position)
|
||||
bit_position, self.constraint_set1_flag = read_bit(buf,bit_position)
|
||||
bit_position, self.constraint_set2_flag = read_bit(buf,bit_position)
|
||||
bit_position, self.constraint_set3_flag = read_bit(buf,bit_position)
|
||||
bit_position, self.constraint_set4_flag = read_bit(buf,bit_position)
|
||||
bit_position, self.constraint_set5_flag = read_bit(buf,bit_position)
|
||||
bit_position, v = read_bits(buf, bit_position, 2)
|
||||
if v!=0:
|
||||
raise ValueError(f'Reserved bits different from 0b00: {v:x}')
|
||||
bit_position, self.level_idc = read_byte(buf, bit_position)
|
||||
bit_position, self.seq_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position)
|
||||
if self.profile_idc in [44, 83, 86, 100, 110, 118, 122, 128, 134, 135, 138, 139, 244]:
|
||||
bit_position, self.chroma_format_idc = read_unsigned_exp_golomb(buf, bit_position)
|
||||
if self.chroma_format_idc==3:
|
||||
bit_position, self.separate_colour_plane_flag=read_bit(buf, bit_position)
|
||||
bit_position, self.bit_depth_luma_minus8 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.bit_depth_chroma_minus8 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.qpprime_y_zero_transform_bypass_flag = read_boolean(buf,
|
||||
bit_position)
|
||||
bit_position, self.seq_scaling_matrix_present_flag = read_boolean(buf, bit_position)
|
||||
if self.seq_scaling_matrix_present_flag:
|
||||
nb_matrices = 12 if self.chroma_format_idc == 3 else 8
|
||||
for i in range(nb_matrices):
|
||||
bit_position, present = read_boolean(buf, bit_position)
|
||||
if present:
|
||||
if i<6:
|
||||
bit_position, matrix = parse_scaling_list(buf, bit_position, 16)
|
||||
self.scaling_list[i] = matrix
|
||||
else:
|
||||
bit_position, matrix = parse_scaling_list(buf, bit_position, 64)
|
||||
self.scaling_list[i] = matrix
|
||||
else:
|
||||
self.scaling_list[i] = []
|
||||
|
||||
bit_position, self.log2_max_frame_num_minus4 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position , self.pic_order_cnt_type = read_unsigned_exp_golomb(buf, bit_position)
|
||||
if self.pic_order_cnt_type == 0:
|
||||
bit_position, self.log2_max_pic_order_cnt_lsb_minus4 =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
elif self.pic_order_cnt_type == 1:
|
||||
bit_position, self.delta_pic_order_always_zero_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.offset_for_non_ref_pic = read_signed_exp_golomb(buf, bit_position)
|
||||
bit_position, self.offset_for_top_to_bottom_field = read_signed_exp_golomb(buf,
|
||||
bit_position)
|
||||
bit_position, self.num_ref_frames_in_pic_order_cnt_cycle =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
for i in range(self.num_ref_frames_in_pic_order_cnt_cycle):
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.offset_for_ref_frame[i]=v
|
||||
|
||||
bit_position, self.max_num_ref_frames = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.gaps_in_frame_num_value_allowed_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.pic_width_in_mbs_minus1 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.pic_height_in_map_units_minus1 = read_unsigned_exp_golomb(buf,
|
||||
bit_position)
|
||||
bit_position, self.frame_mbs_only_flag = read_boolean(buf, bit_position)
|
||||
if not self.frame_mbs_only_flag:
|
||||
bit_position, self.mb_adaptive_frame_field_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.direct_8x8_inference_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.frame_cropping_flag = read_boolean(buf, bit_position)
|
||||
if self.frame_cropping_flag:
|
||||
bit_position, self.frame_crop_left_offset = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.frame_crop_right_offset = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.frame_crop_top_offset = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.frame_crop_bottom_offset = read_unsigned_exp_golomb(buf,
|
||||
bit_position)
|
||||
bit_position, self.vui_parameters_present_flag = read_boolean(buf, bit_position)
|
||||
|
||||
if self.vui_parameters_present_flag:
|
||||
self.vui = VUI()
|
||||
bit_position = self.vui.from_bytes(buf,bit_position)
|
||||
logger.debug('VUI present: %s', self.vui)
|
||||
|
||||
logger.debug('Parse end of SPS. Bit position: %d. Remaining bytes: %s.', bit_position,
|
||||
hexdump.dump(buf[floor(bit_position/8):], sep=':'))
|
||||
bit_position = parse_rbsp_trailing_bits(buf, bit_position)
|
||||
logger.debug('End of SPS: %d. Remaining bytes: %s', bit_position,
|
||||
hexdump.dump(buf[floor(bit_position/8):], sep=':'))
|
||||
return bit_position
|
||||
|
||||
def to_bytes(self):
|
||||
buf = bytearray()
|
||||
bit_position = 0
|
||||
bit_position = write_bit(buf, bit_position,0)
|
||||
bit_position = write_bits(buf, bit_position, 3, 2)
|
||||
bit_position = write_bits(buf, bit_position, 7, 5)
|
||||
bit_position = write_byte(buf, bit_position, self.profile_idc)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set0_flag)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set1_flag)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set2_flag)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set3_flag)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set4_flag)
|
||||
bit_position = write_bit(buf, bit_position, self.constraint_set5_flag)
|
||||
bit_position = write_bits(buf, bit_position, 0, 2)
|
||||
bit_position = write_byte(buf, bit_position, self.level_idc)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.seq_parameter_set_id)
|
||||
if self.profile_idc in [44, 83, 86, 100, 110, 118, 122, 128, 134, 135, 138, 139, 244]:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.chroma_format_idc)
|
||||
if self.chroma_format_idc==3:
|
||||
bit_position = write_bit(buf, bit_position, self.separate_colour_plane_flag)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.bit_depth_luma_minus8)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.bit_depth_chroma_minus8)
|
||||
bit_position = write_boolean(buf, bit_position,
|
||||
self.qpprime_y_zero_transform_bypass_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.seq_scaling_matrix_present_flag)
|
||||
if self.seq_scaling_matrix_present_flag:
|
||||
nb_matrices = 12 if self.chroma_format_idc == 3 else 8
|
||||
for i in range(nb_matrices):
|
||||
matrix = self.scaling_list[i]
|
||||
present = (len(matrix))!=0
|
||||
bit_position = write_boolean(buf, bit_position, present)
|
||||
if present:
|
||||
if i<6:
|
||||
bit_position = write_scaling_list(buf, bit_position, 16, matrix)
|
||||
else:
|
||||
bit_position = write_scaling_list(buf, bit_position, 64, matrix)
|
||||
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.log2_max_frame_num_minus4)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_order_cnt_type)
|
||||
if self.pic_order_cnt_type == 0:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.log2_max_pic_order_cnt_lsb_minus4)
|
||||
elif self.pic_order_cnt_type == 1:
|
||||
bit_position = write_boolean(buf, bit_position, self.delta_pic_order_always_zero_flag)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position, self.offset_for_non_ref_pic)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position,
|
||||
self.offset_for_top_to_bottom_field)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.num_ref_frames_in_pic_order_cnt_cycle)
|
||||
for i in range(self.num_ref_frames_in_pic_order_cnt_cycle):
|
||||
v = self.offset_for_ref_frame[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_num_ref_frames)
|
||||
bit_position = write_boolean(buf, bit_position, self.gaps_in_frame_num_value_allowed_flag)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_width_in_mbs_minus1)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.pic_height_in_map_units_minus1)
|
||||
bit_position = write_boolean(buf, bit_position, self.frame_mbs_only_flag)
|
||||
if not self.frame_mbs_only_flag:
|
||||
bit_position = write_boolean(buf, bit_position, self.mb_adaptive_frame_field_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.direct_8x8_inference_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.frame_cropping_flag)
|
||||
if self.frame_cropping_flag:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_left_offset)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.frame_crop_right_offset)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_top_offset)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.frame_crop_bottom_offset)
|
||||
bit_position = write_boolean(buf, bit_position, self.vui_parameters_present_flag)
|
||||
if self.vui_parameters_present_flag:
|
||||
logger.debug('SPS has VUI. Writing VUI at position: %d', bit_position)
|
||||
bit_position = self.vui.to_bytes(buf, bit_position)
|
||||
logger.debug('VUI written. New bit position: %d', bit_position)
|
||||
|
||||
bit_position = write_rbsp_trailing_bits(buf, bit_position)
|
||||
|
||||
return buf
|
||||
|
||||
@dataclass
|
||||
class PPS:
|
||||
pic_parameter_set_id:int=0
|
||||
seq_parameter_set_id:int=0
|
||||
entropy_coding_mode_flag:bool=False
|
||||
bottom_field_pic_order_in_frame_present_flag:bool=False
|
||||
num_slice_groups_minus1:int=0
|
||||
slice_group_map_type:int=0
|
||||
run_length_minus1:dict = field(default_factory=dict)
|
||||
top_left:dict = field(default_factory=dict)
|
||||
bottom_right:dict = field(default_factory=dict)
|
||||
slice_group_change_direction_flag:bool=False
|
||||
slice_group_change_rate_minus1:int=0
|
||||
pic_size_in_map_units_minus1:int=0
|
||||
slice_group_id:dict = field(default_factory=dict)
|
||||
num_ref_idx_l0_default_active_minus1:int=0
|
||||
num_ref_idx_l2_default_active_minus1:int=0
|
||||
weighted_pred_flag:bool=False
|
||||
weighted_bipred_idc:int=0
|
||||
pic_init_qp_minus26:int=0
|
||||
pic_init_qs_minus26:int=0
|
||||
chroma_qp_index_offset:int=0
|
||||
deblocking_filter_control_present_flag:bool=False
|
||||
constrained_intra_pred_flag:bool=False
|
||||
redundant_pic_cnt_present_flag:bool=False
|
||||
transform_8x8_mode_flag:bool=False
|
||||
pic_scaling_matrix_present_flag:bool=False
|
||||
pic_scaling_list:list[list[int]] = field(default_factory=list)
|
||||
second_chroma_qp_index_offset:int=0
|
||||
|
||||
def __init__(self):
|
||||
self.run_length_minus1={}
|
||||
self.top_left={}
|
||||
self.bottom_right={}
|
||||
self.slice_group_id={}
|
||||
self.pic_scaling_list=[]
|
||||
|
||||
# PPS are located at byte boundary
|
||||
def from_bytes(self, buf, chroma_format_idc):
|
||||
logger.debug('Parsing: %s', (hexdump.dump(buf,sep=':')))
|
||||
|
||||
bit_position=0
|
||||
# NAL Unit PPS
|
||||
bit_position, zero = read_bit(buf, bit_position)
|
||||
if zero != 0:
|
||||
raise ValueError(f'Reserved bit is not equal to 0: {zero:d}')
|
||||
bit_position, nal_ref_idc = read_bits(buf, bit_position,2)
|
||||
if nal_ref_idc != 3:
|
||||
raise ValueError(f'NAL ref idc is not equal to 3: {nal_ref_idc:d}')
|
||||
bit_position, nal_unit_type = read_bits(buf, bit_position,5)
|
||||
if nal_unit_type != 8:
|
||||
raise ValueError(f'NAL unit type is not a PPS: {nal_unit_type:d}')
|
||||
|
||||
bit_position, self.pic_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.seq_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.entropy_coding_mode_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.bottom_field_pic_order_in_frame_present_flag =\
|
||||
read_boolean(buf, bit_position)
|
||||
bit_position, self.num_slice_groups_minus1 = read_unsigned_exp_golomb(buf, bit_position)
|
||||
|
||||
if self.num_slice_groups_minus1>0:
|
||||
bit_position, self.slice_group_map_type = read_unsigned_exp_golomb(buf, bit_position)
|
||||
if self.slice_group_map_type == 0:
|
||||
for i in range(self.num_slice_groups_minus1):
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.run_length_minus1[i]=v
|
||||
elif self.slice_group_map_type == 2:
|
||||
for i in range(self.num_slice_groups_minus1):
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.top_left[i] = v
|
||||
bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
|
||||
self.bottom_right[i] = v
|
||||
elif self.slice_group_map_type in [3,4,5]:
|
||||
bit_position, self.slice_group_change_direction_flag = read_boolean(buf,
|
||||
bit_position)
|
||||
bit_position, self.slice_group_change_rate_minus1 =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
elif self.slice_group_map_type == 6:
|
||||
bit_position, self.pic_size_in_map_units_minus1 =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
nb_bits = ceil(log(self.num_slice_groups_minus1+1))
|
||||
for i in range(self.pic_size_in_map_units_minus1):
|
||||
bit_position, v = read_bits(buf, bit_position, nb_bits)
|
||||
self.slice_group_id[i]=v
|
||||
|
||||
bit_position, self.num_ref_idx_l0_default_active_minus1 =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.num_ref_idx_l2_default_active_minus1 =\
|
||||
read_unsigned_exp_golomb(buf, bit_position)
|
||||
bit_position, self.weighted_pred_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.weighted_bipred_idc = read_bits(buf, bit_position, 2)
|
||||
bit_position, self.pic_init_qp_minus26 = read_signed_exp_golomb(buf, bit_position)
|
||||
bit_position, self.pic_init_qs_minus26 = read_signed_exp_golomb(buf, bit_position)
|
||||
bit_position, self.chroma_qp_index_offset = read_signed_exp_golomb(buf, bit_position)
|
||||
bit_position, self.deblocking_filter_control_present_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.constrained_intra_pred_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.redundant_pic_cnt_present_flag = read_boolean(buf, bit_position)
|
||||
|
||||
if more_rbsp_data(buf, bit_position):
|
||||
bit_position, self.transform_8x8_mode_flag = read_boolean(buf, bit_position)
|
||||
bit_position, self.pic_scaling_matrix_present_flag = read_boolean(buf, bit_position)
|
||||
if self.pic_scaling_matrix_present_flag:
|
||||
nb_matrices = 6 if chroma_format_idc == 3 else 2
|
||||
if self.transform_8x8_mode_flag:
|
||||
nb_matrices+=6
|
||||
else:
|
||||
nb_matrices = 6
|
||||
for i in range(nb_matrices):
|
||||
bit_position, present = read_boolean(buf, bit_position)
|
||||
if present:
|
||||
if i<6:
|
||||
bit_position, matrix = parse_scaling_list(buf, bit_position, 16)
|
||||
self.pic_scaling_list.append(matrix)
|
||||
else:
|
||||
bit_position, matrix = parse_scaling_list(buf, bit_position, 64)
|
||||
self.pic_scaling_list.append(matrix)
|
||||
else:
|
||||
self.pic_scaling_list.append([])
|
||||
bit_position, self.second_chroma_qp_index_offset = read_signed_exp_golomb(buf,
|
||||
bit_position)
|
||||
|
||||
logger.info("parse RBSP")
|
||||
bit_position = parse_rbsp_trailing_bits(buf, bit_position)
|
||||
|
||||
return bit_position
|
||||
|
||||
def to_bytes(self, chroma_format_idc):
|
||||
buf = bytearray()
|
||||
bit_position = 0
|
||||
# NAL Unit PPS
|
||||
bit_position = write_bit(buf, bit_position, 0)
|
||||
bit_position = write_bits(buf, bit_position, 3, 2)
|
||||
bit_position = write_bits(buf, bit_position, 8, 5)
|
||||
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_parameter_set_id)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.seq_parameter_set_id)
|
||||
bit_position = write_boolean(buf, bit_position, self.entropy_coding_mode_flag)
|
||||
bit_position = write_boolean(buf, bit_position,\
|
||||
self.bottom_field_pic_order_in_frame_present_flag)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.num_slice_groups_minus1)
|
||||
|
||||
if self.num_slice_groups_minus1>0:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, self.slice_group_map_type)
|
||||
if self.slice_group_map_type == 0:
|
||||
for i in range(self.num_slice_groups_minus1):
|
||||
v = self.run_length_minus1[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
elif self.slice_group_map_type == 2:
|
||||
for i in range(self.num_slice_groups_minus1):
|
||||
v = self.top_left[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
v = self.bottom_right[i]
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position, v)
|
||||
elif self.slice_group_map_type in [3,4,5]:
|
||||
bit_position = write_boolean(buf, bit_position,
|
||||
self.slice_group_change_direction_flag)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.slice_group_change_rate_minus1)
|
||||
elif self.slice_group_map_type == 6:
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.pic_size_in_map_units_minus1)
|
||||
nb_bits = ceil(log(self.num_slice_groups_minus1+1))
|
||||
for i in range(self.pic_size_in_map_units_minus1):
|
||||
v = self.slice_group_id[i]
|
||||
bit_position, v = write_bits(buf, bit_position, v, nb_bits)
|
||||
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.num_ref_idx_l0_default_active_minus1)
|
||||
bit_position = write_unsigned_exp_golomb(buf, bit_position,
|
||||
self.num_ref_idx_l2_default_active_minus1)
|
||||
bit_position = write_boolean(buf, bit_position, self.weighted_pred_flag)
|
||||
bit_position = write_bits(buf, bit_position, self.weighted_bipred_idc, 2)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position, self.pic_init_qp_minus26)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position, self.pic_init_qs_minus26)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position, self.chroma_qp_index_offset)
|
||||
bit_position = write_boolean(buf, bit_position, self.deblocking_filter_control_present_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.constrained_intra_pred_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.redundant_pic_cnt_present_flag)
|
||||
|
||||
bit_position = write_boolean(buf, bit_position, self.transform_8x8_mode_flag)
|
||||
bit_position = write_boolean(buf, bit_position, self.pic_scaling_matrix_present_flag)
|
||||
if self.pic_scaling_matrix_present_flag:
|
||||
nb_matrices = 6 if chroma_format_idc == 3 else 2
|
||||
if self.transform_8x8_mode_flag:
|
||||
nb_matrices+=6
|
||||
else:
|
||||
nb_matrices = 6
|
||||
for i in range(nb_matrices):
|
||||
matrix = self.pic_scaling_list[i]
|
||||
logger.info("Retrieved pic scaling matrix: %s %d", matrix, len(matrix))
|
||||
present = len(matrix)!=0
|
||||
logger.info("Matrix is present: %s", present)
|
||||
bit_position = write_boolean(buf, bit_position, present)
|
||||
if present:
|
||||
if i<6:
|
||||
logger.info("Writing matrix: %s", matrix)
|
||||
bit_position = write_scaling_list(buf, bit_position, 16, matrix)
|
||||
else:
|
||||
logger.info("Writing matrix: %s", matrix)
|
||||
bit_position = write_scaling_list(buf, bit_position, 64, matrix)
|
||||
bit_position = write_signed_exp_golomb(buf, bit_position,
|
||||
self.second_chroma_qp_index_offset)
|
||||
|
||||
bit_position = write_rbsp_trailing_bits(buf, bit_position)
|
||||
|
||||
return buf
|
||||
@@ -0,0 +1,43 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
from typing import IO
|
||||
|
||||
import hexdump
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import InvalidMediaError
|
||||
from tscut.h264.avc import AVCDecoderConfiguration, get_avc_config_from_h264
|
||||
from tscut.matroska.ebml import get_ebml_length
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@typechecked
|
||||
def dump_codec_private_data(avc_decoder_configuration: AVCDecoderConfiguration) -> bytearray:
|
||||
# Rebuild a Matroska Codec Private Element
|
||||
res = bytearray()
|
||||
# Code private element
|
||||
res.extend(b'\x63\xA2')
|
||||
buf = avc_decoder_configuration.to_bytes()
|
||||
logger.debug('AVC configuration bitstream: %s (length: %d))', hexdump.dump(buf, sep=':'),
|
||||
len(buf))
|
||||
|
||||
embl_length = get_ebml_length(len(buf))
|
||||
if embl_length is None:
|
||||
raise InvalidMediaError
|
||||
logger.debug('EMBL encoded length: %s', hexdump.dump(embl_length, sep=':'))
|
||||
res.extend(embl_length)
|
||||
res.extend(buf)
|
||||
|
||||
return res
|
||||
|
||||
@typechecked
|
||||
def get_codec_private_data_from_h264(input_file: IO[bytes]) -> bytearray:
|
||||
avcconfig = get_avc_config_from_h264(input_file)
|
||||
res = dump_codec_private_data(avcconfig)
|
||||
|
||||
return res
|
||||
@@ -0,0 +1,165 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import logging
|
||||
from os import (
|
||||
SEEK_SET,
|
||||
fstat,
|
||||
ftruncate,
|
||||
lseek,
|
||||
read,
|
||||
write,
|
||||
)
|
||||
from typing import IO
|
||||
|
||||
import hexdump
|
||||
from typeguard import typechecked
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#MKV is formatted as an EBML file (Extended Binary Markup Langage).
|
||||
# cf http://matroska-org.github.io/libebml/specs.html
|
||||
# It is a Type, Length, Value (TLV) kind of binary file.
|
||||
# Types are encoded as follows:
|
||||
# 1xxx xxxx - Class A IDs (2^7 -1 possible values)
|
||||
# 01xx xxxx xxxx xxxx - Class B IDs (2^14-1 possible values)
|
||||
# 001x xxxx xxxx xxxx xxxx xxxx - Class C IDs (2^21-1 possible values)
|
||||
# 0001 xxxx xxxx xxxx xxxx xxxx xxxx xxxx - Class D IDs (2^28-1 possible values)
|
||||
# Lengths are encoded as follows:
|
||||
# 1xxx xxxx
|
||||
# value 0 to 2^7-2
|
||||
# 01xx xxxx xxxx xxxx
|
||||
# value 0 to 2^14-2
|
||||
# 001x xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^21-2
|
||||
# 0001 xxxx xxxx xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^28-2
|
||||
# 0000 1xxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^35-2
|
||||
# 0000 01xx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^42-2
|
||||
# 0000 001x xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^49-2
|
||||
# 0000 0001 xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx xxxx
|
||||
# value 0 to 2^56-2
|
||||
|
||||
@typechecked
|
||||
def get_ebml_length(length:int) -> bytes|None:
|
||||
if 0 <= length <= 2**7-2:
|
||||
size = 1
|
||||
elif length <= 2**14-2:
|
||||
size = 2
|
||||
elif length <= 2**21-2:
|
||||
size = 3
|
||||
elif length <= 2**28-2:
|
||||
size = 4
|
||||
elif length <= 2**35-2:
|
||||
size = 5
|
||||
elif length <= 2**42-2:
|
||||
size = 6
|
||||
elif length <= 2**49-2:
|
||||
size = 7
|
||||
elif length <= 2**56-2:
|
||||
size = 8
|
||||
elif length < 0:
|
||||
logger.error('Impossible to encode a negative length with EBML.')
|
||||
return None
|
||||
else:
|
||||
logger.error('Impossible to encode a length larger than 2^56-2 with EBML.')
|
||||
return None
|
||||
|
||||
encoded_length = length + ((128>>(size-1))<<((size-1)*8))
|
||||
res = (encoded_length).to_bytes(size, byteorder='big')
|
||||
return res
|
||||
|
||||
|
||||
@typechecked
|
||||
def change_ebml_element_size(input_file: IO[bytes], position:int, addendum:int) -> int:
|
||||
initial_position = position
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, position, SEEK_SET)
|
||||
|
||||
buf = read(infd, 1)
|
||||
element_type = int.from_bytes(buf, byteorder='big')
|
||||
mask=128
|
||||
found = False
|
||||
for i in range(1,5):
|
||||
if element_type&mask:
|
||||
type_size = i
|
||||
found = True
|
||||
break
|
||||
mask = mask>>1
|
||||
|
||||
if not found:
|
||||
logger.error('Size of element type cannot be determined: %d', element_type)
|
||||
raise ValueError()
|
||||
|
||||
# We seek to size
|
||||
position+=type_size
|
||||
lseek(infd, position, SEEK_SET)
|
||||
|
||||
buf = read(infd, 1)
|
||||
size_head = int.from_bytes(buf, byteorder='big')
|
||||
logger.info('First byte of size: %x', size_head)
|
||||
mask=128
|
||||
found = False
|
||||
for i in range(1,9):
|
||||
if size_head&mask:
|
||||
size_of_data_size = i
|
||||
found = True
|
||||
break
|
||||
mask = mask>>1
|
||||
|
||||
if not found:
|
||||
logger.error('Size of data size cannot be determined: %d', size_head)
|
||||
raise ValueError()
|
||||
logger.info('Size of data size: %d.', size_of_data_size)
|
||||
|
||||
lseek(infd, position, SEEK_SET)
|
||||
old_size_buf = read(infd, size_of_data_size)
|
||||
max_size = 2**(size_of_data_size*7)-2
|
||||
size_of_data = int.from_bytes(old_size_buf, byteorder='big')
|
||||
logger.info('Size of data with mask: %x mask: %d.', size_of_data, mask)
|
||||
size_of_data-= (mask<<((size_of_data_size-1)*8))
|
||||
logger.info('Found element at position: %d, size of type: %d size of data: %d \
|
||||
maximal size: %d.', initial_position, type_size, size_of_data, max_size)
|
||||
|
||||
new_size = size_of_data+addendum
|
||||
delta = 0
|
||||
if new_size > max_size:
|
||||
# TODO: Test this code ...
|
||||
new_encoded_size = get_ebml_length(new_size)
|
||||
if new_encoded_size is None:
|
||||
raise ValueError()
|
||||
size_of_new_encoded_size = len(new_encoded_size)
|
||||
if size_of_new_encoded_size <= size_of_data_size:
|
||||
logger.error('New encoded size is smaller (%d) or equal than previous size (%d).\
|
||||
This should not happen.', size_of_new_encoded_size, size_of_data_size)
|
||||
raise ValueError()
|
||||
# The difference of length between old size field and new one.
|
||||
delta = size_of_new_encoded_size - size_of_data_size
|
||||
file_length = fstat(infd).st_size
|
||||
# We seek after actual length field
|
||||
lseek(infd, position+size_of_data_size, SEEK_SET)
|
||||
# We read the rest of file
|
||||
tail = read(infd, file_length-(position+size_of_data_size))
|
||||
# We increase file length
|
||||
ftruncate(infd, file_length+delta)
|
||||
# We go to the beginning of length field
|
||||
lseek(infd, position, SEEK_SET)
|
||||
# We write the new length field
|
||||
write(infd, new_encoded_size)
|
||||
# We overwrite the rest of file with its previous content that has been offset.
|
||||
write(infd, tail)
|
||||
else:
|
||||
size = new_size + ((128>>(size_of_data_size-1))<<((size_of_data_size-1)*8))
|
||||
new_size_buf = (size).to_bytes(size_of_data_size, byteorder='big')
|
||||
|
||||
logger.info('Old encoded size: %s New encoded size: %s', hexdump.dump(old_size_buf,sep=':'),
|
||||
hexdump.dump(new_size_buf, sep=':'))
|
||||
lseek(infd, position, SEEK_SET)
|
||||
write(infd, new_size_buf)
|
||||
|
||||
# We return the potential increase in size of the file if the length field had to be increased.
|
||||
return delta
|
||||
@@ -0,0 +1,61 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from enum import IntEnum, unique
|
||||
from pathlib import Path
|
||||
from typing import IO, BinaryIO
|
||||
|
||||
from tscut.h264.avc import AVCDecoderConfiguration
|
||||
|
||||
|
||||
@unique
|
||||
class SupportedFormat(IntEnum):
|
||||
TS = 1
|
||||
MP4 = 2
|
||||
MATROSKA = 3
|
||||
|
||||
def __str__(self):
|
||||
match self:
|
||||
case SupportedFormat.TS:
|
||||
return 'mpegts'
|
||||
case SupportedFormat.MP4:
|
||||
return 'mov,mp4,m4a,3gp,3g2,mj2'
|
||||
case SupportedFormat.MATROSKA:
|
||||
return 'matroska,webm'
|
||||
case _:
|
||||
return 'Unsupported format'
|
||||
|
||||
@dataclass
|
||||
class ProcessingOptions:
|
||||
input_file: Path
|
||||
output_file: Path
|
||||
parts: list[tuple[timedelta, timedelta]]
|
||||
tools_paths: dict[str, str]
|
||||
all_optional_tools: bool
|
||||
framerate: int|None = None
|
||||
threshold: int = 0
|
||||
coarse: bool = False
|
||||
subtitles_ocr: bool = True
|
||||
verbose: bool = False
|
||||
dump_memory: bool = False
|
||||
keep_temporaries: bool = False
|
||||
|
||||
@dataclass
|
||||
class PreparedMedia:
|
||||
basename: str
|
||||
movie: BinaryIO
|
||||
duration: timedelta
|
||||
framerate: float
|
||||
streams: list
|
||||
width: int
|
||||
height: int
|
||||
avc_config: AVCDecoderConfiguration
|
||||
|
||||
@dataclass
|
||||
class CutResult:
|
||||
filename: str
|
||||
movie: IO[bytes]
|
||||
check_positions: list[timedelta]
|
||||
@@ -0,0 +1,253 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from shutil import copyfile, move
|
||||
from typing import BinaryIO
|
||||
|
||||
import hexdump
|
||||
|
||||
from tscut.cutting import cut_recording
|
||||
from tscut.exceptions import InvalidMediaError, TemporaryFileError
|
||||
from tscut.h264.avc import parse_codec_private
|
||||
from tscut.matroska.codec import dump_codec_private_data
|
||||
from tscut.models import PreparedMedia, ProcessingOptions, SupportedFormat
|
||||
from tscut.subtitles.ocr import do_ocr, extract_srt, get_tesseract_supported_lang
|
||||
from tscut.temporaries import TemporaryFiles
|
||||
from tscut.tools.ffmpeg import ffmpeg_convert
|
||||
from tscut.tools.ffprobe import (
|
||||
find_subtitles_tracks,
|
||||
get_format,
|
||||
get_frame_rate,
|
||||
get_movie_duration,
|
||||
get_streams,
|
||||
)
|
||||
from tscut.tools.mkvtoolnix import (
|
||||
get_codec_private_data_from_mkv,
|
||||
remux_srt_subtitles,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def detect_supported_format(format_name: str) -> SupportedFormat:
|
||||
for media_format in SupportedFormat:
|
||||
if format_name == str(media_format):
|
||||
return media_format
|
||||
|
||||
raise InvalidMediaError(
|
||||
f"Unsupported media format: {format_name}"
|
||||
)
|
||||
|
||||
def process_recording(options: ProcessingOptions, temporaries: TemporaryFiles) -> None:
|
||||
nb_parts = len(options.parts)
|
||||
basename = os.path.splitext(os.path.basename(options.input_file))[0]
|
||||
mp4filename = basename+'.mp4'
|
||||
mkvfilename = basename+'.mkv'
|
||||
|
||||
try:
|
||||
input_file = open(options.input_file, mode='rb')
|
||||
logger.debug("Type of input file: %s", type(input_file))
|
||||
except OSError as e:
|
||||
raise InvalidMediaError(f"Impossible to open {options.input_file}") from e
|
||||
|
||||
format_of_file = get_format(options.tools_paths['ffprobe'], input_file)
|
||||
|
||||
if format_of_file is None:
|
||||
raise InvalidMediaError(f"Impossible to retrieve format of file: {input_file}")
|
||||
if 'format_name' not in format_of_file:
|
||||
raise InvalidMediaError(f"Impossible to retrieve format of file: {input_file}")
|
||||
if 'duration' not in format_of_file:
|
||||
raise InvalidMediaError(f"Impossible to retrieve duration of file: {input_file}")
|
||||
format_name = format_of_file['format_name']
|
||||
|
||||
duration = timedelta(seconds=float(format_of_file['duration']))
|
||||
logger.info("Durée de l'enregistrement: %s", duration)
|
||||
|
||||
if options.framerate is None:
|
||||
framerate = get_frame_rate(options.tools_paths['ffprobe'], input_file)
|
||||
if framerate is None:
|
||||
raise InvalidMediaError('Impossible to estimate frame rate !')
|
||||
else:
|
||||
framerate = options.framerate
|
||||
|
||||
logger.info('Frame rate: %.1f fps', framerate)
|
||||
|
||||
final_format_of_file = detect_supported_format(format_name)
|
||||
|
||||
mkv: BinaryIO
|
||||
|
||||
if final_format_of_file == SupportedFormat.TS:
|
||||
logger.info("Converting TS to MP4 (to fix timestamps).")
|
||||
try:
|
||||
with open(mp4filename, 'wb+') as mp4:
|
||||
ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'],
|
||||
input_file, 'mpegts', mp4, 'mp4', duration)
|
||||
temporaries.add(mp4)
|
||||
logger.info("Converting MP4 to MKV.")
|
||||
try:
|
||||
mkv = open(mkvfilename, 'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create: {mkvfilename}") from e
|
||||
|
||||
ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'],
|
||||
mp4, 'mp4', mkv, 'matroska', duration)
|
||||
if nb_parts > 0:
|
||||
temporaries.add(mkv)
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create: {mp4filename}") from e
|
||||
|
||||
|
||||
elif final_format_of_file == SupportedFormat.MP4:
|
||||
logger.info("Converting MP4 to MKV")
|
||||
try:
|
||||
mkv = open(mkvfilename, 'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create: {mkvfilename}") from e
|
||||
ffmpeg_convert(options.tools_paths['ffmpeg'], options.tools_paths['ffprobe'],
|
||||
input_file, 'mp4', mkv, 'matroska', duration)
|
||||
if nb_parts > 0:
|
||||
temporaries.add(mkv)
|
||||
else:
|
||||
logger.info("Already in MKV")
|
||||
mkv = input_file
|
||||
|
||||
streams = get_streams(options.tools_paths['ffprobe'], mkv)
|
||||
if streams is None:
|
||||
raise InvalidMediaError(f"No streams found in file: {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:
|
||||
raise InvalidMediaError("Impossible to find main video stream.")
|
||||
|
||||
# We retrieve the main private codec data
|
||||
_, main_codec_private_data = get_codec_private_data_from_mkv(
|
||||
mkvinfo_path=options.tools_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:
|
||||
raise InvalidMediaError(f"AVC configurations are different:\
|
||||
{main_avc_config}\n{iso_avc_config}\n")
|
||||
|
||||
prepared_movie = PreparedMedia(
|
||||
basename = basename,
|
||||
movie = mkv,
|
||||
duration = duration,
|
||||
framerate = framerate,
|
||||
streams = streams,
|
||||
width = int(width),
|
||||
height = int(height),
|
||||
avc_config = main_avc_config
|
||||
)
|
||||
|
||||
cut_movie = cut_recording(prepared_movie, options, temporaries)
|
||||
|
||||
final_with_video_name = cut_movie.filename
|
||||
final_with_video = cut_movie.movie
|
||||
checks = cut_movie.check_positions
|
||||
|
||||
if options.subtitles_ocr:
|
||||
if not options.all_optional_tools:
|
||||
logger.warning("Missing tools for extracting subtitles.")
|
||||
move(final_with_video_name, options.output_file)
|
||||
else:
|
||||
# Final cut is not any more the final step.
|
||||
temporaries.add(final_with_video)
|
||||
duration = get_movie_duration(options.tools_paths['ffprobe'], final_with_video)
|
||||
supported_langs = get_tesseract_supported_lang(options.tools_paths['tesseract'])
|
||||
logger.info('Supported lang: %s', supported_langs)
|
||||
logger.info('Find subtitles tracks and language.')
|
||||
subtitles = find_subtitles_tracks(options.tools_paths['ffprobe'], final_with_video)
|
||||
logger.info(subtitles)
|
||||
sts: dict[str, list[int]] = {}
|
||||
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(options.tools_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 OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to open {idx_name}") from e
|
||||
try:
|
||||
sub = open(sub_name,'rb')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to open {sub_name}") from e
|
||||
|
||||
temporaries.add(idx)
|
||||
temporaries.add(sub)
|
||||
|
||||
ocr = do_ocr(options.tools_paths['vobsubocr'], list_of_subtitles, duration,
|
||||
temporaries, options.dump_memory)
|
||||
logger.info(ocr)
|
||||
|
||||
# Remux SRT subtitles
|
||||
remux_srt_subtitles(options.tools_paths['mkvmerge'], final_with_video,
|
||||
options.output_file, ocr)
|
||||
else:
|
||||
copyfile(final_with_video_name, options.output_file)
|
||||
else:
|
||||
move(final_with_video_name, options.output_file)
|
||||
|
||||
d = datetime(1,1,1)
|
||||
for c in checks:
|
||||
logger.info("Please check cut smoothness at %s", (c+d).strftime("%H:%M:%S"))
|
||||
@@ -0,0 +1,191 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from datetime import timedelta
|
||||
from io import TextIOWrapper
|
||||
from os import (
|
||||
SEEK_SET,
|
||||
fstat,
|
||||
lseek,
|
||||
memfd_create,
|
||||
read,
|
||||
write,
|
||||
)
|
||||
from subprocess import PIPE, Popen
|
||||
from typing import Sequence
|
||||
|
||||
from iso639 import Lang
|
||||
from iso639.exceptions import InvalidLanguageValue
|
||||
from tqdm import tqdm
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import ExternalToolError
|
||||
from tscut.temporaries import TemporaryFiles
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@typechecked
|
||||
def get_tesseract_supported_lang(tesseract_path:str) -> dict[Lang, str]:
|
||||
"""
|
||||
Retrieves the set of natural languages supported by the Tesseract OCR tool.
|
||||
|
||||
This function runs the Tesseract binary with the --list-langs option and parses the output
|
||||
to extract the supported languages.
|
||||
|
||||
Args:
|
||||
tesseract_path (str): The path to the Tesseract binary.
|
||||
|
||||
Returns:
|
||||
dict[Lang, str] | None:
|
||||
- A dictionary mapping Lang objects to their corresponding language codes
|
||||
(e.g., "eng" for English)
|
||||
- None if an error occurs while running the Tesseract binary
|
||||
"""
|
||||
res = {}
|
||||
|
||||
with Popen([tesseract_path, '--list-langs'], stdout=PIPE) as tesseract:
|
||||
assert tesseract.stdout is not None
|
||||
for line in tesseract.stdout:
|
||||
line = line.decode('utf8')
|
||||
p = re.compile('(?P<lang>[a-z]{3})\n')
|
||||
m = re.match(p,line)
|
||||
if m is not None:
|
||||
try:
|
||||
lang = m.group('lang')
|
||||
key = Lang(lang)
|
||||
res[key] = lang
|
||||
except InvalidLanguageValue as e:
|
||||
logger.warning('Invalid language: %s', e)
|
||||
|
||||
tesseract.wait()
|
||||
|
||||
if tesseract.returncode != 0:
|
||||
raise ExternalToolError("Tesseract returns an error code: %d",tesseract.returncode)
|
||||
|
||||
return res
|
||||
|
||||
|
||||
|
||||
@typechecked
|
||||
def extract_srt(mkvextract:str, filename:str, subtitles:dict[str, list[int]],
|
||||
langs:dict[Lang,str]) -> list[tuple[str,str,str,str]]:
|
||||
params = [mkvextract, filename, 'tracks']
|
||||
|
||||
res = []
|
||||
|
||||
for lang in subtitles:
|
||||
iso = Lang(lang)
|
||||
|
||||
if iso in langs:
|
||||
ocrlang = langs[iso]
|
||||
else:
|
||||
logger.warning("Language not supported by Tesseract: %s", iso.name)
|
||||
ocrlang ='osd'
|
||||
|
||||
if len(subtitles[lang]) == 1:
|
||||
params.append(f'{subtitles[lang][0]:d}:{lang}')
|
||||
res.append((f'{lang}.idx', f'{lang}.sub', lang, ocrlang))
|
||||
else:
|
||||
count = 1
|
||||
for track in subtitles[lang]:
|
||||
params.append(f'{track:d}:{lang}-{count:d}')
|
||||
res.append((f'{lang}-{count:d}.idx', f'{lang}-{count:d}.sub', lang, ocrlang))
|
||||
count = count+1
|
||||
|
||||
logger.debug('Executing %s', params)
|
||||
|
||||
env = {**os.environ, 'LANG': 'C'}
|
||||
with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract:
|
||||
assert extract.stdout is not None
|
||||
pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%',
|
||||
desc='Extraction:')
|
||||
for line in pb:
|
||||
if line.startswith('Progress :'):
|
||||
p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
pb.update(int(m['progress'])-pb.n)
|
||||
pb.update(100-pb.n)
|
||||
pb.refresh()
|
||||
pb.close()
|
||||
|
||||
extract.wait()
|
||||
|
||||
# mkvextract returns 0, 1 or 2 as error code.
|
||||
match extract.returncode:
|
||||
case 0:
|
||||
logger.info('Subtitle tracks were succesfully extracted.')
|
||||
case 1:
|
||||
logger.warning('Mkvextract returns warning')
|
||||
case 2:
|
||||
raise ExternalToolError('Mkvextract returns an error code: %d', extract.returncode)
|
||||
|
||||
return res
|
||||
|
||||
@typechecked
|
||||
def do_ocr(vobsubocr:str, idxs: Sequence[tuple[str,str,str,str]], duration:timedelta,
|
||||
temporaries:TemporaryFiles, dump_mem_fd:bool=False):
|
||||
res = []
|
||||
|
||||
for idx_name, _, lang, iso in idxs:
|
||||
srtname = f'{os.path.splitext(idx_name)[0]}.srt'
|
||||
# Tesseract seems to recognize the three dots ... as "su"
|
||||
ldots = re.compile('^su\n$')
|
||||
# Timestamps produced by vobsubocr: 01:52:19,861 --> 01:52:21,641
|
||||
timestamps = re.compile((r'^[0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3} \-\-> (?P<hours>[0-9]{2}):'
|
||||
r'(?P<minutes>[0-9]{2}):(?P<seconds>[0-9]{2}),[0-9]{3}$'))
|
||||
|
||||
srtfd = memfd_create(srtname, flags=0)
|
||||
with Popen([vobsubocr, '--lang', iso, idx_name], stdout=PIPE) as ocr:
|
||||
assert ocr.stdout is not None
|
||||
pb = tqdm(TextIOWrapper(ocr.stdout, encoding="utf-8"), total=
|
||||
int(duration/timedelta(seconds=1)), unit='s', desc='OCR')
|
||||
for line in pb:
|
||||
m = re.match(ldots,line)
|
||||
if m is not None:
|
||||
write(srtfd, '...'.encode(encoding='UTF-8'))
|
||||
else:
|
||||
write(srtfd, line.encode(encoding='UTF-8'))
|
||||
|
||||
m = re.match(timestamps, line)
|
||||
if m is not None:
|
||||
hours = int(m.group('hours'))
|
||||
minutes = int(m.group('minutes'))
|
||||
seconds = int(m.group('seconds'))
|
||||
ts = timedelta(hours=hours, minutes=minutes, seconds=seconds)
|
||||
pb.n = int(ts/timedelta(seconds=1))
|
||||
pb.update()
|
||||
|
||||
status = ocr.wait()
|
||||
|
||||
if status != 0:
|
||||
logger.error('OCR failed with status code: %d', status)
|
||||
|
||||
if dump_mem_fd:
|
||||
try:
|
||||
with open(srtname,'w', encoding='utf8') as dump_srt:
|
||||
lseek(srtfd, 0, SEEK_SET)
|
||||
srt_length = fstat(srtfd).st_size
|
||||
buf = read(srtfd, srt_length)
|
||||
outfd = dump_srt.fileno()
|
||||
pos = 0
|
||||
while pos < srt_length:
|
||||
pos+=write(outfd, buf[pos:])
|
||||
temporaries.add(dump_srt)
|
||||
except OSError:
|
||||
logger.error('Impossible to create file: %s', srtname)
|
||||
return None
|
||||
|
||||
|
||||
srt_length = fstat(srtfd).st_size
|
||||
if srt_length > 0:
|
||||
res.append((srtfd, lang))
|
||||
|
||||
return res
|
||||
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
"""A class to manage (and clean) all temporary files created during conversion"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import IO, Any, Self
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TemporaryFiles:
|
||||
def __init__(self, keep: bool = False) -> None:
|
||||
self._files: list[IO[Any]] = []
|
||||
self._keep = keep
|
||||
|
||||
def add(self, file: IO[Any]) -> None:
|
||||
self._files.append(file)
|
||||
|
||||
def cleanup(self) -> None:
|
||||
if self._keep:
|
||||
return
|
||||
|
||||
for f in self._files:
|
||||
path = os.path.realpath(f.name)
|
||||
logger.info("Removing: %s", path)
|
||||
try:
|
||||
f.close()
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
logger.exception("Unable to remove temporary file: %s", path)
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
return self
|
||||
|
||||
def __exit__(self) -> None:
|
||||
self.cleanup()
|
||||
@@ -0,0 +1,51 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
from shutil import which
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import MissingToolError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@typechecked
|
||||
def check_required_tools() -> tuple[bool,dict[str,str]]:
|
||||
"""
|
||||
Checks if all required external tools are installed.
|
||||
|
||||
This function verifies the presence of required and optional external tools on the system.
|
||||
It returns a tuple containing a boolean indicating whether all optional tools are installed,
|
||||
along with a dictionary containing the paths to all tools.
|
||||
|
||||
Args:
|
||||
None
|
||||
|
||||
Returns:
|
||||
tuple[bool, dict[str, str]]:
|
||||
- bool: True if all optional tools are installed, False otherwise
|
||||
- dict[str, str]: dictionary containing the paths to all tools
|
||||
"""
|
||||
all_optional_tools = True
|
||||
paths = {}
|
||||
required = ['ffmpeg', 'ffprobe', 'mkvmerge', 'mkvinfo']
|
||||
optional = ['mkvextract', 'vobsubocr','tesseract']
|
||||
for tool in required:
|
||||
path = which(tool)
|
||||
if path is None:
|
||||
logger.error('Required tool: %s is missing.',tool)
|
||||
raise MissingToolError(tool)
|
||||
paths[tool] = path
|
||||
for tool in optional:
|
||||
path = which(tool)
|
||||
if path is None:
|
||||
logger.info('Optional tool: %s is missing.',tool)
|
||||
all_optional_tools = False
|
||||
else:
|
||||
paths[tool] = path
|
||||
|
||||
return all_optional_tools, paths
|
||||
@@ -0,0 +1,388 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
from io import TextIOWrapper
|
||||
from math import ceil, floor, log
|
||||
from os import SEEK_SET, close, lseek, memfd_create, read, set_inheritable, write
|
||||
from subprocess import PIPE, Popen
|
||||
from typing import IO, BinaryIO
|
||||
|
||||
from tqdm import tqdm
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import ExternalToolError, InvalidMediaError, TemporaryFileError
|
||||
from tscut.temporaries import TemporaryFiles
|
||||
from tscut.tools.ffprobe import (
|
||||
get_frames_in_stream,
|
||||
get_video_dimensions,
|
||||
with_subtitles,
|
||||
)
|
||||
from tscut.tools.ppm import dump_ppm
|
||||
from tscut.tools.timeframe import get_packet_duration, parse_timestamp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@typechecked
|
||||
def ffmpeg_convert(ffmpeg_path:str, ffprobe_path:str, input_file: IO[bytes], input_format:str,
|
||||
output_file: IO[bytes], output_format:str, duration: timedelta) -> None:
|
||||
width, height = get_video_dimensions(ffprobe_path, input_file)
|
||||
if width is None or height is None:
|
||||
return
|
||||
|
||||
subtitles = with_subtitles(ffprobe_path, input_file)
|
||||
|
||||
infd = input_file.fileno()
|
||||
outfd = output_file.fileno()
|
||||
set_inheritable(infd, True)
|
||||
set_inheritable(outfd, True)
|
||||
|
||||
log_level = [] if logger.getEffectiveLevel() == logging.DEBUG else ['-loglevel', 'quiet']
|
||||
|
||||
params = [ffmpeg_path, '-y',]+log_level+['-progress', '/dev/stdout', '-canvas_size',
|
||||
f'{width:d}x{height:d}', '-f', input_format,
|
||||
'-i', f'/proc/self/fd/{infd:d}', '-map', '0:v',
|
||||
'-map', '0:a']
|
||||
if subtitles:
|
||||
params.extend(['-map', '0:s'])
|
||||
params.extend(['-bsf:v', 'h264_mp4toannexb,dump_extra=freq=keyframe', '-vcodec', 'copy',
|
||||
'-acodec', 'copy'])
|
||||
if subtitles:
|
||||
params.extend(['-scodec', 'dvdsub'])
|
||||
params.extend(['-r:0', '25', '-f', output_format, f'/proc/self/fd/{outfd:d}'])
|
||||
|
||||
logger.debug('Executing %s', params)
|
||||
|
||||
with Popen(params, stdout=PIPE, close_fds=False) as ffmpeg:
|
||||
assert ffmpeg.stdout is not None
|
||||
pb = tqdm(TextIOWrapper(ffmpeg.stdout, encoding="utf-8"),
|
||||
total=int(duration/timedelta(seconds=1)), unit='s', desc='Conversion')
|
||||
for line in pb:
|
||||
if line.startswith('out_time='):
|
||||
ts_str = line.split('=')[1].strip()
|
||||
ts = parse_timestamp(ts_str)
|
||||
if ts is not None:
|
||||
pb.n = int(ts/timedelta(seconds=1))
|
||||
pb.update()
|
||||
status = ffmpeg.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"Conversion failed with status code: {status:d}")
|
||||
|
||||
|
||||
@typechecked
|
||||
def extract_pictures(ffmpeg_path:str, input_file:IO[bytes], begin:timedelta, nb_frames:int,
|
||||
width:int=640, height:int=480) -> tuple[bytes,int]:
|
||||
"""
|
||||
Extract pictures from a video file using FFmpeg.
|
||||
|
||||
This function runs the FFmpeg binary to extract a specified number of frames from a video file,
|
||||
starting at a given time.
|
||||
The extracted frames are stored in memory as PPM images and returned as a tuple containing
|
||||
the image data and a file descriptor to the memory created by memfd_create.
|
||||
|
||||
Args:
|
||||
ffmpeg_path (str): The path to the FFmpeg binary.
|
||||
input_file (IO[bytes]): The input video file.
|
||||
begin (timedelta): The start time of the extraction.
|
||||
nb_frames (int): The number of frames to extract.
|
||||
width (int, optional): The width of the extracted images. Defaults to 640.
|
||||
height (int, optional): The height of the extracted images. Defaults to 480.
|
||||
|
||||
Returns:
|
||||
tuple[bytes, int] | tuple[None, None]:
|
||||
- A tuple containing the extracted image data as bytes and a file descriptor
|
||||
- A tuple containing None, None if the extraction fails
|
||||
"""
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
outfd = memfd_create('pictures', flags=0)
|
||||
set_inheritable(outfd, True)
|
||||
# PPM header
|
||||
# "P6\nWIDTH HEIGHT\n255\n"
|
||||
header_len=2+1+ceil(log(width, 10))+1+ceil(log(height, 10))+1+3+1
|
||||
logger.debug('Header length: %d', header_len)
|
||||
image_length = width*height*3+header_len
|
||||
length = image_length*nb_frames
|
||||
logger.debug("Estimated length: %d", length)
|
||||
|
||||
command = [ffmpeg_path, '-loglevel', 'quiet' ,'-y', '-ss', f'{begin}', '-i',
|
||||
f'/proc/self/fd/{infd}', '-s', f'{width:d}x{height:d}', '-vframes', f'{nb_frames:d}',
|
||||
'-c:v', 'ppm','-f', 'image2pipe', f'/proc/self/fd/{outfd:d}']
|
||||
logger.debug('Executing: %s', command)
|
||||
|
||||
images = b''
|
||||
with Popen(command, stdout=PIPE, close_fds=False) as ffmpeg:
|
||||
status = ffmpeg.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"Conversion failed with status code: {status:d}")
|
||||
|
||||
lseek(outfd, 0, SEEK_SET)
|
||||
images = read(outfd,length)
|
||||
if len(images) != length:
|
||||
raise InvalidMediaError(f"Received {len(images)} bytes but {length} were expected.")
|
||||
|
||||
lseek(outfd, 0, SEEK_SET)
|
||||
return images, outfd
|
||||
|
||||
@typechecked
|
||||
def extract_sound(ffmpeg_path:str, input_file: IO[bytes], begin:timedelta, output_filename:str,
|
||||
packet_duration:int, sub_channel:int=0,
|
||||
nb_packets:int=0, sample_rate:int=48000,
|
||||
nb_channels:int=2) -> tuple[bytes,int]:
|
||||
outfd = memfd_create(output_filename, flags=0)
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
set_inheritable(outfd, True)
|
||||
sound = b''
|
||||
length = int(nb_channels*sample_rate*4*nb_packets*packet_duration/1000)
|
||||
|
||||
command = [ffmpeg_path, '-y', '-loglevel', 'quiet', '-ss', f'{begin}',
|
||||
'-i', f'/proc/self/fd/{infd}', f'-frames:a:{sub_channel:d}', f'{nb_packets+1:d}',
|
||||
'-c:a', 'pcm_s32le', '-sample_rate', f'{sample_rate:d}',
|
||||
'-channels', f'{nb_channels:d}', '-f', 's32le', f'/proc/self/fd/{outfd:d}']
|
||||
logger.debug('Executing: %s', command)
|
||||
|
||||
with Popen(command, stdout=PIPE, close_fds=False) as ffmpeg:
|
||||
status = ffmpeg.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"Sound extraction returns error code: {status}")
|
||||
|
||||
lseek(outfd, 0, SEEK_SET)
|
||||
sound = read(outfd, length)
|
||||
|
||||
if len(sound) != length:
|
||||
raise InvalidMediaError(f"Received {len(sound)} bytes but {length} were expected (\
|
||||
channels={nb_channels}, freq={sample_rate} packets={nb_packets},\
|
||||
duration={packet_duration} ms).")
|
||||
|
||||
return sound, outfd
|
||||
|
||||
@typechecked
|
||||
def extract_all_streams(ffmpeg_path:str, ffprobe_path:str, input_file:IO[bytes], begin:timedelta,
|
||||
end:timedelta, streams, files_prefix, nb_frames:int, framerate:float,
|
||||
width:int, height:int, temporaries:TemporaryFiles,
|
||||
dump_mem_fd:bool=False) -> tuple[BinaryIO|None,
|
||||
TextIOWrapper|None,
|
||||
BinaryIO|None]:
|
||||
# The command line for encoding only video track
|
||||
video_encoder_params = [ ffmpeg_path, '-y', '-loglevel', 'quiet']
|
||||
video_input_params = []
|
||||
video_codec_params = []
|
||||
|
||||
# The command line to create a MKV file with the rest of tracks
|
||||
generic_encoder_params = [ ffmpeg_path, '-y', '-loglevel', 'quiet' ]
|
||||
generic_input_params = []
|
||||
generic_codec_params = []
|
||||
|
||||
if begin < end:
|
||||
video_id=0
|
||||
audio_id=0
|
||||
subtitle_id=0
|
||||
memfds = []
|
||||
for stream in streams:
|
||||
if stream['codec_type'] == 'video':
|
||||
logger.info("Extracting %d frames of video stream v:%d", nb_frames, video_id)
|
||||
sar = stream['sample_aspect_ratio']
|
||||
dar = stream['display_aspect_ratio']
|
||||
pixel_format = stream['pix_fmt']
|
||||
color_range = stream['color_range']
|
||||
color_space =stream['color_space']
|
||||
color_transfer = stream['color_transfer']
|
||||
color_primaries = stream['color_primaries']
|
||||
level_int = int(stream['level'])
|
||||
level = f'{floor(level_int/10):d}.{level_int%10:d}'
|
||||
chroma_location = stream['chroma_location']
|
||||
field_order = stream
|
||||
match field_order:
|
||||
case 'progressive':
|
||||
interlaced_options = ['-field_order', '0']
|
||||
case 'tt':
|
||||
interlaced_options = ['-top', '1', f'-flags:v:{video_id:d}', '+ilme+ildct',
|
||||
'-field_order', '1']
|
||||
case 'bb':
|
||||
interlaced_options = ['-top', '0', f'-flags:v:{video_id:d}', '+ilme+ildct',
|
||||
'-field_order','2']
|
||||
case 'tb':
|
||||
interlaced_options = ['-top', '1', f'-flags:v:{video_id:d}', '+ilme+ildct',
|
||||
'-field_order', '3']
|
||||
case 'bt':
|
||||
interlaced_options = ['-top', '0', f'-flags:v:{video_id:d}', '+ilme+ildct',
|
||||
'-field_order', '4']
|
||||
case _:
|
||||
interlaced_options = []
|
||||
|
||||
# ======================================= #
|
||||
# TODO: adjust SAR and DAR
|
||||
# https://superuser.com/questions/907933/correct-aspect-ratio-without-re-encoding-video-file
|
||||
# SAR: -aspect width:height
|
||||
# DAR: -bsf:v sample_aspect_ratio=1:video_format
|
||||
logger.warning('Missing SAR adjustment for: %s', sar)
|
||||
logger.warning('Missing DAR adjustment for: %s', dar)
|
||||
logger.warning('Missing treatment for chroma location: %s', chroma_location)
|
||||
codec = stream['codec_name']
|
||||
images_bytes, memfd = extract_pictures(ffmpeg_path, input_file=input_file,
|
||||
begin=begin, nb_frames=nb_frames,
|
||||
width=width, height=height)
|
||||
|
||||
memfds.append(memfd)
|
||||
if dump_mem_fd:
|
||||
dump_ppm(images_bytes, f'{files_prefix}-{video_id:d}', temporaries)
|
||||
|
||||
# We rewind to zero the memory file descriptor
|
||||
lseek(memfd, 0, SEEK_SET)
|
||||
set_inheritable(memfd, True)
|
||||
|
||||
video_input_params.extend(['-framerate', f'{framerate:f}', '-f', 'image2pipe', '-i',
|
||||
f'/proc/self/fd/{memfd:d}'])
|
||||
video_codec_params.extend([f'-c:v:{video_id:d}', codec, f'-level:v:{video_id:d}',
|
||||
level, '-pix_fmt', pixel_format])
|
||||
video_codec_params.extend(interlaced_options)
|
||||
video_codec_params.extend([f'-colorspace:v:{video_id:d}', color_space,
|
||||
f'-color_primaries:v:{video_id:d}', color_primaries,
|
||||
f'-color_trc:v:{video_id:d}', color_transfer,
|
||||
f'-color_range:v:{video_id:d}', color_range])
|
||||
video_id=video_id+1
|
||||
elif stream['codec_type'] == 'audio':
|
||||
logger.debug('Audio stream: %s', stream)
|
||||
sample_rate = int(stream['sample_rate'])
|
||||
nb_channels = int(stream['channels'])
|
||||
bit_rate = int(stream['bit_rate']) if 'bit_rate' in stream else 128000
|
||||
codec = stream['codec_name']
|
||||
if 'tags' in stream and 'language' in stream['tags']:
|
||||
generic_codec_params.extend([f'-metadata:s:a:{audio_id:d}',
|
||||
f"language={stream['tags']['language']}"])
|
||||
packets = get_frames_in_stream(ffprobe_path, input_file=input_file, begin=begin,
|
||||
end=end, stream_kind='a', sub_stream_id=audio_id)
|
||||
if packets is None:
|
||||
raise InvalidMediaError("Impossible to retrieve audio packets")
|
||||
nb_packets = len(packets)
|
||||
logger.debug("Found %d packets to be extracted from audio track.", nb_packets)
|
||||
if nb_packets > 0:
|
||||
packet_duration = get_packet_duration(packets[0])
|
||||
else:
|
||||
packet_duration = 0
|
||||
|
||||
logger.info("Extracting %d packets of audio stream: a:%d" , nb_packets, audio_id)
|
||||
tmpname = f'{files_prefix}-{audio_id:d}.pcm'
|
||||
|
||||
sound_bytes, memfd = extract_sound(ffmpeg_path=ffmpeg_path, input_file=input_file,
|
||||
begin=begin, nb_packets=nb_packets,
|
||||
packet_duration=packet_duration,
|
||||
output_filename=tmpname,
|
||||
sample_rate=sample_rate, nb_channels=nb_channels)
|
||||
|
||||
memfds.append(memfd)
|
||||
|
||||
if dump_mem_fd:
|
||||
try:
|
||||
with open(tmpname,'wb') as output:
|
||||
temporaries.add(output)
|
||||
outfd = output.fileno()
|
||||
pos = 0
|
||||
while pos < len(sound_bytes):
|
||||
pos+=write(outfd, sound_bytes[pos:])
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {tmpname}") from e
|
||||
|
||||
# We rewind to zero the memory file descriptor
|
||||
lseek(memfd, 0, SEEK_SET)
|
||||
set_inheritable(memfd, True)
|
||||
|
||||
generic_input_params.extend(['-f', 's32le', '-ar', f'{sample_rate:d}', '-ac',
|
||||
f'{nb_channels:d}', '-i', f'/proc/self/fd/{memfd:d}'])
|
||||
generic_codec_params.extend([f'-c:a:{audio_id:d}', codec, f'-b:a:{audio_id:d}',
|
||||
f'{bit_rate:d}'])
|
||||
audio_id=audio_id+1
|
||||
elif stream['codec_type'] == 'subtitle':
|
||||
logger.info("Extracting a subtitle stream: s:%d", subtitle_id)
|
||||
codec = stream['codec_name']
|
||||
generic_input_params.extend(['-i', './empty.idx'])
|
||||
if 'tags' in stream and 'language' in stream['tags']:
|
||||
generic_codec_params.extend([f'-metadata:s:s:{subtitle_id:d}',
|
||||
f"language={stream['tags']['language']}"])
|
||||
generic_codec_params.extend([f'-c:s:{subtitle_id:d}', 'copy'])
|
||||
subtitle_id=subtitle_id+1
|
||||
else:
|
||||
logger.error("Unknown stream type: %s", stream['codec_type'])
|
||||
|
||||
# Create a new MKV movie with all streams (except videos) that have been extracted.
|
||||
generic_encoder_params.extend(generic_input_params)
|
||||
|
||||
for index in range(audio_id+subtitle_id):
|
||||
generic_encoder_params.extend(['-map', f'{index:d}'])
|
||||
generic_encoder_params.extend(generic_codec_params)
|
||||
|
||||
mkv_filename = f'{files_prefix}.mkv'
|
||||
try:
|
||||
mkv_output = open(mkv_filename,'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {mkv_filename}") from e
|
||||
|
||||
mkvoutfd = mkv_output.fileno()
|
||||
set_inheritable(mkvoutfd, True)
|
||||
generic_encoder_params.extend(['-f', 'matroska', f'/proc/self/fd/{mkvoutfd:d}'])
|
||||
|
||||
logger.info('Encoding all streams (except video) into a MKV file: %s', mkv_filename)
|
||||
logger.debug('Executing: %s', generic_encoder_params)
|
||||
with Popen(generic_encoder_params, stdout=PIPE, close_fds=False) as ffmpeg:
|
||||
status = ffmpeg.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"Encoding failed with status code: {status}")
|
||||
|
||||
temporaries.add(mkv_output)
|
||||
|
||||
h264_filename = f'{files_prefix}.h264'
|
||||
try:
|
||||
h264_output = open(h264_filename,'wb+')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file {h264_filename}") from e
|
||||
|
||||
h264outfd = h264_output.fileno()
|
||||
set_inheritable(h264outfd, True)
|
||||
|
||||
video_encoder_params.extend(video_input_params)
|
||||
video_encoder_params.extend(video_codec_params)
|
||||
|
||||
video_encoder_params.extend([ '-x264opts', f'keyint=1:sps-id={1:d}','-bsf:v',
|
||||
'h264_mp4toannexb,dump_extra=freq=keyframe,h264_metadata=\
|
||||
overscan_appropriate_flag=1:sample_aspect_ratio=1:video_format=\
|
||||
0:chroma_sample_loc_type=0','-f', 'h264',
|
||||
f'/proc/self/fd/{h264outfd:d}'])
|
||||
|
||||
logger.info('Encoding video into a H264 file: %s', h264_filename)
|
||||
logger.debug('Executing: %s', video_encoder_params)
|
||||
with Popen(video_encoder_params, stdout=PIPE, close_fds=False) as ffmpeg:
|
||||
status = ffmpeg.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"Encoding failed with status code: {status:d}")
|
||||
|
||||
temporaries.add(h264_output)
|
||||
|
||||
h264_ts_filename = f'{files_prefix}-ts.txt'
|
||||
try:
|
||||
h264_ts_output = open(h264_ts_filename,'w+', encoding='utf8')
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file: {h264_ts_filename}") from e
|
||||
|
||||
h264_ts_output.write('# timestamp format v2\n')
|
||||
ts = 0
|
||||
for _ in range(nb_frames):
|
||||
ts = ts+ceil(1000/framerate)
|
||||
h264_ts_output.write(f'{ts:d}\n')
|
||||
h264_ts_output.flush()
|
||||
h264_ts_output.seek(0)
|
||||
|
||||
temporaries.add(h264_ts_output)
|
||||
|
||||
for memfd in memfds:
|
||||
close(memfd)
|
||||
|
||||
return h264_output, h264_ts_output, mkv_output
|
||||
|
||||
# Nothing to be done. We are already at a i-frame boundary.
|
||||
return None, None, None
|
||||
|
||||
@@ -0,0 +1,411 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from io import BytesIO
|
||||
from math import floor
|
||||
from os import (
|
||||
SEEK_SET,
|
||||
lseek,
|
||||
set_inheritable,
|
||||
)
|
||||
from subprocess import PIPE, Popen
|
||||
from typing import IO
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import InvalidMediaError, ExternalToolError
|
||||
from tscut.tools.timeframe import get_ts_frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@typechecked
|
||||
def get_frame_rate(ffprobe_path:str, input_file: IO[bytes]) -> float|None:
|
||||
"""
|
||||
Retrieves the frame rate of a video file using the ffprobe tool.
|
||||
|
||||
This function runs the ffprobe binary with the specified input file and parses the output
|
||||
to extract the frame rate.
|
||||
It uses two methods to calculate the frame rate: one based on the timestamp of the frames
|
||||
and another based on the duration of the frames.
|
||||
If the two calculated frame rates are significantly different, the function returns an error
|
||||
|
||||
Args:
|
||||
ffprobe_path (str): The path to the ffprobe binary.
|
||||
input_file (IO[bytes]): The input video file.
|
||||
|
||||
Returns:
|
||||
float | None:
|
||||
- The frame rate of the video file as a floating-point number
|
||||
- None if an error occurs while running the ffprobe binary or if the calculated
|
||||
frame rates are inconsistent
|
||||
"""
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
|
||||
mean_duration = 0.
|
||||
nb_frames1 = 0
|
||||
nb_frames2 = 0
|
||||
min_ts = None
|
||||
max_ts = None
|
||||
interlaced = False
|
||||
|
||||
params = [ffprobe_path, '-loglevel', 'quiet', '-select_streams', 'v', '-show_frames',
|
||||
'-read_intervals', '00%+30', '-of', 'json', f'/proc/self/fd/{infd:d}']
|
||||
env = {**os.environ, 'LANG': 'C'}
|
||||
with Popen(params, stdout=PIPE, close_fds=False, env=env) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'frames' in out:
|
||||
for frame in out['frames']:
|
||||
if 'interlaced_frame' in frame and frame['interlaced_frame'] == 1:
|
||||
interlaced = True
|
||||
if 'pts_time' in frame:
|
||||
ts = float(frame['pts_time'])
|
||||
if min_ts is None:
|
||||
min_ts = ts
|
||||
if max_ts is None:
|
||||
max_ts = ts
|
||||
min_ts = min(min_ts, ts)
|
||||
max_ts = max(max_ts, ts)
|
||||
nb_frames1+=1
|
||||
if 'duration_time' in frame:
|
||||
mean_duration+=float(frame['duration_time'])
|
||||
nb_frames2+=1
|
||||
else:
|
||||
return None
|
||||
|
||||
ffprobe.wait()
|
||||
|
||||
if ffprobe.returncode != 0:
|
||||
logger.error("ffprobe returns an error code: %d", ffprobe.returncode)
|
||||
return None
|
||||
|
||||
if max_ts is None or min_ts is None:
|
||||
logger.error("Impossible to determine maximum or minimum timestamps.")
|
||||
return None
|
||||
|
||||
frame_rate1 = nb_frames1/(max_ts-min_ts)
|
||||
frame_rate2 = nb_frames2 / mean_duration
|
||||
|
||||
if abs(frame_rate1 - frame_rate2) > 0.2:
|
||||
if not interlaced:
|
||||
logger.error('Video is not interlaced and the disperancy between frame rates is too \
|
||||
big: %f / %f', frame_rate1, frame_rate2)
|
||||
return None
|
||||
if abs(frame_rate1*2 - frame_rate2) < 0.2:
|
||||
return frame_rate2/2
|
||||
logger.error('Video is interlaced and the disperancy between frame rates is too big:\
|
||||
%f / %f', frame_rate1, frame_rate2)
|
||||
return None
|
||||
|
||||
return frame_rate2
|
||||
|
||||
@typechecked
|
||||
def get_subtitles_tracks(ffprobe_path:str, mkv_path: str) -> dict[str,list[str]]|None:
|
||||
tracks={}
|
||||
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-select_streams', 's', '-show_entries',
|
||||
'stream=index,codec_name:stream_tags=language', '-of', 'json', mkv_path],
|
||||
stdout=PIPE) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'streams' in out:
|
||||
for stream in out['streams']:
|
||||
index = stream['index']
|
||||
codec = stream['codec']
|
||||
lang:str = stream['tags']['language']
|
||||
if codec == 'dvd_subtitle':
|
||||
if lang not in tracks:
|
||||
tracks[lang] = [index]
|
||||
else:
|
||||
current_langs = tracks[lang]
|
||||
current_langs.append(index)
|
||||
tracks[lang] = current_langs
|
||||
else:
|
||||
return None
|
||||
|
||||
ffprobe.wait()
|
||||
|
||||
if ffprobe.returncode != 0:
|
||||
logger.error("ffprobe returns an error code: %d", ffprobe.returncode)
|
||||
return None
|
||||
|
||||
return tracks
|
||||
|
||||
@typechecked
|
||||
def get_format(ffprobe_path:str, input_file: IO[bytes]) -> dict|None:
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-show_format', '-of', 'json', '-i',
|
||||
f'/proc/self/fd/{infd:d}'], stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'format' in out:
|
||||
return out['format']
|
||||
logger.error('Impossible to retrieve format of file')
|
||||
|
||||
return None
|
||||
|
||||
@typechecked
|
||||
def get_movie_duration(ffprobe_path:str, input_file: IO[bytes]) -> timedelta:
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-show_format', '-of', 'json', '-i',
|
||||
f'/proc/self/fd/{infd:d}'], stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'format' in out and 'duration' in out['format']:
|
||||
duration = floor(float(out['format']['duration']))
|
||||
return timedelta(seconds=duration)
|
||||
raise InvalidMediaError("Impossible to retrieve duration of movie")
|
||||
|
||||
# ffprobe -loglevel quiet -select_streams v:0 -show_entries stream=width,height -of json sample.ts
|
||||
@typechecked
|
||||
def get_video_dimensions(ffprobe_path:str, input_file: IO[bytes]) -> tuple[int,int]:
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-select_streams', 'v:0', '-show_entries',\
|
||||
'stream=width,height', '-of', 'json', '-i', f'/proc/self/fd/{infd:d}'],\
|
||||
stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'streams' in out:
|
||||
video = out['streams'][0]
|
||||
if ('width' in video) and ('height' in video):
|
||||
return int(video['width']), int(video['height'])
|
||||
|
||||
raise InvalidMediaError('Impossible to retrieve dimensions of video')
|
||||
|
||||
@typechecked
|
||||
def get_streams(ffprobe_path:str, input_file: IO[bytes]) -> list|None:
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-show_streams', '-of', 'json', '-i',
|
||||
f'/proc/self/fd/{infd:d}'], stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'streams' in out:
|
||||
return out['streams']
|
||||
logger.error('Impossible to retrieve streams inside file')
|
||||
|
||||
return None
|
||||
|
||||
@typechecked
|
||||
def with_subtitles(ffprobe_path:str, input_file: IO[bytes]) -> bool:
|
||||
"""
|
||||
Checks if a media file contains subtitles using the ffprobe tool.
|
||||
|
||||
This function runs the ffprobe binary with the specified input file and parses the output
|
||||
to determine if the file contains subtitles.
|
||||
It returns True if at least one subtitle stream is found, False otherwise.
|
||||
|
||||
Args:
|
||||
ffprobe_path (str): The path to the ffprobe binary.
|
||||
input_file (IO[bytes]): The input media file.
|
||||
|
||||
Returns:
|
||||
bool:
|
||||
- True if the media file contains at least one subtitle stream
|
||||
- False if:
|
||||
- the media file does not contain any subtitle streams
|
||||
- an error occurs while running the ffprobe binary
|
||||
- the streams information cannot be retrieved from the media file
|
||||
"""
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-show_streams', '-of', 'json', '-i',
|
||||
f'/proc/self/fd/{infd:d}'], stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'streams' in out:
|
||||
streams = out['streams']
|
||||
for stream in streams:
|
||||
if 'codec_type' in stream and stream['codec_type'] == 'subtitle':
|
||||
return True
|
||||
else:
|
||||
logger.error('Impossible to retrieve streams inside file')
|
||||
|
||||
return False
|
||||
|
||||
def find_subtitles_tracks(ffprobe_path:str, input_file: IO[bytes]) -> dict:
|
||||
infd = input_file.fileno()
|
||||
lseek(infd, 0, SEEK_SET)
|
||||
set_inheritable(infd, True)
|
||||
|
||||
command = [ffprobe_path, '-loglevel','quiet', '-i', f'/proc/self/fd/{infd:d}',
|
||||
'-select_streams', 's', '-show_entries', 'stream=index:stream_tags=language',
|
||||
'-of', 'json']
|
||||
logger.debug('Executing: %s', command)
|
||||
|
||||
with Popen(command, stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
out = json.load(BytesIO(out))
|
||||
if 'streams' in out:
|
||||
return out['streams']
|
||||
|
||||
ffprobe.wait()
|
||||
raise InvalidMediaError('Impossible to retrieve format of file')
|
||||
|
||||
@typechecked
|
||||
def get_frames_in_stream(ffprobe_path: str, input_file: IO[bytes], begin:timedelta, end:timedelta,
|
||||
stream_kind:str, sub_stream_id:int=0) -> list[dict]|None:
|
||||
infd = input_file.fileno()
|
||||
set_inheritable(infd, True)
|
||||
|
||||
command = [ffprobe_path, '-loglevel', 'quiet', '-read_intervals', f'{begin}%{end}',
|
||||
'-show_entries', 'frame', '-select_streams',
|
||||
f'{stream_kind}:{sub_stream_id:d}','-of', 'json', f'/proc/self/fd/{infd:d}']
|
||||
logger.debug('Executing: %s', command)
|
||||
|
||||
with Popen(command, stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
frames = json.load(BytesIO(out))
|
||||
status = ffprobe.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"ffprobe failed with status code: {status:d}")
|
||||
|
||||
# Sort frames by timestamp
|
||||
tmp = {}
|
||||
if 'frames' in frames:
|
||||
frames = frames['frames']
|
||||
for frame in frames:
|
||||
ts = get_ts_frame(frame)
|
||||
if begin <= ts <= end:
|
||||
tmp[ts]=frame
|
||||
|
||||
res = []
|
||||
for ts in sorted(tmp):
|
||||
res.append(tmp[ts])
|
||||
return res
|
||||
|
||||
logger.error('Impossible to retrieve frames inside file around [%s,%s]', begin, end)
|
||||
return None
|
||||
|
||||
# TODO: Finish implementation of this function and use it.
|
||||
@typechecked
|
||||
def get_nearest_idr_frame(ffprobe_path: str, input_file: IO[bytes], timestamp:timedelta,
|
||||
before: bool=True, delta: timedelta=timedelta(seconds=2)):
|
||||
zero = timedelta()
|
||||
tbegin = timestamp-delta
|
||||
tend = timestamp+delta
|
||||
tbegin = max(tbegin, zero)
|
||||
|
||||
infd = input_file.fileno()
|
||||
set_inheritable(infd, True)
|
||||
|
||||
logger.debug('Looking for IDR frame in [%s, %s]', tbegin, tend)
|
||||
|
||||
idrs = []
|
||||
|
||||
# Retains only IDR frame
|
||||
with Popen([ffprobe_path, '-loglevel', 'quiet', '-read_intervals', f'{tbegin}%{tend}',
|
||||
'-skip_frame', 'nokey', '-show_entries', 'frame', '-select_streams', 'v:0',
|
||||
'-of', 'json', f'/proc/self/fd/{infd:d}'], stdout=PIPE, close_fds=False) as ffprobe:
|
||||
out, _ = ffprobe.communicate()
|
||||
frames = json.load(BytesIO(out))
|
||||
status = ffprobe.wait()
|
||||
if status != 0:
|
||||
raise ExternalToolError(f"ffprobe failed with status code: {status:d}")
|
||||
|
||||
if 'frames' in frames:
|
||||
frames = frames['frames']
|
||||
for frame in frames:
|
||||
ts = get_ts_frame(frame)
|
||||
if tbegin <= ts <= tend:
|
||||
idrs.append(frame)
|
||||
else:
|
||||
logger.error('Impossible to retrieve IDR frames inside file around [%s,%s]',
|
||||
tbegin, tend)
|
||||
return
|
||||
|
||||
return
|
||||
|
||||
@typechecked
|
||||
def get_nearest_iframe(ffprobe_path:str, input_file: IO[bytes],
|
||||
timestamp:timedelta, before:bool=True,
|
||||
delta_max:timedelta=timedelta(seconds=15))-> tuple[int,dict] | None:
|
||||
infd = input_file.fileno()
|
||||
set_inheritable(infd, True)
|
||||
|
||||
delta = timedelta(seconds=1)
|
||||
|
||||
iframe = None
|
||||
while delta < delta_max:
|
||||
zero = timedelta()
|
||||
tbegin = timestamp - delta if before else timestamp
|
||||
tend = timestamp + delta if not before else timestamp
|
||||
tbegin = max(tbegin, zero)
|
||||
logger.debug('Looking for an iframe in [%s, %s]', tbegin, tend)
|
||||
|
||||
frames = get_frames_in_stream(ffprobe_path, input_file=input_file, begin=tbegin, end=tend,
|
||||
stream_kind='v')
|
||||
if frames is None:
|
||||
logger.debug('Found no frame in [%s, %s]', tbegin, tend)
|
||||
delta+=timedelta(seconds=1)
|
||||
continue
|
||||
|
||||
iframes = []
|
||||
for frame in frames:
|
||||
if frame['pict_type'] == 'I':
|
||||
iframes.append(frame)
|
||||
|
||||
found = False
|
||||
for frame in iframes:
|
||||
try:
|
||||
ts = get_ts_frame(frame)
|
||||
except InvalidMediaError:
|
||||
logger.warning('I-frame with no timestamp: %s', frame)
|
||||
continue
|
||||
|
||||
if before and ts <= timestamp:
|
||||
found = True
|
||||
iframe = frame
|
||||
if not before and ts >= timestamp:
|
||||
found = True
|
||||
iframe = frame
|
||||
break
|
||||
|
||||
if found:
|
||||
logger.info("Found i-frame at: %s", iframe)
|
||||
break
|
||||
|
||||
delta+=timedelta(seconds=1)
|
||||
|
||||
if iframe is not None and frames is not None:
|
||||
its = get_ts_frame(iframe)
|
||||
if its is None:
|
||||
logger.error("Impossible to retrieve timestamp of i-frame !")
|
||||
return None
|
||||
nb_frames = 0
|
||||
for frame in frames:
|
||||
try:
|
||||
ts = get_ts_frame(frame)
|
||||
except InvalidMediaError:
|
||||
logger.warning('Frame without timestamp: %s', frame)
|
||||
continue
|
||||
|
||||
if before:
|
||||
if its <= ts <= timestamp:
|
||||
logger.info("Retrieve a frame between %s and %s at %s", its, timestamp, ts)
|
||||
nb_frames = nb_frames+1
|
||||
else:
|
||||
if timestamp <= ts <= its:
|
||||
logger.info("Retrieve a frame between %s and %s at %s", timestamp, ts, its)
|
||||
nb_frames = nb_frames+1
|
||||
else:
|
||||
logger.error("Impossible to find I-frame between: %s and %s", tbegin, tend)
|
||||
return None
|
||||
|
||||
return(nb_frames, iframe)
|
||||
@@ -0,0 +1,467 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from datetime import timedelta
|
||||
from io import TextIOWrapper
|
||||
from os import SEEK_SET, fstat, ftruncate, lseek, read, set_inheritable, write
|
||||
from pathlib import Path
|
||||
from subprocess import PIPE, Popen
|
||||
from typing import IO, Sequence
|
||||
|
||||
import hexdump
|
||||
from tqdm import tqdm
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import InvalidMediaError, ExternalToolError
|
||||
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]:
|
||||
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<hexdump>([0-9a-f]{2} )+)at (?P<position>[0-9]+) size (?P<size>[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
|
||||
|
||||
raise InvalidMediaError("Impossible to retrieve codec private data from MKV !")
|
||||
|
||||
@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<root>\+)|(\|(?P<depth>[ ]*\+))).*at (?P<position>[0-9]+)"
|
||||
r" size (?P<size>[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 | bytearray) -> 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:
|
||||
raise InvalidMediaError("Impossible to retrieve private data from 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:
|
||||
raise InvalidMediaError("Impossible to retrieve the key of codec private data")
|
||||
|
||||
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:
|
||||
assert mkvmerge.stdout is not None
|
||||
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<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
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:
|
||||
raise ExternalToolError("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: Sequence[IO[bytes]], output_name:str,
|
||||
concatenate: bool=True,
|
||||
timestamps: dict[int, IO[str]] | None = None) -> IO[bytes]:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if timestamps is None:
|
||||
timestamps = {}
|
||||
|
||||
fds = []
|
||||
try:
|
||||
out = open(output_name, 'wb+')
|
||||
except OSError as e:
|
||||
logger.error('Impossible to create file: %s', output_name)
|
||||
raise e
|
||||
|
||||
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:
|
||||
assert mkvmerge.stdout is not None
|
||||
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<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
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:
|
||||
assert extract.stdout is not None
|
||||
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<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
pb.update(int(m['progress'])-pb.n)
|
||||
pb.update(100-pb.n)
|
||||
pb.refresh()
|
||||
pb.close()
|
||||
|
||||
extract.wait()
|
||||
|
||||
if extract.returncode != 0:
|
||||
raise ExternalToolError(f"Mkvextract returns an error code: {extract.returncode:d}")
|
||||
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:
|
||||
assert remove.stdout is not None
|
||||
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<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
pb.update(int(m['progress'])-pb.n)
|
||||
pb.update(100-pb.n)
|
||||
pb.refresh()
|
||||
pb.close()
|
||||
|
||||
remove.wait()
|
||||
|
||||
if remove.returncode != 0:
|
||||
raise ExternalToolError(f"Mkvmerge returns an error code: {remove.returncode:d}")
|
||||
else:
|
||||
logger.info('Video tracks were succesfully extracted.')
|
||||
|
||||
|
||||
@typechecked
|
||||
def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_file: Path,
|
||||
subtitles) -> None:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
out = open(output_file, 'w', encoding='utf8')
|
||||
except OSError:
|
||||
logger.error('Impossible to create file: %s', output_file)
|
||||
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:
|
||||
assert mkvmerge.stdout is not None
|
||||
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<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
else:
|
||||
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
|
||||
@@ -0,0 +1,86 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
import re
|
||||
from io import BytesIO
|
||||
from math import ceil, log
|
||||
from os import write
|
||||
from typing import IO
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import TemporaryFileError
|
||||
from tscut.temporaries import TemporaryFiles
|
||||
|
||||
|
||||
@typechecked
|
||||
def dump_ppm(pictures: bytes, prefix: str, temporaries: TemporaryFiles) -> None:
|
||||
"""
|
||||
Dump PPM pictures from a bytes buffer to files.
|
||||
|
||||
This function takes a bytes buffer containing PPM pictures, a prefix for the output file names,
|
||||
and a list of temporary files.
|
||||
It extracts each PPM picture from the buffer, checks its validity, and writes it to a file.
|
||||
The output files are named according to the prefix and a zero-padded three-digit number.
|
||||
|
||||
Args:
|
||||
pictures (bytes): The bytes buffer containing the PPM pictures.
|
||||
prefix (str): The prefix for the output file names.
|
||||
temporaries (list[IO[bytes]]): A list of temporary files that will be used to store
|
||||
the output files.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
None, but logs errors if:
|
||||
- the PPM picture is not valid (e.g. wrong magic number, dimensions, or color encoding)
|
||||
- an I/O error occurs while creating or writing to an output file
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# "P6\nWIDTH HEIGHT\n255\n"
|
||||
pos = 0
|
||||
picture = 0
|
||||
|
||||
logger.debug('Dumping %d pictures: %s', len(pictures),prefix)
|
||||
|
||||
while pos<len(pictures):
|
||||
filename = f'{prefix}-{picture:03d}.ppm'
|
||||
header = BytesIO(pictures[pos:])
|
||||
magic = header.readline().decode('utf8')
|
||||
dimensions = header.readline().decode('utf8')
|
||||
max_value = int(header.readline().decode('utf8'))
|
||||
if magic == 'P6\n':
|
||||
pattern = re.compile('^(?P<width>[0-9]+) (?P<height>[0-9]+)\n$')
|
||||
m = pattern.match(dimensions)
|
||||
if m is not None:
|
||||
width = int(m['width'])
|
||||
height = int(m['height'])
|
||||
else:
|
||||
logger.error('Impossible to parse dimensions of picture')
|
||||
return
|
||||
else:
|
||||
logger.error('Not a PPM picture')
|
||||
return
|
||||
|
||||
if max_value != 255:
|
||||
logger.error('Not a valid PPM picture. Color are not encoded on byte. Max value: %d',
|
||||
max_value)
|
||||
|
||||
header_len=2+1+ceil(log(width, 10))+1+ceil(log(height, 10))+1+3+1
|
||||
try:
|
||||
with open(filename, 'wb') as out:
|
||||
temporaries.add(out)
|
||||
outfd = out.fileno()
|
||||
length=header_len+3*width*height
|
||||
nb_bytes = 0
|
||||
while nb_bytes < length:
|
||||
nb_bytes+=write(outfd, pictures[pos+nb_bytes:pos+length])
|
||||
pos+=length
|
||||
picture+=1
|
||||
except OSError as e:
|
||||
raise TemporaryFileError(f"Impossible to create file {filename}") from e
|
||||
@@ -0,0 +1,213 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
import re
|
||||
from datetime import timedelta
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
from tscut.exceptions import InvalidMediaError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@typechecked
|
||||
def parse_timestamp(ts:str) -> timedelta|None:
|
||||
"""
|
||||
Parse a timestamp string into a timedelta object.
|
||||
|
||||
This function takes a string representing a timestamp in the format HH:MM:SS[.us] and returns
|
||||
a timedelta object representing the corresponding time interval.
|
||||
The timestamp string can have an optional microsecond component.
|
||||
|
||||
Args:
|
||||
ts (str): The timestamp string to parse.
|
||||
|
||||
Returns:
|
||||
timedelta | None:
|
||||
- A timedelta object representing the parsed timestamp
|
||||
- None if:
|
||||
- the timestamp string is not in the correct format
|
||||
- the timestamp values are out of range (e.g. hour > 23, minute > 59, etc.)
|
||||
"""
|
||||
ts_reg_exp = (r'^(?P<hour>[0-9]{1,2}):(?P<minute>[0-9]{1,2})'
|
||||
r':(?P<second>[0-9]{1,2})(\.(?P<us>[0-9]{1,6}))?$')
|
||||
p = re.compile(ts_reg_exp)
|
||||
m = p.match(ts)
|
||||
if m is None:
|
||||
logger.warning("Impossible to parse timestamp: %s", ts)
|
||||
return None
|
||||
|
||||
values = m.groupdict()
|
||||
hour = 0
|
||||
minute = 0
|
||||
second = 0
|
||||
us = 0
|
||||
if values['hour'] is not None:
|
||||
hour = int(values['hour'])
|
||||
if values['minute'] is not None:
|
||||
minute = int(values['minute'])
|
||||
if values['second'] is not None:
|
||||
second = int(values['second'])
|
||||
if values['us'] is not None:
|
||||
us = int(values['us'])
|
||||
|
||||
if hour < 0 or hour > 23:
|
||||
logger.error("hour must be in [0,24[")
|
||||
return None
|
||||
if minute < 0 or minute > 59:
|
||||
logger.error("minute must be in [0,60[")
|
||||
return None
|
||||
if second < 0 or second > 59:
|
||||
logger.error("second must be in [0,60[")
|
||||
return None
|
||||
if us < 0 or us > 1000000:
|
||||
logger.error("milliseconds must be in [0,1000000[")
|
||||
return None
|
||||
|
||||
return timedelta(hours=hour, minutes=minute, seconds=second, microseconds=us)
|
||||
|
||||
@typechecked
|
||||
def parse_time_interval(interval: str) -> tuple[timedelta, timedelta] | tuple[None, None]:
|
||||
"""
|
||||
Parse a time interval string into a tuple of two timedelta objects.
|
||||
|
||||
This function takes a string representing a time interval in the format
|
||||
HH:MM:SS[.ms]-HH:MM:SS[.ms] and returns a tuple of two timedelta objects representing
|
||||
the start and end times of the interval.
|
||||
The time interval string can have an optional millisecond component.
|
||||
|
||||
Args:
|
||||
interval (str): The time interval string to parse.
|
||||
|
||||
Returns:
|
||||
tuple[timedelta, timedelta] | None:
|
||||
- A tuple of two timedelta objects representing the start and end times of the interval
|
||||
- None if:
|
||||
- the time interval string is not in the correct format
|
||||
- the time values are out of range (e.g. hour > 23, minute > 59, etc.)
|
||||
- the end time is before the start time (non-monotonic interval)
|
||||
"""
|
||||
interval_reg_exp = (r'^(?P<hour1>[0-9]{1,2}):(?P<minute1>[0-9]{1,2}):(?P<second1>[0-9]{1,2})'
|
||||
r'(\.(?P<ms1>[0-9]{1,3}))?-(?P<hour2>[0-9]{1,2}):(?P<minute2>[0-9]{1,2})'
|
||||
r':(?P<second2>[0-9]{1,2})(\.(?P<ms2>[0-9]{1,3}))?$')
|
||||
p = re.compile(interval_reg_exp)
|
||||
m = p.match(interval)
|
||||
if m is None:
|
||||
logger.error("Impossible to parse time interval")
|
||||
return None, None
|
||||
|
||||
values = m.groupdict()
|
||||
hour1 = 0
|
||||
minute1 = 0
|
||||
second1 = 0
|
||||
ms1 = 0
|
||||
hour2 = 0
|
||||
minute2 = 0
|
||||
second2 = 0
|
||||
ms2 = 0
|
||||
if values['hour1'] is not None:
|
||||
hour1 = int(values['hour1'])
|
||||
if values['minute1'] is not None:
|
||||
minute1 = int(values['minute1'])
|
||||
if values['second1'] is not None:
|
||||
second1 = int(values['second1'])
|
||||
if values['ms1'] is not None:
|
||||
ms1 = int(values['ms1'])
|
||||
if values['hour2'] is not None:
|
||||
hour2 = int(values['hour2'])
|
||||
if values['minute2'] is not None:
|
||||
minute2 = int(values['minute2'])
|
||||
if values['second2'] is not None:
|
||||
second2 = int(values['second2'])
|
||||
if values['ms2'] is not None:
|
||||
ms2 = int(values['ms2'])
|
||||
|
||||
if hour1 < 0 or hour1 > 23:
|
||||
logger.error("hour must be in [0,24[")
|
||||
return None, None
|
||||
if minute1 < 0 or minute1 > 59:
|
||||
logger.error("minute must be in [0,60[")
|
||||
return None, None
|
||||
if second1 < 0 or second1 > 59:
|
||||
logger.error("second must be in [0,60[")
|
||||
return None, None
|
||||
if ms1 < 0 or ms1 > 1000:
|
||||
logger.error("milliseconds must be in [0,1000[")
|
||||
return None, None
|
||||
|
||||
if hour2 < 0 or hour2 > 23:
|
||||
logger.error("hour must be in [0,24[")
|
||||
return None, None
|
||||
if minute2 < 0 or minute2 > 59:
|
||||
logger.error("minute must be in [0,60[")
|
||||
return None, None
|
||||
if second2 < 0 or second2 > 59:
|
||||
logger.error("second must be in [0,60[")
|
||||
return None, None
|
||||
if ms2 < 0 or ms2 > 1000:
|
||||
logger.error("milliseconds must be in [0,1000[")
|
||||
return None, None
|
||||
|
||||
ts1 = timedelta(hours=hour1, minutes=minute1, seconds=second1, microseconds=ms1*1000)
|
||||
ts2 = timedelta(hours=hour2, minutes=minute2, seconds=second2, microseconds=ms2*1000)
|
||||
|
||||
if ts2 < ts1:
|
||||
logger.error("Non monotonic interval")
|
||||
return None,None
|
||||
|
||||
return (ts1, ts2)
|
||||
|
||||
@typechecked
|
||||
def compare_time_interval(interval1: tuple[timedelta, timedelta],
|
||||
interval2: tuple[timedelta, timedelta]) -> int:
|
||||
"""
|
||||
Compare two time intervals.
|
||||
|
||||
This function compares two time intervals represented by tuples of two timedelta objects.
|
||||
It returns an integer indicating the relationship between the two intervals:
|
||||
- -1 if interval 1 is before interval 2
|
||||
- 1 if interval 1 is after interval 2
|
||||
- 0 if the two intervals overlap or are equal
|
||||
|
||||
Args:
|
||||
interval1 (tuple[timedelta, timedelta]): The first time interval
|
||||
interval2 (tuple[timedelta, timedelta]): The second time interval
|
||||
|
||||
Returns:
|
||||
int: The relationship between the two time intervals
|
||||
"""
|
||||
ts11,ts12 = interval1
|
||||
ts21,ts22 = interval2
|
||||
|
||||
if ts12 < ts21:
|
||||
return -1
|
||||
if ts22 < ts11:
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
@typechecked
|
||||
def get_ts_frame(frame: dict) -> timedelta:
|
||||
if 'pts_time' in frame:
|
||||
pts_time = float(frame['pts_time'])
|
||||
elif 'pkt_pts_time' in frame:
|
||||
pts_time = float(frame['pkt_pts_time'])
|
||||
else:
|
||||
raise InvalidMediaError(f"Impossible to find timestamp of frame {frame}")
|
||||
|
||||
return timedelta(seconds=pts_time)
|
||||
|
||||
@typechecked
|
||||
def get_packet_duration(packet: dict) -> int:
|
||||
if 'duration' in packet:
|
||||
duration = int(packet['duration'])
|
||||
elif 'pkt_duration' in packet:
|
||||
duration = int(packet['pkt_duration'])
|
||||
else:
|
||||
raise InvalidMediaError("Impossible to find duration of packet {packet}")
|
||||
|
||||
return duration
|
||||
-3297
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,37 @@
|
||||
from tscut.h264.bitstream import (
|
||||
read_bit,
|
||||
read_byte,
|
||||
read_word
|
||||
)
|
||||
|
||||
def test_read_bit_1():
|
||||
buf = bytes(10)
|
||||
for p in range(80):
|
||||
pos, b = read_bit(buf, p)
|
||||
assert pos == p+1 and b == 0
|
||||
|
||||
def test_read_bit_2():
|
||||
buf = b'\xFF'*10
|
||||
for p in range(80):
|
||||
pos, b = read_bit(buf, p)
|
||||
assert pos == p+1 and b == 1
|
||||
|
||||
def test_read_byte_1():
|
||||
buf = bytes(10)
|
||||
for p in range(10):
|
||||
pos, b = read_byte(buf, p*8)
|
||||
assert pos == (p+1)*8 and b == 0
|
||||
|
||||
def test_read_byte_2():
|
||||
buf = bytes(10)
|
||||
buf = b'\xFF'*10
|
||||
for p in range(10):
|
||||
pos, b = read_byte(buf, p*8)
|
||||
assert pos == (p+1)*8 and b == 0xFF
|
||||
|
||||
def test_read_word_1():
|
||||
buf = bytes(20)
|
||||
for p in range(10):
|
||||
pos, b = read_word(buf, p*16)
|
||||
assert pos == (p+1)*16 and b == 0
|
||||
|
||||
Reference in New Issue
Block a user