diff --git a/removeads.py b/removeads.py index 2c02b24..8dfbe04 100755 --- a/removeads.py +++ b/removeads.py @@ -10,6 +10,7 @@ from os import mkdir, set_inheritable from io import BytesIO, TextIOWrapper import json from enum import Enum, IntEnum, unique, auto +import shutil @unique class SupportedFormat(IntEnum): @@ -256,7 +257,13 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid codecsParams.extend(['-c:a:%d' % audioID, codec, '-b:a:%d' % audioID, '%d' % bitRate]) audioID=audioID+1 elif stream['codec_type'] == 'subtitle': - # TODO: what can be done ? + print("Extracting a subtitle stream: %s" % stream) + codec = stream['codec_name'] + inputParams.extend(['-i', './empty.idx']) + if 'tags' in stream: + if 'language' in stream['tags']: + codecsParams.extend(['-metadata:s:s:%d' % subTitleID, 'language=%s' % stream['tags']['language']]) + codecsParams.extend(['-c:s:%d' % subTitleID, 'copy']) subTitleID=subTitleID+1 else: logger.info("Unknown stream type: %s" % stream['codec_type']) @@ -266,12 +273,13 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid # Create a new MKV movie with all streams that have been extracted. encoderParams.extend(inputParams) - for index in range(0,videoID+audioID+subTitleID-1): + for index in range(0,videoID+audioID+subTitleID): encoderParams.extend(['-map', '%d' % index]) encoderParams.extend(codecsParams) - # output = open('out.mkv','w') - # outfd = output.fileno() - encoderParams.extend(['-f', 'matroska', 'out.mkv']) + output = open('%s.mkv' % filesPrefix,'w') + outfd = output.fileno() + set_inheritable(outfd, True) + encoderParams.extend(['-top', '1', '-bsf:v', 'h264_mp4toannexb,dump_extra=keyframe', '-f', 'matroska', '/proc/self/fd/%d' % outfd]) print(encoderParams) @@ -279,11 +287,45 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"): print(line, end='') + return output else: # Nothing to be done. We are already at a i-frame boundary. - pass + return None + +# Merge a list of mkv files passed as input, and produce a new MKV as output +def mergeMKVs(inputs, outputName): + fds = [] + out = open(outputName, 'w') + # TODO: Check success or failure + outfd = out.fileno() + fds.append(outfd) + set_inheritable(outfd, True) + mergeParams = ['mkvmerge'] + first = True + for mkv in inputs: + if mkv !=None: + fd = mkv.fileno() + fds.append(fd) + set_inheritable(fd, True) + if first: + mergeParams.append('/proc/self/fd/%d' % fd) + first = False + else: + mergeParams.append('+/proc/self/fd/%d' % fd) + + mergeParams.extend(['-o', '/proc/self/fd/%d' % outfd]) + + # We merge all files. + with Popen(mergeParams, stdout=PIPE, close_fds=False) as mkvmerge: + for line in TextIOWrapper(mkvmerge.stdout, encoding="utf-8"): + print(line, end='') + + for fd in fds: + set_inheritable(fd, False) + + return out def parseTimeInterval(interval): logger = logging.getLogger(__name__) @@ -460,7 +502,17 @@ def main(): # Pour chaque portion partnum = 0 + mkvparts = [] + for ts1, ts2 in parts: + # 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. + # 'B/P' 'B/P'* 'I' 'I' 'B/P'* 'B/P' + # Si la trame de début est déjà 'I', il n'y a rien à faire (idem pour la fin). + # Sinon on extrait les trames 'B' ou 'P' depuis le début jusqu'à la trame 'I' non incluse + partnum = partnum + 1 headFrames = getNearestIFrame(mkv, ts1, before=False) @@ -480,25 +532,32 @@ def main(): headIFrameTS = timedelta(seconds=float(headIFrame['pts_time'])) tailIFrameTS = timedelta(seconds=float(tailIFrame['pts_time'])) - extractAllStreams(inputFile=mkv, begin=ts1, end=headIFrameTS, nbFrames=nbHeadFrames, filesPrefix='part-%d-head' % (partnum), streams=streams, width=width, height=height) - extractAllStreams(inputFile=mkv, begin=tailIFrameTS, end=ts2, nbFrames=nbTailFrames, filesPrefix='part-%d-tail' % (partnum), streams=streams, width=width, height=height) + + subparts = [] + + head = extractAllStreams(inputFile=mkv, begin=ts1, end=headIFrameTS, nbFrames=nbHeadFrames, filesPrefix='part-%d-head' % (partnum), streams=streams, width=width, height=height) + subparts.append(head) + + tail = extractAllStreams(inputFile=mkv, begin=tailIFrameTS, end=ts2, nbFrames=nbTailFrames, filesPrefix='part-%d-tail' % (partnum), streams=streams, width=width, height=height) # Creating MKV file that corresponds to current part between I-frames - with open('part-%d-internal.mkv' % partnum, 'w') as partmkv: - extractMKVPart(inputFile=mkv, outputFile=partmkv, begin=headIFrameTS, end=tailIFrameTS) + internal = open('part-%d-internal.mkv' % partnum, 'w') + # TODO: test if failure + extractMKVPart(inputFile=mkv, outputFile=internal, begin=headIFrameTS, end=tailIFrameTS) + subparts.append(internal) - # 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. - # 'B/P' 'B/P'* 'I' 'I' 'B/P'* 'B/P' - # Si la trame de début est déjà 'I', il n'y a rien à faire (idem pour la fin). - # Sinon on extrait les trames 'B' ou 'P' depuis le début jusqu'à la trame 'I' non incluse - # Fabriquer une courte vidéo au format MKV reprenant les mêmes codecs que la vidéo originale avec les fichiers extraits précedemment. - # mkvmerge() => création d'un fichier part-%d.mkv + subparts.append(tail) + part = mergeMKVs(inputs=subparts, outputName="part-%d.mkv" % partnum) + mkvparts.append(part) - # Appeler mkvmerge pour fusionner tous les fichiers part-%d.mkv - + nbParts = len(mkvparts) + if nbParts > 1: + mergeMKVs(inputs=mvkparts, outputName=args.output) + elif nbParts == 1: + print("A single part") + else: + print("Nothing produced !") + if __name__ == "__main__": main()