Advertisement
AzraelNewtype

Parallel Encoding 8/10bit thing that may work?

Dec 9th, 2011
199
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 8.83 KB | None | 0 0
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8-*-
  3. """ Port of parallelencoding.cmd to python, cleaned from the ghastly batch
  4.    script standards for PEP8 compliance (not yet). Also adds 10bit
  5.    functionality.
  6. """
  7.  
  8. import optparse, os, re, shutil, subprocess, sys, tempfile, time
  9.  
  10. # user params
  11. avs2yuv_path = 'C:/enc_tools/avs2yuv.exe'
  12. x264_path = 'C:/enc_tools/x264.exe'
  13. x264_10_path =  'C:/enc_tools/x264-10bit.exe'
  14. ffms2_10bit_path = 'C:/encoding/oldplugins/ffms2-10bit.dll'
  15. ffms2_8bit_path = 'C:/avsfilters/ffms2.dll'
  16.  
  17. x264_extra_params='--preset ultrafast --subme 1 --input-depth 16'
  18.  
  19. def final_script_suffix(avs):
  20. """ Use this function to add things to the final script that you always
  21.    want but do not directly pertain to the split/join functionality.
  22. """
  23.     avs.write("spline64resize(848,480) #SD")
  24.     avs.write('scxvid("out.stats") #WR')
  25.     avs.write('#TextSub("{0}") #SD'.format(script_out_path))
  26.  
  27. # generate_parallel_avs creates trimmed files for parallel encoding
  28. def generate_parallel_avs(avs_out, main_avs, avs_mem, total_threads, thread_number):
  29.     parallel_avs = open(avs_out, 'w')
  30.     parallel_avs.write('SetMemoryMax({0})\n'.format(avs_mem))
  31.     parallel_avs.write('Import("{0}")\n'.format(main_avs))
  32.     parallel_avs.write('start = (FrameCount() / {0}) * {1}\n'.format(total_threads, thread_num - 1))
  33.     if(thread_num == total_threads):
  34.         parallel_avs.write('end = FrameCount()\n')
  35.     else:
  36.         parallel_avs.write('end = start + (FrameCount() / {0}) + 100\n'.format(total_threads))
  37.     parallel_avs.write('Trim(start,end)\n')
  38.  
  39. #generate_joined_avs joins lossless files from parallel encoding
  40. def generate_joined_avs(output_avs, lossless, avs_mem, total_threads, depth):
  41.     joined_avs = open(final_avs, 'w')
  42.     if depth == 10:
  43.         filter_path = os.path.normpath(ffms2_10bit_path)
  44.     else
  45.         filter_path = os.path.normpath(ffms2_8bit_path)
  46.     joined_avs.write('LoadPlugin("{0}")\n'.format(filter_path))
  47.     joined_avs.write('#SetMemoryMax({0})\n'.format(avs_mem))
  48.     for thread in range(1, total_threads + 1):
  49.         write_source_line(joined_avs, lossless, thread, depth)
  50.         if (thread == 1):
  51.             joined_avs.write('total1 = tmp.Trim(0,tmp.FrameCount() - 51)\n')
  52.         elif (thread == total_threads - 1):
  53.             #final thread
  54.             joined_avs.write('total1 = total1 + tmp.Trim(51,tmp.FrameCount())\n')
  55.         else:
  56.             joined_avs.write('total1 = total1 + tmp.Trim(51,tmp.FrameCount() - 51)\n')
  57.     joined_avs.write('total1\n')
  58.     final_script_suffix(joined_avs)
  59.  
  60. def write_source_line(avs, lossless, num, depth):
  61.     lossless_out = lossless.replace('[NUM]', str(num))
  62.     if depth == 10:
  63.         colorspace = 'YV12_10-bit_hack'
  64.     else:
  65.         colorspace = 'YV12'
  66.     avs.write('tmp = FFVideoSource("{0}",colorspace="{1}",track=-1,pp="")\n'.format(lossless_out,
  67.                                                                                     colorspace))
  68.  
  69.  
  70. #count_frames counts number of frames in AviSynth Script
  71. def count_frames(script_in, proc):
  72.     frame = [0, 0, 0, 0]
  73.     a,tempYUV=tempfile.mkstemp()
  74.     os.close(a)
  75.     frames_cmd = '"{0}" -raw -frames 1 "[input]" -o "{1}"'.format(os.path.normpath(avs2yuv_path), tempYUV)
  76.     if(options.usewine):
  77.         frames_cmd = 'wine ' + frames_cmd.replace('[input]', 'Z:' + script_in.replace('/','\\'))
  78.         proc = subprocess.Popen(frames_cmd,shell=True,stdout=subprocess.PIPE,universal_newlines=True,stderr=subprocess.STDOUT)
  79.     else:
  80.         frames_cmd = frames_cmd.replace('[input]', script_in)
  81.         proc = subprocess.Popen('"' + frames_cmd + '"',shell=True,stdout=subprocess.PIPE,universal_newlines=True,stderr=subprocess.STDOUT)
  82.     proc.wait()
  83.     p = re.compile ('.+: ([0-9]+)x([0-9]+), ([0-9]+/[0-9]+) fps, ([0-9]+) frames')
  84.     result = False
  85.     for line in proc.stdout:
  86.         m = p.search(line)
  87.         if m:
  88.             frame = [m.group(1), m.group(2), m.group(3), m.group(4)]
  89.             result = True
  90.             break
  91.     if not result:
  92.         print('Error: Could not count number of frames.')
  93.         frame[3] = -1
  94.     os.unlink(tempYUV)
  95.     return (frame, proc)
  96.  
  97. parser = optparse.OptionParser(usage="usage: %prog [options] input.avs")
  98. parser.add_option('-t', '--threads', type='int', dest='threads', default=4,
  99.                   help="Number of parallel encodes to spawn")
  100. parser.add_option('-m', '--max-memory', type='int', dest='AVS_Mem_Per_Thread', default=512,
  101.                   help="Value for SetMemoryMax() in threads")
  102. parser.add_option('-w', '--wine', action='store_true', dest='usewine', default=False,
  103.                   help="Encoding on linux, so use wine")
  104. parser.add_option('-n', '--no-avs2yuv', action='store_false', dest='useavs2yuv', default=True,
  105.                   help="Do not use avs2yuv. Strange default action requires explicitly turning off.")
  106. parser.add_option('-d', '--depth', type='int', dest='depth', default=8,
  107.                   help="Bit depth for lossless [default=8]")
  108. (options, args) = parser.parse_args()
  109.  
  110. if(options.usewine):
  111.     options.useavs2yuv = True
  112.  
  113. if (options.threads < 2):
  114.     print("I'm sorry, but there is currently no special case for a single thread.")
  115.     raise SystemExit
  116.  
  117. total_threads = options.threads
  118. avs_mem = options.AVS_Mem_Per_Thread
  119.  
  120. if len(args) < 1:
  121.     print('No input file given. Use -h or --help for usage.')
  122.     raise SystemExit
  123.  
  124. # Looping for multiple input is kinda stupid, loop elsewhere if you want multiple input
  125. infile = args[0]
  126. (fname, ext) = os.path.splitext(infile)
  127. # ensure positional argument is a real avs file.
  128. if(not os.path.exists(infile) or not os.path.isfile(infile) or ext.lower().find('.avs') == -1):
  129.     print('Input file does not exist or is not an avisynth script.')
  130.     raise SystemExit
  131.  
  132. # script vars
  133. avs_in = os.path.abspath(infile)
  134. proj_name = os.path.basename(fname)
  135. script_out_path = os.path.dirname(os.path.abspath(infile)) + os.sep + proj_name + os.sep
  136. split_script = script_out_path + proj_name + '.[NUM].avs'
  137. final_avs = script_out_path + proj_name + '.joined.avs'
  138. lossless_path = script_out_path + 'Lossless' + os.sep
  139. split_output = lossless_path + proj_name + '.[NUM].mkv'
  140. depth = options.depth
  141.  
  142. # remove old scripts and files before script is run
  143. if os.path.isdir(script_out_path):
  144.     shutil.rmtree(script_out_path)
  145.  
  146. # check to make sure dirs exist
  147. for dir in (script_out_path, lossless_path):
  148.     if not os.path.isdir(dir):
  149.         os.makedirs(dir)
  150.  
  151. # create trimmed scripts
  152. print('Creating trimmed scripts.')
  153. if options.useavs2yuv:
  154.     split_script_frames = list(range(total_threads))
  155.     proc = list(range(total_threads))
  156. for thread in range(1,total_threads + 1):
  157.     new_split_script = split_script.replace('[NUM]', str(thread))
  158.     if(options.usewine):
  159.         generate_parallel_avs(new_split_script,'Z:' + avs_in.replace('/','\\'),avs_mem,total_threads,_CurrentThread)
  160.     else:
  161.         generate_parallel_avs(new_split_script, avs_in, avs_mem, total_threads, thread)
  162.     if(options.useavs2yuv):
  163.         print('Counting frames in script {0}'.format(thread))
  164.         (split_script_frames[thread-1],proc[thread-1]) = count_frames(new_split_script, proc[thread-1])
  165.  
  166. if(options.useavs2yuv):
  167.     for thread in range(total_threads):
  168.         proc[thread].wait()
  169.  
  170. # create cmd batch files
  171. enc_cmd = ''
  172. cmd_input = '"[input]"'
  173. cmd_output = '"[output]"'
  174. if(options.useavs2yuv):
  175.     enc_cmd = enc_cmd + '"' + os.path.normpath(avs2yuv_path) + '" -raw ' + cmd_input + ' -o - | '
  176.     cmd_input = '--frames [frames] --demuxer raw -'
  177.     if(split_script_frames[0][3] > -1):
  178.         cmd_input = '--input-res ' + str(int(split_script_frames[0][0])//2) + 'x' + split_script_frames[0][1] + ' --fps ' + split_script_frames[0][2] + ' ' + cmd_input
  179. enc_cmd = enc_cmd + '"' + os.path.normpath(x264_path) + '" ' + x264_extra_params + ' --crf 0 --threads 1 --thread-input --output ' + cmd_output + ' ' + cmd_input
  180. proc = list(range(total_threads))
  181. for thread in range(1, total_threads + 1):
  182.     new_split_script = split_script.replace('[NUM]', str(thread))
  183.     new_lossless = split_output.replace('[NUM]', str(thread))
  184.     new_cmd = enc_cmd.replace('[input]', new_split_script)
  185.     new_cmd = new_cmd.replace('[output]', new_lossless)
  186.     if(options.useavs2yuv):
  187.         new_cmd = new_cmd.replace('[frames]', split_script_frames[thread-1][3])
  188.     if(options.usewine):
  189.         new_cmd = 'wine ' + new_cmd
  190.         print(new_cmd + '\n')
  191.         proc[thread] = subprocess.Popen(new_cmd,shell=True)
  192.     else:
  193.         print(new_cmd + '\n')
  194.         proc[thread] = subprocess.Popen('"' + new_cmd + '"',shell=True)
  195.  
  196. for thread in range(total_threads):
  197.     proc[thread].wait()
  198.  
  199. # create joined lossless script
  200. print('Generating joined script.')
  201.  
  202.  
  203. if(options.usewine):
  204.     split_output = 'Z:' + split_output.replace('/','\\')
  205.  
  206. generate_joined_avs(final_avs, split_output, avs_mem, total_threads, depth)
  207.  
  208.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement