diff --git a/removeads.py b/removeads.py index 1acb168..251683b 100755 --- a/removeads.py +++ b/removeads.py @@ -16,7 +16,7 @@ import shutil from tqdm import tqdm, trange from select import select from math import floor, ceil, log -from shutil import copyfile, which +from shutil import copyfile, which, move import hexdump from iso639 import Lang from iso639.exceptions import InvalidLanguageValue @@ -107,9 +107,10 @@ def getSubTitlesTracks(ffprobe, mkvPath): return tracks -def extractSRT(mkvextract, mkvPath, subtitles, langs): +def extractSRT(mkvextract, fileName, subtitles, langs): logger = logging.getLogger(__name__) - params = [mkvextract, mkvPath, 'tracks'] + + params = [mkvextract, fileName, 'tracks'] res = [] @@ -133,8 +134,20 @@ def extractSRT(mkvextract, mkvPath, subtitles, langs): count = count+1 env = {**os.environ, 'LANG': 'C'} - with Popen(params, env=env) as extract: - extract.wait() + with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract: + 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[0-9]{1,3})%$') + m = p.match(line) + if m == None: + logger.error('Impossible to parse progress') + pb.update(int(m['progress'])-pb.n) + pb.update(100-pb.n) + pb.refresh() + pb.close() + + extract.wait() if extract.returncode != 0: logger.error('Mkvextract returns an error code: %d' % extract.returncode) @@ -981,6 +994,8 @@ def dumpPPM(pictures, prefix, temporaries): nbBytes+=write(outfd, pictures[pos+nbBytes:pos+length]) pos+=length picture+=1 + + def extractAllStreams(ffmpeg, ffprobe, inputFile, begin, end, streams, filesPrefix, nbFrames, width, height, temporaries, dumpMemFD=False): logger = logging.getLogger(__name__) @@ -988,6 +1003,7 @@ def extractAllStreams(ffmpeg, ffprobe, inputFile, begin, end, streams, filesPref inputParams = [] codecsParams = [] + if begin < end: videoID=0 audioID=0 @@ -1001,6 +1017,8 @@ def extractAllStreams(ffmpeg, ffprobe, inputFile, begin, end, streams, filesPref m = pattern.match(frameRate) if m != None: frameRate = float(m['numerator']) / float(m['denominator']) + # TODO: Framerate estimation seems broken ... + frameRate = 25. sar = stream['sample_aspect_ratio'] dar = stream['display_aspect_ratio'] pixelFormat = stream['pix_fmt'] @@ -1217,6 +1235,52 @@ def extractSubTitleTrack(mkvmerge, inputFileName, index, lang): for lines in out: logger.info(out) +def remuxSRTSubtitles(mkvmerge, inputFile, outputFileName, subtitles): + logger = logging.getLogger(__name__) + + try: + out = open(outputFileName, 'w') + except IOError: + logger.error('Impossible to create file: %s' % outputFileName) + return None + + outfd = out.fileno() + infd = inputFile.fileno() + lseek(infd, 0, SEEK_SET) + set_inheritable(infd, True) + set_inheritable(outfd, True) + + mkvmergeParams = [mkvmerge, '/proc/self/fd/%d' % infd] + for fd, lang in subtitles: + lseek(fd, 0, SEEK_SET) + set_inheritable(fd, True) + mkvmergeParams.extend(['--language', '0:%s' % lang, '/proc/self/fd/%d' % fd]) + + mkvmergeParams.extend(['-o', '/proc/self/fd/%d' % outfd]) + + warnings = [] + env = {**os.environ, 'LANG': 'C'} + logger.info('Remux subtitles: %s' % mkvmergeParams) + with Popen(mkvmergeParams, stdout=PIPE, close_fds=False, env=env) as mkvmerge: + 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[0-9]{1,3})%$') + m = p.match(line) + if m == None: + logger.error('Impossible to parse progress') + 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') def main(): @@ -1431,30 +1495,35 @@ def main(): # We need to check the end also checks.append(pos) + finalCutName = '%s-cut.mkv' % basename + nbMKVParts = len(mkvparts) if nbMKVParts > 1: logger.info('Merging: %s' % mkvparts) - mergeMKVs(inputs=mkvparts, outputName=args.outputFile) + mergeMKVs(inputs=mkvparts, outputName=finalCutName) elif nbMKVParts == 1: - copyfile('part-1.mkv', args.outputFile) + copyfile('part-1.mkv', finalCutName) else: logger.info("Nothing else to do.") + + try: + finalCut = open(finalCutName, mode='r') + except IOError: + logger.error("Impossible to open %s to finalize processing." % finalCutName) + exit(-1) if args.srt: if not allOptionalTools: logger.warning("Missing tools for extracting subtitles.") + move(finalCutName, args.outputFile) else: - try: - final = open(args.outputFile, mode='r') - except IOError: - logger.error("Impossible to open %s to finalize processing." % args.outputFile) - exit(-1) - - duration = getMovieDuration(paths['ffprobe'], final) + # Final cut is not any more the final step. + temporaries.append(finalCut) + duration = getMovieDuration(paths['ffprobe'], finalCut) supportedLangs = getTesseractSupportedLang(paths['tesseract']) logger.info('Supported lang: %s' % supportedLangs) logger.info('Find subtitles tracks and language.') - subtitles = findSubtitlesTracks(paths['ffprobe'], final) + subtitles = findSubtitlesTracks(paths['ffprobe'], finalCut) logger.info(subtitles) sts = {} for subtitle in subtitles: @@ -1472,7 +1541,7 @@ def main(): logger.error("Dropping subtitle: %s because it is missing language indication") logger.debug(sts) - listOfSubtitles = extractSRT(paths['mkvextract'], args.outputFile, sts, supportedLangs) + listOfSubtitles = extractSRT(paths['mkvextract'], finalCutName, sts, supportedLangs) logger.info(listOfSubtitles) for idxName, subName, _, _ in listOfSubtitles: try: @@ -1491,7 +1560,11 @@ def main(): ocr = doOCR(paths['vobsubocr'], listOfSubtitles, duration, temporaries, args.dump) logger.info(ocr) - + + # Remux SRT subtitles + remuxSRTSubtitles(paths['mkvmerge'], finalCut, args.outputFile, ocr) + else: + move(finalCutName, args.outputFile) if not args.keep: logger.info("Cleaning temporary files")