Public/removeads
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Separate main() function in a dedicated module.

Browse Source
This commit is contained in:
Frédéric Tronel
2026-08-28 16:20:17 +02:00
parent 7e72611b1c
commit 97ef0e972c
4 changed files with 570 additions and 532 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/tscut/tscut.py
-532
View File
@@ -3295,535 +3295,3 @@ def do_coarse_processing(ffmpeg_path:str, ffprobe_path:str, mkvmerge_path:str,
begin=begin, end=end)
temporaries.append(internal_mkv)
def main() -> None:
logger = logging.getLogger(__name__)
coloredlogs.install()
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", dest='input_file', type=str, required=True,
help="Input file to process (can be .ts, .mp4 or .mkv).")
parser.add_argument("-o", "--output", dest='output_file', type=str, required=True,
help="Output MKV file to produce.")
parser.add_argument("-p", "--part", dest='parts', nargs='+', required=False, action='append',
metavar="hh:mm:ss[.mmm]-hh:mm:ss[.mmm]",
help="Extract this exact part of the original file.")
parser.add_argument("-k", "--keep", action='store_true',
help="Do not cleanup temporary files after processing.")
parser.add_argument("-t", "--threshold", action='store', type=int,
help="Suppress headers and trailers that are smaller than the threshold.")
parser.add_argument("-c", "--coarse", action='store_true', dest='coarse',
help="Do not take trailers and headers into account.")
parser.add_argument("--dump-memory", action='store_true', dest='dump',
help="For debug purpose, dump all memory mapping of headers (and trailers)\
before (after) each part. They are kept in memory only otherwise.")
parser.add_argument("-s","--srt", action='store_true', dest='srt',
help="Dump subtitles and make OCR and finally remux them in the movie\
(as SRT).")
parser.add_argument("-v","--verbose", action='store_true', dest='verbose', help="Debug.")
parser.add_argument("-f","--framerate", action='store', type=int,
help="Override frame rate estimator.")
args = parser.parse_args()
logger.info('Arguments: %s', args)
if args.verbose:
logger.info('Setting logging to debug mode')
coloredlogs.set_level(level=logging.DEBUG)
logger.debug('Arguments: %s', args)
if args.coarse and args.threshold is not None:
logger.error('--coarse and threshold arguments are exclusive.')
exit(-1)
if (not args.coarse) and args.threshold is None:
args.threshold = 0
all_optional_tools, paths = check_required_tools()
# Flatten args.parts
intervals = []
if args.parts is not None:
for part in args.parts:
for subpart in part:
intervals.append(subpart)
parts=[]
# Parse each interval
for interval in intervals:
ts1, ts2 = parse_time_interval(interval)
if ts1 is None or ts2 is None:
logger.error("Illegal time interval: %s", interval)
exit(-1)
parts.append((ts1,ts2))
# Sort intervals
parts.sort(key=cmp_to_key(compare_time_interval))
# Check that no intervals are overlapping
prevts = timedelta(0)
for part in parts:
ts1, ts2 = part
if prevts > ts1:
logger.error('Intervals are overlapping')
exit(-1)
prevts = ts2
nb_parts = len(parts)
temporaries = []
basename = os.path.splitext(os.path.basename(args.input_file))[0]
mp4filename = basename+'.mp4'
mkvfilename = basename+'.mkv'
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")
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:
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"))
if __name__ == "__main__":
main()