Better quoting for --dump

[transcoding] / encode.py

#!/usr/bin/env python

import optparse
import re
import subprocess
import sys
import os
import shutil
import tempfile

class FatalException(Exception):
        pass

def mkarg(arg):
        if re.match("^[a-zA-Z0-9\-\\.,/@_:=]*$", arg):
                return arg

        if "'" not in arg:
                return "'%s'" % arg
        out = "\""
        for c in arg:
                if c in "\\$\"`":
                        out += "\\"
                out += c
        out += "\""
        return out

class Command(object):
        def __init__(self, profile, opts):
                self.profile = profile
                self.opts = opts
        
        def print_install_message(self):
                print >>sys.stderr, "Problem with command: %s", self.name
                if self.package:
                        print >>sys.stderr, "Try running:\n# aptitude install %s", self.package
        
        def check_command(self, cmd):
                if self.opts.dump:
                        return
                if subprocess.Popen(["which", cmd], stdout=open("/dev/null", "w")).wait() != 0:
                        raise FatalException("Command '%s' is required" % cmd)
        
        def check_no_file(self, path):
                if os.path.exists(path):
                        raise FatalException("Output file '%s' exists." % path)

        def do_exec(self, args):
                if self.opts.dump:
                        print " ".join(map(mkarg, args))
                else:
                        if subprocess.Popen(args).wait() != 0:
                                raise FatalException("Failure executing command: %s" % args)


class MP4Box(Command):
        codec2exts = {
                "xvid": "m4v",
                "x264": "h264",
                "faac": "aac",
        }

        def check(self):
                self.check_command("mencoder")
                self.check_command("MP4Box")
                self.check_no_file(self.opts.output + ".mp4")

        def run(self):
                p = self.profile
                video = "video.%s" % self.codec2exts[p.vcodec]
                audio = "audio.%s" % self.codec2exts[p.acodec]
                input = self.opts.output + ".avi" # From Mencoder command
                output = self.opts.output + ".mp4"
                mencoder = ["mencoder", input, "-ovc", "copy", "-oac", "copy", "-of"]
                self.do_exec(["rm", "-f", output])
                self.do_exec(mencoder + ["rawvideo", "-o", video])
                self.do_exec(mencoder + ["rawaudio", "-o", audio])
                self.do_exec(["MP4Box", "-add", video, "-add", audio, output])
                self.do_exec(["rm", "-f", video, audio, input])



class MKVMerge(Command):
        def check(self):
                self.check_command("mkvmerge")
                self.check_no_file(self.opts.output + ".mkv")

        def run(self):
                input = self.opts.output + ".avi" # From Mencoder command
                output = self.opts.output + ".mkv"
                self.do_exec(["mkvmerge", "-o", output, input])
                self.do_exec(["rm", "-f", input])



class Mencoder(Command):
        codec2opts = {
                "lavc": "-lavcopts",
                "xvid": "-xvidencopts",
                "x264": "-x264encopts",
                "faac": "-faacopts",
                "mp3lame": "-lameopts",
        }

        def insert_options(self, cmd):
                def try_opt(opt, var):
                        if var is not None:
                                cmd.append(opt)
                                cmd.append(var)
                if self.opts.deinterlace:
                        cmd += ["-vf-add", "pp=lb"]
                try_opt("-ss", self.opts.startpos)
                try_opt("-endpos", self.opts.endpos)
                try_opt("-dvd-device", self.opts.dvd)
                try_opt("-chapter", self.opts.chapter)
                try_opt("-aid", self.opts.audioid)
                try_opt("-sid", self.opts.subtitleid)
                try_opt("-vf-add", self.opts.vfilters)
                try_opt("-af", self.opts.afilters)

        def subst_values(self, cmd, vpass):
                subst = {
                        "vbitrate": self.opts.vbitrate,
                        "abitrate": self.opts.abitrate,
                        "input": self.opts.input,
                        "output": self.opts.output + ".avi",
                        "vpass": vpass,
                }

                return [x % subst for x in cmd]
        
        def pass1(self):
                p = self.profile
                cmd = []
                cmd += ["mencoder", "%(input)s", "-o", "/dev/null"]
                self.insert_options(cmd)
                cmd += ["-ovc", p.vcodec, self.codec2opts[p.vcodec], p.vopts]
                cmd += ["-oac", "copy"]
                cmd += self.profile.extra + self.profile.extra1
                cmd = self.subst_values(cmd, vpass=1)
                return cmd

        def pass2(self):
                p = self.profile
                cmd = []
                cmd += ["mencoder", "%(input)s", "-o", "%(output)s"]
                self.insert_options(cmd)
                cmd += ["-ovc", p.vcodec, self.codec2opts[p.vcodec], p.vopts]
                cmd += ["-oac", p.acodec, self.codec2opts[p.acodec], p.aopts]
                if self.opts.episode_name:
                        cmd += ["-info", "name='%s'" % self.opts.episode_name]
                cmd += self.profile.extra + self.profile.extra2
                cmd = self.subst_values(cmd, vpass=2)
                return cmd

        def check(self):
                self.check_command("mencoder")
                self.check_no_file(self.opts.output + ".avi")

        def run(self):
                self.do_exec(self.pass1())
                self.do_exec(self.pass2())



class Profile(object):
        def __init__(self, commands, **kwargs):
                self.default_opts = {
                        "vbitrate": 1000,
                        "abitrate": 192,
                }
                self.extra = []
                self.extra1 = []
                self.extra2 = []
                self.commands = commands
                self.__dict__.update(kwargs)

        def __contains__(self, keyname):
                return hasattr(self, keyname)


profiles = {
        "qt7" :
        Profile(
                commands=[Mencoder, MP4Box],
                vcodec="x264",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:me=umh:partitions=all:trellis=1:subq=7:bframes=1:direct_pred=auto",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
        ),

        "x264" :
        Profile(
                commands=[Mencoder, MKVMerge],
                vcodec="x264",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:subq=6:frameref=6:me=umh:partitions=all:bframes=4:b_adapt:qcomp=0.7:keyint=250",
                acodec="mp3lame",
                aopts="abr:br=%(abitrate)d",
        ),

        "xvid" :
        Profile(
                commands=[Mencoder],
                vcodec="xvid",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:vhq=4:autoaspect",
                acodec="mp3lame",
                aopts="abr:br=%(abitrate)d",
                extra2=["-ffourcc", "DX50"],
        ),

        "ipodxvid" :
        Profile(
                commands=[Mencoder, MP4Box],
                vcodec="xvid",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:vhq=4:autoaspect:max_bframes=0",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
                extra=["-vf-add", "scale=480:-10"],
        ),

        "ipodx264" :
        Profile(
                commands=[Mencoder, MP4Box],
                vcodec="x264",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:vbv_maxrate=1500:vbv_bufsize=2000:nocabac:me=umh:partitions=all:trellis=1:subq=7:bframes=0:direct_pred=auto:level_idc=30:turbo",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
                extra=["-vf-add", "scale=480:-10"],
                extra2=["-channels", "2", "-srate", "48000"],
        ),

        "nokiax264" :
        Profile(
                commands=[Mencoder, MP4Box],
                default_opts={
                        "vbitrate": 256,
                        "abitrate": 96,
                },
                vcodec="x264",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:nocabac:me=umh:partitions=all:trellis=1:subq=7:bframes=0:direct_pred=auto",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
                extra=["-vf-add", "scale=320:-10"],
        ),

        "n97xvid" :
        Profile(
                commands=[Mencoder, MP4Box],
                default_opts={
                        "vbitrate": 1000,
                        "abitrate": 96,
                },
                vcodec="xvid",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:vhq=4:autoaspect:max_bframes=0",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
                extra=["-vf-add", "scale=640:-10"],
        ),

        "n97x264" :
        Profile(
                commands=[Mencoder, MP4Box],
                default_opts={
                        "vbitrate": 1000,
                        "abitrate": 96,
                },
                vcodec="x264",
                vopts="pass=%(vpass)d:bitrate=%(vbitrate)d:vbv_maxrate=2000:vbv_bufsize=2000:nocabac:me=umh:partitions=all:trellis=1:subq=7:bframes=0:direct_pred=auto:level_idc=20",
                acodec="faac",
                aopts="br=%(abitrate)d:mpeg=4:object=2",
                extra=["-vf-add", "scale=640:-10"],
        ),
}




def parse_args():
        for profile_name in profiles.keys():
                if sys.argv[0].find(profile_name) >= 0:
                        break
        else:
                profile_name = "xvid"

        parser = optparse.OptionParser(usage="%prog [options] input [output]")
        parser.add_option("--dvd", action="store", dest="dvd")
        parser.add_option("--deinterlace", action="store_true", dest="deinterlace")
        parser.add_option("--vfilters", action="store", dest="vfilters")
        parser.add_option("--afilters", action="store", dest="afilters")
        parser.add_option("--vbitrate", action="store", dest="vbitrate", type="int")
        parser.add_option("--abitrate", action="store", dest="abitrate", type="int")
        parser.add_option("--chapter", action="store", dest="chapter")
        parser.add_option("--startpos", action="store", dest="startpos")
        parser.add_option("--endpos", action="store", dest="endpos")
        parser.add_option("--audioid", action="store", dest="audioid")
        parser.add_option("--subtitleid", action="store", dest="subtitleid")
        parser.add_option("--profile", action="store", dest="profile_name", default=profile_name)
        parser.add_option("--episode-name", action="store", dest="episode_name")
        parser.add_option("--dump", action="store_true", dest="dump")
        try:
                opts, args = parser.parse_args(sys.argv[1:])
                if len(args) == 1:
                        input = args[0]
                        output = os.path.splitext(os.path.basename(input))[0]
                elif len(args) == 2:
                        input, output = args
                else:
                        raise ValueError
        except Exception:
                parser.print_usage()
                sys.exit(1)
        
        if "://" not in input:
                opts.input = os.path.abspath(input)
        else:
                if opts.dvd:
                        opts.dvd = os.path.abspath(opts.dvd)
                opts.input = input

        opts.output = os.path.abspath(output)

        return opts

def main():
        opts = parse_args()

        # Find our profile
        try:
                profile = profiles[opts.profile_name]
        except KeyError:
                print >>sys.stderr, "Profile '%s' not found!" % opts.profile_name
                sys.exit(1)

        # Pull in default option values from the profile
        for key, value in profile.default_opts.iteritems():
                if getattr(opts, key) is None:
                        setattr(opts, key, value)

        # Run in a temp dir so that multiple instances can be run simultaneously
        tempdir = tempfile.mkdtemp()
        try:
                os.chdir(tempdir)

                try:
                        commands = []
                        for CommandClass in profile.commands:
                                command = CommandClass(profile, opts)
                                commands.append(command)
                                command.check()
                        for command in commands:
                                command.run()

                except FatalException, e:
                        print >>sys.stderr, "Error:", e.message
                        sys.exit(1)

        finally:
                os.chdir("/")
                shutil.rmtree(tempdir)

if __name__ == "__main__":
        main()