ts2mkv.py

#!/usr/bin/env python

import subprocess
from eit import readeit, eitinfo
import os, shlex

def filter_lines(data, search):
    ret = []
    for line in data.split("\n"):
        if line.find(search) == -1:
            continue
        ret.append(line)
    return "\n".join(ret)

def run_command(command):
    process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE)
    while True:
        output = process.stdout.readline()
        if output == '' and process.poll() is not None:
            break
        if output:
            print output.strip()
    rc = process.poll()
    return rc

class ts2mkv(object):
    def __init__(self, crf=19, tune='film', scaleto_720p=True, rename=False):
        self.msg_prepare = ""
        self.msg_eit = ""
        self.msg_ffmpeg = ""
        self.command = None

        self.scaleto_720p = scaleto_720p
        self.rename = rename

        self.video_options = [
            "-c:v libx264",
            "-preset faster", # slow
            "-tune %s" % tune, # film / animation
            "-crf %i" % crf # 21, better 19
        ]
        self.audio_options = [
            "-c:a copy",
        ]

    def get_stream_index(self, data):
        idx = data.find("Stream #")
        if idx == -1:
            return ""
        idx += 8
        self.msg_prepare += "Selecting: %s\n" % data
        return data[idx:idx+3]

    def get_movie_description(self):
        if not self.filename:
            return
        # read the EIT file
        # TODO: fallback to meta file if no EIT
        # TODO: is there a way to get the imdb for the movie automagically?
        # http://www.omdbapi.com/apikey.aspx
        filename = os.path.splitext(self.filename)[0] + ".eit"
        self.msg_eit = readeit(filename)
        if not self.rename or not self.msg_eit:
            return
        info = eitinfo(filename)
        name = info.eit.get("name")
        if name == "":
            # cancel rename, no movie title found!
            return
        genre = info.eit.get("genre")
        if genre != "":
            name = "%s (%s)" % (name, genre)
        # build new filename
        name = name.replace(' : ', ' - ')
        name = name.replace(': ', ' - ')
        name = name.replace(':', '-')
        name = name.replace('/', '')
        name = name.replace('\\', '')
        name = name.replace('?', '')
        name = name.replace('*', '')
        name = name.replace('\"', '\'')

        self.outfilebase = os.path.join(
            os.path.dirname(filename),
            name
            )


    def get_crop_option(self):
        lines = filter_lines(self.msg_ffmpeg, "[Parsed_cropdetect").split("\n")
        option = None
        for line in lines:
            tmp = line[line.find(" crop="):].strip()
            #print "DEBUG: " + tmp
            if not option:
                option = tmp
            else:
                if option != tmp:
                    self.msg_prepare += "WARNING: cropdetect inconsistent over scan time, disabling autocrop\n"
                    return None
        self.msg_prepare += "Crop detected: %s\n" % option
        return option

    def get_ffmpeg_command(self):
        if not self.filename:
            return None

        commands = []
        fn = "\\'".join(p for p in self.filename.split("'"))
        fn = fn.replace(" ", "\\ ")
        outfn = self.outfilebase + ".mkv"
        outfn = "\\'".join(p for p in outfn.split("'"))
        outfn = outfn.replace(" ", "\\ ")

        #  ffmpeg -ss 00:05:00 -t 2 -i testfiles/chappie.ts -vf "cropdetect=24:16:0" -f null -

        cmd = ["ffmpeg",
                "-ss 00:05:00", "-t 1", # search to 5 minutes, analyze 1 second
                "-i %s" % fn,
                "-vf \"cropdetect=24:2:0\"", # detect black bar crop on top and bottom
                "-f null", "-" # no output file
                ]
        print " ".join(cmd)
        p = subprocess.Popen(shlex.split(" ".join(cmd)), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = p.communicate()
        print "Command return code: ", p.poll()
        self.msg_ffmpeg = out + "\n" + err
        self.msg_ffmpeg = self.msg_ffmpeg[self.msg_ffmpeg.find("Input #0"):]

        # find "Stream #0:" lines
        info = filter_lines(self.msg_ffmpeg, "Stream #0:")


        # CRAP! ffmpeg cannot decode dvb_teletext streams to srt :(
        # are there any subtitle streams?!
        #s = filter_lines(info, "Subtitle:")
        #s = filter_lines(s, "(deu):")
        #if s != "":
        #    s = self.get_stream_index(s.split("\n")[0])
        #    commands.append([
        #        "ffmpeg",
        #        "-txt_format text",
        #        "-i %s" % fn,
        #        "-map %s" % s,
        #        "%s.srt" % os.path.splitext(fn)[0]
        #        ])

        v = self.get_stream_index(
            filter_lines(info, "Video:"))
        if v == "":
            print "No video stream found"
            return None

        audiomap = []
        audioall = filter_lines(info, "Audio:")
        audio = filter_lines(audioall, "(deu):")
        a = self.get_stream_index(
            filter_lines(audio, "ac3"))
        # TODO: wenn kein ac3 stream dann dts oder mpeg fallback
        if a == "":
            print audioall
            print "No AC3 german audio stream found"
            # try to find the first german audio stream
            a = self.get_stream_index(audio.split("\n")[0])
            if a == "":
                print "No other german audio streams, trying english ac3..."
            else:
                print "Selecting first german stream."
                audiomap.append(a)
        else:
            audiomap.append(a)

        audio = filter_lines(audioall, "(eng):")
        a = self.get_stream_index(
            filter_lines(audio, "ac3"))
        if a != "":
            # append english audio too!
            print "Selecting english ac3 stream."
            audiomap.append(a)

        if len(audiomap) == 0:
            print "No suitable audio stream found, aborting."
            return None


        self.msg_prepare += "Video Stream selected: Stream #%s\n" % v
        cmd = [
            "ffmpeg",
            "-i %s" % fn,
            "-map %s" % v,
            ]
        flt = []
        crop = self.get_crop_option()
        if crop:
            flt.append(crop)
        if self.scaleto_720p:
            flt.append("scale='min(1280,iw)':-2'") # -2 ensures division by two for codec
            self.msg_prepare += "Scaling cropped output stream to 720p\n"
        if len(flt) > 0:
            # append video filters
            cmd.append('-filter:v "%s"' % ",".join(flt))
        for a in audiomap:
            self.msg_prepare += "Audio Stream selected: Stream #%s\n" % a
            cmd.append("-map %s" % a)
        cmd.extend(self.video_options)
        cmd.extend(self.audio_options)
        cmd.append(outfn)

        commands.append(" ".join(cmd))
        return commands

    def load(self, filename):
        self.filename = filename
        self.outfilebase = os.path.splitext(filename)[0]
        self.get_movie_description()
        self.command = self.get_ffmpeg_command()

    def convert(self):
        if not self.command:
            return None
        fd = open(self.outfilebase + ".txt", "wb")
        fd.write(self.msg_eit)
        fd.write("\n\n# ---DEBUG---\n\n")
        fd.write(self.msg_prepare)
        fd.write(self.msg_ffmpeg)
        fd.close()
        #print self.msg_ffmpeg

        for cmd in mkv.command:
            print "Executing ffmpeg:\n%s\n" % cmd
            return run_command(cmd)


if __name__ == "__main__":
    # parse command line options
    import argparse, sys, glob

    parser = argparse.ArgumentParser(description = 'DVB-TS to MKV kung-fu')
    parser.add_argument('--crf', type=int, default=19,
        help='h264 crf (default 19)')
    parser.add_argument('--tune', default='film',
        help='ffmpeg tune preset [film, animation] (default is film)')
    parser.add_argument('--ns', action='store_true', default=False,
        help='no rescaling (default is scale to 720p)')
    parser.add_argument('--rename', action='store_true', default=False,
        help='rename file basename to name and genre from EIT file if present')
    parser.add_argument('input', metavar='input', nargs='+',
        help='one or more files, glob style syntax')

    args = parser.parse_args()
    mkv = ts2mkv(
        crf = args.crf,
        tune = args.tune,
        scaleto_720p = (not args.ns),
        rename = args.rename
        )

    #os.system('cls' if os.name == 'nt' else 'clear')

    for srcstr in args.input:
        src = glob.glob(srcstr)
        for filename in src:
            print filename
            mkv.load(filename)
            mkv.convert()