SHARE
TWEET

Untitled

a guest Aug 19th, 2017 45 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #twopassqcomp.py version
  2. twppassqcompver = '0.0.1'
  3.  
  4. import sys, os, subprocess, shlex, argparse, time
  5.  
  6. def log(line,logfile='twopassqcomp.log'):
  7.     with open(logfile,'a') as l:
  8.         l.write(line+'\n')
  9.     l.closed
  10.     print line
  11.  
  12. def setrange(txt,step=0,count=0,seek=0):
  13.     with open('test.avs', 'w') as out:
  14.         for line in txt:
  15.             out.write(line)
  16.         if not (step == 0 or count == 0):
  17.             out.write('\nSelectRangeEvery('+str(step)+','+str(count)+','+str(seek)+')')
  18.             log('\nSetting SelectRangeEvery('+str(step)+','+str(count)+','+str(seek)+')\n')
  19.     out.closed
  20.    
  21. def readavs(source):
  22.     """Read the .avs file into memory so we can use it to create our own test.avs"""
  23.     res = []
  24.     with open(source,'r') as avs:
  25.         while (True):
  26.             line = avs.readline()
  27.             if line == '':
  28.                 break
  29.             elif (line[:16].lower() == 'selectrangeevery'):
  30.                 log('\nYOUR .AVS FILE CONTAINS A SELECTRANGEEVERY() LINE.  REMOVE IT AND TRY AGAIN.')
  31.                 exit()
  32.             else:
  33.                 res.append(line)
  34.     avs.closed
  35.     return res
  36.    
  37. def floatrange(i,j,k):
  38.     x = range(int(100*i),int(100*j),int(100*k))
  39.     y = list([z/100.0 for z in x])
  40.     return y
  41.  
  42. class TwoPassTests:
  43.     def __init__(self, keep=False, parameters={}):
  44.         self.keeptests = keep
  45.         self.results = {}
  46.         self.params = parameters
  47.         with open('twopassqcomp.log','w') as logfile:
  48.             logfile.write('2passqcomp version '+twppassqcompver+'\n')
  49.         logfile.closed
  50.     def set(self, p, v=''):
  51.         self.params[p] = str(v)
  52.     def rem(self, p):
  53.         if p in self.params:
  54.             del self.params[p]
  55.     def setrange(txt,step=0,count=0,seek=0):
  56.         with open('test.avs', 'w') as out:
  57.             for line in txt:
  58.                 out.write(line)
  59.             if not (step == 0 or count == 0):
  60.                 out.write('\nSelectRangeEvery('+str(step)+','+str(count)+','+str(seek)+')')
  61.                 log('\nSetting SelectRangeEvery('+str(step)+','+str(count)+','+str(seek)+')\n')
  62.         out.closed
  63.     def string(self):
  64.         param_string = ''
  65.         for k,v in self.params.iteritems():
  66.             param_string += ' --' + k
  67.             if v != '': param_string += ' ' + v
  68.         return param_string
  69.     def override(self,setting,tests,msg=''):
  70.         if setting in tests:
  71.             print 'Input a value to use for '+setting+', or press ENTER to use the computed value.'
  72.             override = raw_input()
  73.             print
  74.             if override != '':
  75.                 self.set(setting,override)
  76.                 log('User chose '+setting+' '+override)
  77.         else:
  78.             override = ''
  79.             while override == '':
  80.                 print 'Input a value to use for '+setting+'. '+msg
  81.                 override = raw_input()
  82.             self.set(setting,override)
  83.             log('User chose '+setting+' '+override)
  84.     def test_encode(self,outfile='test.mkv',use_previous=True,verbose=True):
  85.         self.rem('output')
  86.         if self.string() in self.results and use_previous:
  87.             info = self.results[self.string()]
  88.             log('Using previously-recorded data from the following parameters:\n' + self.string())
  89.             log('\nPreviously-recorded results:\nframes:\t\t' + str(info['frames']) +
  90.                 '\nssim:\t\t' + str(info['ssim']) + '\ndb:\t\t' + str(info['db']) + '\nbitrate:\t' +
  91.                 str(info['bitrate']) + '\nfps:\t\t' + str(info['fps']) + '\n')
  92.             return info
  93.         if self.keeptests or outfile.startswith('psy-rd'):
  94.             self.set('output',outfile)
  95.         else:
  96.             self.set('output','test.mkv')
  97.         info = {}            
  98.         if verbose: log('Starting test encode with the following parameters:\n' + self.string())
  99.         with open('workfile.txt', 'w') as workfile:
  100.             param_string = self.string()
  101.             param_string = 'pipebuf avs2yuv test.avs -o - : x264 - --stdin y4m' + param_string
  102.             args = shlex.split(param_string)
  103.             ###Find the current time (which is a funky long number that docs.python.org chapter 15 describes as "Return the time as a floating point number expressed in seconds since the epoch, in UTC." We will then subtract the end time to find duration of the test.###
  104.             start = time.time()
  105.             p = subprocess.call(args,bufsize=-1,stdout=workfile,stderr=subprocess.STDOUT)
  106.             ###Find the current time after the test is over to go with the time we found just before the test started so that we can compute how long this test took to run.###
  107.             end = time.time()
  108.             ###Subtract the end time from the start time (note that they are both measured in seconds) to find how long this test took to run and save it as info['testtime'].
  109.             info['testtime'] = (end - start)
  110.         workfile.closed
  111.         with open('workfile.txt', 'r') as workfile:
  112.             while (True):
  113.                 line = workfile.readline()
  114.                 if line == '':
  115.                     break
  116.                 elif 'consecutive B-frames' in line:
  117.                     bframe_list = [float(s[:-1]) for s in shlex.split(line)[4:]]
  118.                     info['bframes'] = [0,bframe_list]
  119.                     for i in range(len(bframe_list)):
  120.                         if bframe_list[i] >= 1.0: info['bframes'][0] = i
  121.                 elif 'SSIM Mean' in line:
  122.                     metric_list = shlex.split(line)
  123.                     info['ssim'] = float(metric_list[4][2:])
  124.                     info['db'] = float(metric_list[5][1:-3])
  125.                 elif line.startswith('encoded'):
  126.                     rate_list = shlex.split(line)
  127.                     info['frames'] = int(rate_list[1])
  128.                     info['fps'] = float(rate_list[3])
  129.                     info['bitrate'] = float(rate_list[5])
  130.         workfile.closed
  131.         if info == {}:
  132.             log('\nTEST ENCODE FAILED.')
  133.             exit()
  134.         log('\nFinished test encode with the following results:\nframes:\t\t' + str(info['frames']) + '\nssim:\t\t' + str(info['ssim']) +
  135.             '\ndb:\t\t' + str(info['db']) + '\nbitrate:\t' + str(info['bitrate']) +
  136.             '\nfps:\t\t' + str(info['fps']) + '\ntest time:\t' + str(info['testtime']) + '\nfile size:\t' + str(info['filesize']) + '\n')
  137.         self.rem('output')
  138.         self.results[self.string()] = info
  139.         return info
  140.  
  141. def twopass(x264,q):
  142.     x264.set('qcomp',q)
  143.     log('\nRunning twopass tests with qcomp = '+str(q)+'...\n')
  144.  
  145.     tests = []
  146.     #test bitrate in the range of start to end, in increments of step.
  147.     ##############################
  148.     ##     BITRATE TEST VALUES  ##
  149.     ##############################
  150.     start = 2541
  151.     end = 3141
  152.     step = -150
  153.     ##############################
  154.     for br in floatrange(start, end+step, step):
  155.         x264.set('bitrate',br)
  156.         for p in ['1','2']:
  157.             self.set('pass',p)
  158.             log('bitrate (in Kbps) = ' + str(br) + ', pass = ' + p)
  159.             test = x264.test_encode('qcomp_'+str(q)+'_bitrate_'+str(br)+'.mkv')
  160.             with open('twopassqcomp.txt','a') as csv:
  161.                 csv.write(str(q)+','+str(br)+','+str(p)+','+str(test['ssim'])+','+str(test['db'])+','+str(test['bitrate'])+'\n')
  162.             csv.closed
  163.  
  164. if __name__ == '__main__':
  165.     ################## CLI OPTIONS ##################
  166.     supported_tests = set(['twopass'])
  167.  
  168.     parser = argparse.ArgumentParser(description='me.py - Automated testing of various me encoding options', usage='%(prog)s infile.avs [options]')
  169.     parser.add_argument('--version', action='version', version='%(prog)s version '+twppassqcompver)
  170.     parser.add_argument('infile', help='source avisynth script')
  171.     parser.add_argument('--hd', action='store_true', help='set this flag if the source is high definition')
  172.     parser.add_argument('--sar', default=argparse.SUPPRESS, help='sample aspect ratio - mandatory for SD sources')
  173.     parser.add_argument('--keep', action='store_true', help = 'set this flag if you want to keep all test encodes for manual \
  174.                         comparison (this may take up lots of space)')
  175.     parser.add_argument('--only', metavar='test', dest='test', nargs='+', help='this option overrides the standard procedure \
  176.                     and only runs the specified tests')
  177.  
  178.     cli = vars(parser.parse_args())
  179.  
  180.  
  181.     ################## CLI CHECKS ##################
  182.     if not os.path.isfile(cli['infile']):
  183.         parser.exit(status=2, message='\nError: Input file ' + cli['infile'] + ' does not exist.')
  184.     elif not cli['infile'].endswith('.avs'):
  185.         parser.exit(status=2, message='\nError: Input file must be an avisynth script (*.avs).')
  186.     else:
  187.         source = cli['infile']
  188.  
  189.     if (not cli['hd']) and ('sar' not in cli):
  190.         parser.exit(status=2, message='\nError: SAR required for SD sources.')
  191.    
  192.     if cli['test'] is None:
  193.         cli['test'] = set(['twopass'])
  194.        
  195.     else:
  196.         for x in cli['test']:
  197.             if x not in supported_tests: parser.exit(status=2, message='\nError: unsupported test '+x+'.')
  198.         cli['test'] = set(cli['test'])
  199.  
  200.  
  201.     # read the avs file, then add quotes around it so we can use it in the command line
  202.     avstext = readavs(source)
  203.     source = '"' + source + '"'
  204.  
  205.  
  206.     ################# FIND WIDTH, HEIGHT, FRAME COUNT #################
  207.     with open('workfile.txt', 'w') as workfile:
  208.         args = 'pipebuf avs2yuv ' + source + ' -frames 1 -o - : x264 - --stdin y4m --output test.mkv'
  209.         args = shlex.split(args)
  210.         p = subprocess.call(args,bufsize=-1,stdout=workfile,stderr=subprocess.STDOUT)
  211.     workfile.closed
  212.     with open('workfile.txt', 'r') as workfile:
  213.         while (True):
  214.             line = workfile.readline()
  215.             if line == '':
  216.                 break
  217.             elif line.endswith('frames\n'):
  218.                 frames = int(line.split()[-2])
  219.                 res = line.split()[-5][:-1]
  220.                 res = res.split('x')
  221.                 width = int(res[0])
  222.                 height = int(res[1])
  223.     workfile.closed
  224.  
  225.  
  226.     ################# INITIAL X264 PARAMETERS #################
  227.     ref = min(16,(8388608/(width*height)))
  228.  
  229.     x264 = TwoPassTests(keep=cli['keep'])
  230.     if 'sar' in cli:
  231.         x264.set('sar',cli['sar'])
  232.     x264.set('ref',ref)
  233.     x264.set('ssim')
  234.     x264.set('direct','auto')
  235.     x264.set('bframes',7)
  236.     x264.set('b-adapt',2)
  237.     x264.set('b-pyramid','normal')
  238.     x264.set('psy-rd','1.0:0.0')
  239.     x264.set('deblock','-3:-3')
  240.     x264.set('subme',10)
  241.     x264.set('trellis',2)
  242.     x264.set('merange',24)
  243.     x264.set('aq-mode',2)
  244.     x264.set('rc-lookahead',250)
  245.     x264.set('analyse','all')
  246.     x264.set('no-fast-pskip')
  247.     x264.set('thread-input')
  248.     x264.set('passlogfile')
  249.     x264.set('2passtry.log')
  250.    
  251.     from config import *
  252.     config = startingOptions(x264)
  253.     trange = config[1]
  254.  
  255.    
  256.     ############### TEST ENCODING PROCEDURE #################
  257.     prologue = '\nSTARTING twopass/QCOMP TESTING PROCESS FOR:\n' + source + '\n\nCommand line options:\n'+str(sys.argv[1:])+'\n\nParsed options:\n'
  258.     for k,v in cli.iteritems():
  259.         prologue += k + ': ' + str(v) + '\n'
  260.     prologue += '\nSource width: '+str(width)+'\nSource height: '+str(height)+'\nFrame count: '+str(frames)+'\n\nx264 options:\n'+x264.string()+'\n'
  261.     log(prologue)
  262.  
  263.     ###If the 'twopassqcomp' test is set to run then make a .csv file called crfme.csv and write me.py version (with the version number) and infile (with the .avs file path) on the first line, labels on the 2nd line, and data on the lines after that (one line per test)###
  264.     if ('twopassqcomp' in cli['test']):
  265.         setrange(avstext,trange[0],trange[1],trange[2])
  266.         with open('twopassqcomp.csv','w') as csv:
  267.             csv.write('twopassqcomp.py version'+','+twppassqcompver+','+'infile'+','+source+'\nme,crf,pass,SSIM,db,bitrate,fps,test time(in seconds),file size(in bytes)\n')
  268.         csv.closed
  269.  
  270.     ##################
  271.     start = 0.7
  272.     end = 0.6
  273.     step = -0.05
  274.     ##################
  275.  
  276.     for q in floatrange(start,end+step,step):
  277.         x264.set('qcomp',q)
  278.         twopass(x264,q)
  279.  
  280.     ### DONE ###
  281.     x264.rem('ssim')
  282.     x264.rem('output')
  283.     log('\n\nTESTING COMPLETE!')
RAW Paste Data
Top