diff --git a/removeads.py b/removeads.py index 9865bb8..2c02b24 100755 --- a/removeads.py +++ b/removeads.py @@ -25,8 +25,38 @@ class SupportedFormat(IntEnum): elif self is SupportedFormat.Matroska: return 'matroska,webm' else: - return "Unsupported format" + return 'Unsupported format' +@unique +class ColorSpace(IntEnum): + BT709=0, + FCC=1, + BT601=2, + BT470=3, + BT470BG=4, + SMPTE170M=5, + SMPTE240M=6, + BT2020=7 + + def __str__(self): + if self is ColorSpace.BT709: + return 'bt709' + elif self is ColorSpace.FCC: + return 'fcc' + elif self is ColorSpace.BT601: + return 'bt601' + elif self is ColorSpace.BT470: + return 'bt470' + elif self is ColorSpace.BT470BG: + return 'bt470bg' + elif self is ColorSpace.SMPTE170M: + return 'smpte170m' + elif self is ColorSpace.SMPTE240M: + return 'smpte240m' + elif self is ColorSpace.BT2020: + return 'bt2020' + else: + return 'Unsupported color space' def getFormat(inputFile): logger = logging.getLogger(__name__) @@ -75,6 +105,26 @@ def ffmpegConvert(inputFile, inputFormat, outputFile, outputFormat): if line.startswith('out_time='): print(line, end='') +def getFramesInStream(inputFile, begin, end, streamKind, subStreamId=0): + logger = logging.getLogger(__name__) + infd = inputFile.fileno() + set_inheritable(infd, True) + + with Popen(['ffprobe', '-loglevel', 'quiet', '-read_intervals', ('%s%%%s' %(begin, end)), '-show_entries', 'frame', '-select_streams', '%s:%d' % (streamKind, subStreamId), '-of', 'json', '/proc/self/fd/%d' % infd], stdout=PIPE, close_fds=False) as ffprobe: + out, _ = ffprobe.communicate() + frames = json.load(BytesIO(out)) + res = [] + if 'frames' in frames: + frames = frames['frames'] + for frame in frames: + ts = timedelta(seconds=float(frame['pts_time'])) + if begin <= ts and ts <= end: + res.append(frame) + return res + else: + logger.error('Impossible to retrieve frames inside file around [%s,%s]' % (begin, end)) + return None + def getNearestIFrame(inputFile, timestamp, before=True, delta=timedelta(seconds=2)): logger = logging.getLogger(__name__) @@ -89,48 +139,43 @@ def getNearestIFrame(inputFile, timestamp, before=True, delta=timedelta(seconds= logger.debug('Looking for iframe in [%s, %s]' % (tbegin, tend)) - with Popen(['ffprobe', '-loglevel', 'quiet', '-read_intervals', ('%s%%%s' %(tbegin, tend)), '-show_entries', 'frame', '-select_streams', 'v', '-of', 'json', '/proc/self/fd/%d' % infd], stdout=PIPE, close_fds=False) as ffprobe: - out, _ = ffprobe.communicate() - frames = json.load(BytesIO(out)) - if 'frames' in frames: - frames = frames['frames'] - iframes = [] - for frame in frames: - if frame['pict_type'] == 'I': - iframes.append(frame) - - found = False - res = None - for frame in iframes: - if before and timedelta(seconds=float(frame['pts_time'])) <= timestamp: - found = True - iframe = frame - if not before and timedelta(seconds=float(frame['pts_time'])) >= timestamp: - found = True - iframe = frame - break + frames = getFramesInStream(inputFile=inputFile, begin=tbegin, end=tend, streamKind='v') + if frames == None: + return None + + iframes = [] + for frame in frames: + if frame['pict_type'] == 'I': + iframes.append(frame) + + found = False + res = None + for frame in iframes: + if before and timedelta(seconds=float(frame['pts_time'])) <= timestamp: + found = True + iframe = frame + if not before and timedelta(seconds=float(frame['pts_time'])) >= timestamp: + found = True + iframe = frame + break - if found: - logger.debug("Found: %s" % res) - - its = timedelta(seconds=float(iframe['pts_time'])) - nbFrames = 0 - for frame in frames: - ts = timedelta(seconds=float(frame['pts_time'])) - if before: - if its <= ts and ts <= timestamp: - nbFrames = nbFrames+1 - else: - if timestamp <= ts and ts <= its: - nbFrames = nbFrames+1 + if found: + logger.debug("Found: %s" % res) + + its = timedelta(seconds=float(iframe['pts_time'])) + nbFrames = 0 + for frame in frames: + ts = timedelta(seconds=float(frame['pts_time'])) + if before: + if its <= ts and ts <= timestamp: + nbFrames = nbFrames+1 else: - logger.error("Impossible to find I-frame around: %s" % timestamp) + if timestamp <= ts and ts <= its: + nbFrames = nbFrames+1 + else: + logger.error("Impossible to find I-frame around: %s" % timestamp) - return(nbFrames-1, iframe) - else: - logger.error('Impossible to retrieve video frames inside file around [%s,%s]' % (tbegin, tend)) - - return None + return(nbFrames-1, iframe) def extractMKVPart(inputFile, outputFile, begin, end): inputFile.seek(0,0) @@ -147,19 +192,99 @@ def extractPictures(inputFile, begin, nbFrames, prefix, width=640, height=480): inputFile.seek(0,0) infd = inputFile.fileno() set_inheritable(infd, True) - with Popen(['ffmpeg', '-loglevel', 'quiet', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-s', '%dx%d'%(width, height), '-vframes', '%d'%nbFrames, '-c:v', 'ppm', '-f', 'image2', '%s-%%03d.ppm' % prefix], stdout=PIPE, close_fds=False) as ffmpeg: + with Popen(['ffmpeg', '-y', '-loglevel', 'quiet', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-s', '%dx%d'%(width, height), '-vframes', '%d'%nbFrames, '-c:v', 'ppm', '-f', 'image2', '%s-%%03d.ppm' % prefix], stdout=PIPE, close_fds=False) as ffmpeg: for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"): print(line, end='') -def extractSound(inputFile, begin, outputFile, channel=0, nbPackets=10, sampleRate=48000, nbChannels=2): +def extractSound(inputFile, begin, outputFile, subChannel=0, nbPackets=0, sampleRate=48000, nbChannels=2): inputFile.seek(0,0) + outputFile.seek(0,0) infd = inputFile.fileno() + outfd = outputFile.fileno() set_inheritable(infd, True) - with Popen(['ffmpeg', '-loglevel', 'quiet', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-frames:a:%d' % channel, '%d' % nbPackets, - '-c:a', 'pcm_s32le', '-sample_rate', '%d' % 48000, '-channels', '%d' % 2, '-f', 's32le', '%s.pcm' % outputFile], stdout=PIPE, close_fds=False) as ffmpeg: + set_inheritable(outfd, True) + with Popen(['ffmpeg', '-y', '-loglevel', 'quiet', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-frames:a:%d' % subChannel, '%d' % nbPackets, + '-c:a', 'pcm_s32le', '-sample_rate', '%d' % sampleRate, '-channels', '%d' % nbChannels, '-f', 's32le', '/proc/self/fd/%d' % outfd], stdout=PIPE, close_fds=False) as ffmpeg: for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"): print(line, end='') +def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, width, height): + logger = logging.getLogger(__name__) + # encoderParams = [ 'ffmpeg', '-y', '-loglevel', 'quiet' ] + encoderParams = [ 'ffmpeg', '-y' ] + inputParams = [] + codecsParams = [] + + if begin < end: + videoID=0 + audioID=0 + subTitleID=0 + audioFiles = {} + for stream in streams: + if stream['codec_type'] == 'video': + print("Extracting video stream: %s" % stream) + frameRate = stream['r_frame_rate'] + sar = stream['sample_aspect_ratio'] + dar = stream['display_aspect_ratio'] + pixelFormat = stream['pix_fmt'] + colorRange = stream['color_range'] + colorSpace =stream['color_space'] + colorTransfer = stream['color_transfer'] + colorPrimaries = stream['color_primaries'] + codec = stream['codec_name'] + extractPictures(inputFile=inputFile, begin=begin, nbFrames=nbFrames, prefix="%s-%d" % (filesPrefix, videoID), width=width, height=height) + inputParams.extend(['-i', '%s-%d-%%03d.ppm' % (filesPrefix, videoID)]) + codecsParams.extend(['-c:v:%d' % videoID, codec, '-pix_fmt', pixelFormat, '-colorspace:v:%d' % videoID, colorSpace, '-color_primaries:v:%d' % videoID, colorPrimaries, + '-color_trc:v:%d' % videoID, colorTransfer, '-color_range:v:%d' % videoID, colorRange ]) + videoID=videoID+1 + elif stream['codec_type'] == 'audio': + print("Extracting audio stream: %s" % stream) + sampleRate = int(stream['sample_rate']) + nbChannels = int(stream['channels']) + bitRate = int(stream['bit_rate']) + codec = stream['codec_name'] + if 'tags' in stream: + if 'language' in stream['tags']: + codecsParams.extend(['-metadata:s:a:%d' % audioID, 'language=%s' % stream['tags']['language']]) + packets = getFramesInStream(inputFile=inputFile, begin=begin, end=end, streamKind='a', subStreamId=audioID) + nbPackets = len(packets) + print("Found %d packets to be extracted from audio track." % nbPackets) + audioFiles[audioID] = open('%s-%d.pcm' % (filesPrefix,audioID), 'w') + # TODO: test if successfully openened + extractSound(inputFile=inputFile, begin=begin, nbPackets=nbPackets, outputFile=audioFiles[audioID], sampleRate=sampleRate, nbChannels=nbChannels) + inputParams.extend(['-f', 's32le', '-ar', '%d'%sampleRate, '-ac', '%d'%nbChannels, '-i', '/proc/self/fd/%d' % audioFiles[audioID].fileno()]) + 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 ? + subTitleID=subTitleID+1 + else: + logger.info("Unknown stream type: %s" % stream['codec_type']) + + # Example: + # ffmpeg -framerate 25.85 -i image-%02d.ppm -f s32le -ar 48000 -ac 2 -i ./audio-1.pcm -c:a eac3 -b:a 128k -c:v libx264 -crf 25.85 -vf "scale=1920:1080,format=yuv420p" -colorspace:v "bt709" -color_primaries:v "bt709" -color_trc:v "bt709" -color_range:v "tv" -top 1 -flags:v +ilme+ildct -bsf:v h264_mp4toannexb,dump_extra=keyframe -metadata MAJOR_BRAND=isom -metadata MINOR_VERSION=512 -movflags +faststart cut-1.mkv + + # Create a new MKV movie with all streams that have been extracted. + encoderParams.extend(inputParams) + for index in range(0,videoID+audioID+subTitleID-1): + encoderParams.extend(['-map', '%d' % index]) + encoderParams.extend(codecsParams) + # output = open('out.mkv','w') + # outfd = output.fileno() + encoderParams.extend(['-f', 'matroska', 'out.mkv']) + + print(encoderParams) + + with Popen(encoderParams, stdout=PIPE, close_fds=False) as ffmpeg: + for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"): + print(line, end='') + + + else: + # Nothing to be done. We are already at a i-frame boundary. + pass + + def parseTimeInterval(interval): logger = logging.getLogger(__name__) @@ -338,58 +463,29 @@ def main(): for ts1, ts2 in parts: partnum = partnum + 1 - preFrame = getNearestIFrame(mkv, ts1, before=False) - if preFrame == None: + headFrames = getNearestIFrame(mkv, ts1, before=False) + if headFrames == None: exit(-1) - postFrame = getNearestIFrame(mkv, ts2, before=True) - if postFrame == None: + tailFrames = getNearestIFrame(mkv, ts2, before=True) + if tailFrames == None: exit(-1) - nbPreFrame, preIFrame = preFrame - nbPostFrame, postIFrame = postFrame + nbHeadFrames, headIFrame = headFrames + nbTailFrames, tailIFrame = tailFrames - print("Found pre I-frame and %d frames between: %s" % (nbPreFrame, preIFrame)) - print("Found I-frame and %d frames between: %s" % (nbPostFrame, postIFrame)) + print("Found head I-frame and %d frames between: %s" % (nbHeadFrames, headIFrame)) + print("Found I-frame and %d frames between: %s" % (nbTailFrames, tailIFrame)) - preIFrameTS = timedelta(seconds=float(preIFrame['pts_time'])) - postIFrameTS = timedelta(seconds=float(postIFrame['pts_time'])) + headIFrameTS = timedelta(seconds=float(headIFrame['pts_time'])) + tailIFrameTS = timedelta(seconds=float(tailIFrame['pts_time'])) - if ts1 < preIFrameTS: - for stream in streams: - if stream['codec_type'] == 'video': - extractPictures(inputFile=mkv, begin=ts1, nbFrames=nbPreFrame, prefix="pre-part-%d" % partnum, width=width, height=height) - elif stream['codec_type'] == 'audio': - print("Extracting audio stream: %s" % stream) - sampleRate = stream['sample_rate'] - nbChannel = stream['channels'] - extractSound(inputFile=mkv, begin=ts1, nbPackets=nbPreFrame, outputFile="pre-part-%d" % partnum, sampleRate=sampleRate, nbChannels=nbChannels) - else: - pass - else: - # Nothing to do - pass - - if postIFrameTS < ts2: - for stream in streams: - if stream['codec_type'] == 'video': - extractPictures(inputFile=mkv, begin=postIFrameTS, nbFrames=nbPostFrame, prefix="post-part-%d" % partnum, width=width, height=height) - elif stream['codec_type'] == 'audio': - print("Extracting audio stream: %s" % stream) - sampleRate = stream['sample_rate'] - nbChannel = stream['channels'] - # TODO: how many packets should be dumped ... - # TODO: take into account multiple sound tracks ... - extractSound(inputFile=mkv, begin=postIFrameTS, nbPackets=nbPostFrame, outputFile="post-part-%d" % partnum, sampleRate=sampleRate, nbChannels=nbChannels) - else: - pass - else: - # Nothing to do ! - pass + 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) # Creating MKV file that corresponds to current part between I-frames - with open('part-%d.mkv' % partnum, 'w') as partmkv: - extractMKVPart(inputFile=mkv, outputFile=partmkv, begin=preIFrameTS, end=postIFrameTS) + with open('part-%d-internal.mkv' % partnum, 'w') as partmkv: + extractMKVPart(inputFile=mkv, outputFile=partmkv, begin=headIFrameTS, end=tailIFrameTS) # 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. @@ -399,8 +495,10 @@ def main(): # 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 + - # Appeler mkvmerge + # Appeler mkvmerge pour fusionner tous les fichiers part-%d.mkv if __name__ == "__main__": main()