Better performances and simplification by removing pipes and using memory file descriptors.

This commit is contained in:
Frédéric Tronel
2023-12-02 17:25:55 +01:00
parent d549311e20
commit 076e3c990b

View File

@@ -8,7 +8,8 @@ from datetime import datetime,timedelta,time
import coloredlogs, logging
from functools import cmp_to_key
from subprocess import Popen, PIPE
import os
from os import read, write, lseek, pipe, set_inheritable, memfd_create, SEEK_SET, close, unlink
import os.path
from io import BytesIO, TextIOWrapper
import json
from enum import Enum, IntEnum, unique, auto
@@ -63,7 +64,7 @@ def getFormat(inputFile):
infd = inputFile.fileno()
inputFile.seek(0,0)
os.set_inheritable(infd, True)
set_inheritable(infd, True)
with Popen(['ffprobe', '-loglevel', 'quiet', '-show_format', '-of', 'json', '-i', '/proc/self/fd/%d' % infd], stdout=PIPE, close_fds=False) as ffprobe:
out, _ = ffprobe.communicate()
out = json.load(BytesIO(out))
@@ -79,7 +80,7 @@ def getStreams(inputFile):
infd = inputFile.fileno()
inputFile.seek(0,0)
os.set_inheritable(infd, True)
set_inheritable(infd, True)
with Popen(['ffprobe', '-loglevel', 'quiet', '-show_streams', '-of', 'json', '-i', '/proc/self/fd/%d' % infd], stdout=PIPE, close_fds=False) as ffprobe:
out, _ = ffprobe.communicate()
out = json.load(BytesIO(out))
@@ -220,8 +221,8 @@ def ffmpegConvert(inputFile, inputFormat, outputFile, outputFormat, duration):
infd = inputFile.fileno()
outfd = outputFile.fileno()
os.set_inheritable(infd, True)
os.set_inheritable(outfd, True)
set_inheritable(infd, True)
set_inheritable(outfd, True)
# TODO: canvas size to be fixed !
with Popen(['ffmpeg', '-y', '-loglevel', 'quiet', '-progress', '/dev/stdout', '-canvas_size', '720x560', '-f', inputFormat, '-i', '/proc/self/fd/%d' % infd,
'-map', '0:v', '-map', '0:a', '-map', '0:s', '-bsf:v', 'h264_mp4toannexb,dump_extra=freq=keyframe', '-vcodec', 'copy', '-acodec', 'copy', '-scodec', 'dvdsub',
@@ -241,7 +242,7 @@ def ffmpegConvert(inputFile, inputFormat, outputFile, outputFormat, duration):
def getFramesInStream(inputFile, begin, end, streamKind, subStreamId=0):
logger = logging.getLogger(__name__)
infd = inputFile.fileno()
os.set_inheritable(infd, True)
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()
@@ -272,7 +273,7 @@ def getNearestIFrame(inputFile, timestamp, before=True, delta=timedelta(seconds=
tbegin = zero
infd = inputFile.fileno()
os.set_inheritable(infd, True)
set_inheritable(infd, True)
logger.debug('Looking for iframe in [%s, %s]' % (tbegin, tend))
@@ -324,8 +325,8 @@ def extractMKVPart(inputFile, outputFile, begin, end):
outputFile.seek(0,0)
infd = inputFile.fileno()
outfd = outputFile.fileno()
os.set_inheritable(infd, True)
os.set_inheritable(outfd, True)
set_inheritable(infd, True)
set_inheritable(outfd, True)
warnings = []
with Popen(['mkvmerge', '-o', '/proc/self/fd/%d' % outfd, '--split', 'parts:%s-%s' % (begin, end), '/proc/self/fd/%d' % infd], stdout=PIPE, close_fds=False) as mkvmerge:
pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%', desc='Extraction')
@@ -354,10 +355,9 @@ def extractPictures(inputFile, begin, nbFrames, width=640, height=480):
inputFile.seek(0,0)
infd = inputFile.fileno()
fdr, fdw = os.pipe()
os.set_inheritable(infd, True)
os.set_inheritable(fdr, False)
os.set_inheritable(fdw, True)
outfd = memfd_create('pictures', flags=0)
set_inheritable(outfd, True)
# PPM header
# "P6\nWIDTH HEIGHT\n255\n"
headerLen=2+1+ceil(log(width, 10))+1+ceil(log(height, 10))+1+3+1
logger.debug('Header length: %d' % headerLen)
@@ -365,67 +365,48 @@ def extractPictures(inputFile, begin, nbFrames, width=640, height=480):
length = imageLength*nbFrames
logger.debug("Estimated length: %d" % length)
pg = trange(nbFrames)
images = bytes()
with Popen(['ffmpeg', '-loglevel', 'quiet' ,'-y', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-s', '%dx%d'%(width, height), '-vframes', '%d'%nbFrames, '-c:v', 'ppm', '-f', 'image2pipe', '/proc/self/fd/%d' % fdw ], stdout=PIPE, close_fds=False) as ffmpeg:
while ffmpeg.poll() == None:
# TODO: understand why this line ends up in reading on an already closed file descriptor
# fds, _, _ = select([fdr, ffmpeg.stdout], [], [], .1)
fds, _, _ = select([fdr], [], [], .1)
if fdr in fds:
buf = os.read(fdr, imageLength)
# print("Read %d bytes of image. ffmpeg finished: %s" % (len(buf), ffmpeg.poll()))
if len(buf) == 0:
break
pg.update(len(buf)/imageLength)
images=images+buf
if ffmpeg.stdout in fds:
for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"):
logger.debug(line)
with Popen(['ffmpeg', '-loglevel', 'quiet' ,'-y', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-s', '%dx%d'%(width, height), '-vframes', '%d'%nbFrames, '-c:v', 'ppm', '-f', 'image2pipe', '/proc/self/fd/%d' % outfd ], stdout=PIPE, close_fds=False) as ffmpeg:
status = ffmpeg.wait()
# Finishing to read residual bytes from pipe
while True:
fd, _, _ = select([fdr], [], [], .1)
if fd != []:
buf = os.read(fdr, imageLength)
# print("Read %d bytes of image" % len(buf))
if len(buf) == 0:
break
pg.update(len(buf)/imageLength)
images=images+buf
else:
# Nothing more to read
break
logger.debug("%d bytes received." % len(images))
os.close(fdr)
os.close(fdw)
if status != 0:
logger.error('Image extraction returns error code: %d' % status)
logger.error('Conversion failed with status code: %d' % status)
return None, None
return images
lseek(outfd, 0, SEEK_SET)
images = read(outfd,length)
if len(images) != length:
logger.info("Received %d bytes but %d were expected." % (len(images), length))
return None, None
def extractSound(inputFile, begin, outputFile, subChannel=0, nbPackets=0, sampleRate=48000, nbChannels=2):
lseek(outfd, 0, SEEK_SET)
return images, outfd
def extractSound(inputFile, begin, outputFileName, packetDuration, subChannel=0, nbPackets=0, sampleRate=48000, nbChannels=2):
logger = logging.getLogger(__name__)
inputFile.seek(0,0)
outputFile.seek(0,0)
outfd = memfd_create(outputFileName, flags=0)
infd = inputFile.fileno()
outfd = outputFile.fileno()
os.set_inheritable(infd, True)
os.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"):
logger.debug(line)
set_inheritable(infd, True)
set_inheritable(outfd, True)
sound = bytes()
length = int(nbChannels*sampleRate*4*nbPackets*packetDuration/1000)
with Popen(['ffmpeg', '-y', '-loglevel', 'quiet', '-ss', '%s'%begin, '-i', '/proc/self/fd/%d' % infd, '-frames:a:%d' % subChannel, '%d' % (nbPackets+1),
'-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:
status = ffmpeg.wait()
if status != 0:
logger.error('Sound extraction returns error code: %d' % status)
return None, None
lseek(outfd, 0, SEEK_SET)
sound = read(outfd, length)
if (len(sound) != length):
logger.info("Received %d bytes but %d were expected (channels=%d, freq=%d, packets=%d, duration=%d ms)." % (len(sound), length, nbChannels, sampleRate, nbPackets, packetDuration))
return None, None
return sound, outfd
def dumpPPM(pictures, prefix, temporaries):
logger = logging.getLogger(__name__)
@@ -463,14 +444,12 @@ def dumpPPM(pictures, prefix, temporaries):
length=headerLen+3*width*height
nbBytes = 0
while nbBytes < length:
nbBytes+=os.write(outfd, pictures[pos+nbBytes:pos+length])
nbBytes+=write(outfd, pictures[pos+nbBytes:pos+length])
pos+=length
picture+=1
def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, width, height, temporaries, dumpPictures=False):
def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, width, height, temporaries, dumpMemFD=False):
logger = logging.getLogger(__name__)
# encoderParams = [ 'ffmpeg', '-y', '-loglevel', 'quiet' ]
encoderParams = [ 'ffmpeg', '-y', '-loglevel', 'quiet' ]
inputParams = []
codecsParams = []
@@ -479,8 +458,6 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid
videoID=0
audioID=0
subTitleID=0
audioFiles = {}
imagesPipes = {}
for stream in streams:
if stream['codec_type'] == 'video':
logger.info("Extracting video stream v:%d" % videoID)
@@ -506,20 +483,18 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid
# TODO: adjust SAR and DAR
# https://superuser.com/questions/907933/correct-aspect-ratio-without-re-encoding-video-file
codec = stream['codec_name']
imagesBytes = extractPictures(inputFile=inputFile, begin=begin, nbFrames=nbFrames, width=width, height=height)
if dumpPictures:
imagesBytes, memfd = extractPictures(inputFile=inputFile, begin=begin, nbFrames=nbFrames, width=width, height=height)
if imagesBytes == None:
sys.exit(-1)
if dumpMemFD:
dumpPPM(imagesBytes, '%s-%d' % (filesPrefix,videoID), temporaries)
# imagesBytes contains now a buffer of bytes that represents the pictures that have been dumped by ffmpeg.
fdr, fdw = os.pipe()
os.set_inheritable(fdr, True)
# The writalbe end of the pipe (fdw) must not be stayed opened in ffmpeg child, otherwise ffmpeg will not be able
# to detect the end of pictures data sent by the other end of the pipe it is reading from (fdr).
# We manually force non inheritance to be sure (although this should be the case since Python 3.4).
os.set_inheritable(fdw, False)
logger.debug("Creating pipes for images: r:%d w:%d" % (fdr,fdw))
imagesPipes[videoID] = (imagesBytes, fdr, fdw)
inputParams.extend(['-framerate', '%f'%frameRate, '-f', 'image2pipe', '-i', '/proc/self/fd/%d' % fdr])
# We rewind to zero the memory file descriptor
lseek(memfd, 0, SEEK_SET)
set_inheritable(memfd, True)
inputParams.extend(['-framerate', '%f'%frameRate, '-f', 'image2pipe', '-i', '/proc/self/fd/%d' % memfd])
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
@@ -535,14 +510,31 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid
packets = getFramesInStream(inputFile=inputFile, begin=begin, end=end, streamKind='a', subStreamId=audioID)
nbPackets = len(packets)
logger.debug("Found %d packets to be extracted from audio track." % nbPackets)
if(nbPackets > 0):
packetDuration = packets[0]['duration']
tmpname = '%s-%d.pcm' % (filesPrefix,audioID)
soundBytes , memfd = extractSound(inputFile=inputFile, begin=begin, nbPackets=nbPackets, packetDuration=packetDuration, outputFileName=tmpname, sampleRate=sampleRate, nbChannels=nbChannels)
if dumpMemFD:
try:
audioFiles[audioID] = open('%s-%d.pcm' % (filesPrefix,audioID), 'w')
output = open(tmpname,'w')
except IOError:
logger.error('Impossible to create file: %s-%d.pcm' % (filesPrefix,audioID))
logger.error('Impossible to create file: %s' % tmpname)
return None
temporaries.append(audioFiles[audioID])
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()])
outfd = output.fileno()
pos = 0
while pos < len(soundBytes):
pos+=write(outfd, soundBytes[pos:])
temporaries.append(output)
# We rewind to zero the memory file descriptor
lseek(memfd, 0, SEEK_SET)
set_inheritable(memfd, True)
inputParams.extend(['-f', 's32le', '-ar', '%d'%sampleRate, '-ac', '%d'%nbChannels, '-i', '/proc/self/fd/%d' % memfd])
codecsParams.extend(['-c:a:%d' % audioID, codec, '-b:a:%d' % audioID, '%d' % bitRate])
audioID=audioID+1
elif stream['codec_type'] == 'subtitle':
@@ -570,40 +562,16 @@ def extractAllStreams(inputFile, begin, end, streams, filesPrefix, nbFrames, wid
return None
outfd = output.fileno()
os.set_inheritable(outfd, True)
set_inheritable(outfd, True)
# TODO: manage interlaced to previous parameters.
encoderParams.extend(['-top', '1', '-flags:v', '+ilme+ildct', '-bsf:v', 'h264_mp4toannexb,dump_extra=freq=keyframe', '-f', 'matroska', '/proc/self/fd/%d' % outfd])
logger.info('Encoding video: %s' % fileName)
with Popen(encoderParams, stdout=PIPE, close_fds=False) as ffmpeg:
pos = {}
totalLength = 0
for vid in range(videoID):
pos[vid]=0
img, fdr, _ = imagesPipes[vid]
# We close the end of the pipe used by ffmepg to read data.
os.close(fdr)
totalLength+=len(img)
length = 0
pg = trange(totalLength)
while length<totalLength:
for vid in range(videoID):
img, _, fdw = imagesPipes[vid]
nbBytes = os.write(fdw, img[pos[vid]:])
pos[vid]=pos[vid]+nbBytes
length+=nbBytes
pg.update(length)
for vid in range(videoID):
_, _, fdw = imagesPipes[vid]
os.close(fdw)
for line in TextIOWrapper(ffmpeg.stdout, encoding="utf-8"):
logger.debug(line)
status = ffmpeg.wait()
if status != 0:
logger.error('Encoding failed with status code: %d' % status)
return None
temporaries.append(output)
@@ -626,7 +594,7 @@ def mergeMKVs(inputs, outputName):
outfd = out.fileno()
fds.append(outfd)
os.set_inheritable(outfd, True)
set_inheritable(outfd, True)
mergeParams = ['mkvmerge']
first = True
@@ -634,7 +602,7 @@ def mergeMKVs(inputs, outputName):
if mkv !=None:
fd = mkv.fileno()
fds.append(fd)
os.set_inheritable(fd, True)
set_inheritable(fd, True)
if first:
mergeParams.append('/proc/self/fd/%d' % fd)
first = False
@@ -667,7 +635,7 @@ def mergeMKVs(inputs, outputName):
logger.error('Extraction returns errors')
for fd in fds:
os.set_inheritable(fd, False)
set_inheritable(fd, False)
return out
@@ -704,7 +672,7 @@ def main():
parser.add_argument("-o", "--output", dest='outputFile', 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("--dump-pictures", action='store_true', dest='dump', help="For debug purpose, dump pictures of headers (and trailers) before (after) each part. They are kept in memory only otherwise.")
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 ")
args = parser.parse_args()
@@ -853,7 +821,7 @@ def main():
if nbHeadFrames > 0:
# We extract all frames between the beginning upto the frame that immediately preceeds the I-frame.
head = extractAllStreams(inputFile=mkv, begin=ts1, end=headIFrameTS, nbFrames=nbHeadFrames-1, filesPrefix='part-%d-head' % (partnum), streams=streams, width=width, height=height, temporaries=temporaries, dumpPictures=args.dump)
head = extractAllStreams(inputFile=mkv, begin=ts1, end=headIFrameTS, nbFrames=nbHeadFrames-1, filesPrefix='part-%d-head' % (partnum), streams=streams, width=width, height=height, temporaries=temporaries, dumpMemFD=args.dump)
subparts.append(head)
# Creating MKV file that corresponds to current part between I-frames
@@ -868,7 +836,7 @@ def main():
if nbTailFrames > 0:
# We extract all frames between the I-frame (including it) upto the end.
tail = extractAllStreams(inputFile=mkv, begin=tailIFrameTS, end=ts2, nbFrames=nbTailFrames, filesPrefix='part-%d-tail' % (partnum), streams=streams, width=width, height=height, temporaries=temporaries, dumpPictures=args.dump)
tail = extractAllStreams(inputFile=mkv, begin=tailIFrameTS, end=ts2, nbFrames=nbTailFrames, filesPrefix='part-%d-tail' % (partnum), streams=streams, width=width, height=height, temporaries=temporaries, dumpMemFD=args.dump)
subparts.append(tail)
logger.info('Merging: %s' % subparts)
@@ -928,7 +896,7 @@ def main():
path = os.path.realpath(f.name)
logger.info("Removing: %s" % path)
f.close()
os.unlink(path)
unlink(path)
for c in checks:
logger.info("Please check cut smoothness at: %s" % c)