convert_hevc.py

#!/usr/bin/python3

import json
import os
import subprocess
import sys
import time


def get_container_name():
  num = '0'
  name = 'transcoder_'
  ps = subprocess.check_output(['docker', 'ps'], universal_newlines=True)
  for line in ps.splitlines():
    if name in line.split()[-1]:
      this_num = str(int(line.split()[-1].split('_')[-1]) + 1)
      if this_num > num:
        num = this_num
  return name + num


def get_conv_cmd(infile, outfile, stream_map):
  (filedir, filename) = get_filepaths(infile)
  maps = []
  for stream_id in stream_map.keys():
    maps.extend(['-map', f'0:{stream_id}'])
  cmd = [
      'docker', 'run', '-it', '--rm', '--device', '/dev/dri:/dev/dri',
      '-v', f'{filedir}:{filedir}', '--name', get_container_name(),
      '-e', 'AV_LOG_FORCE_NOCOLOR=1', 'jrottenberg/ffmpeg:4.1-vaapi'
  ] + [
      '-hide_banner', '-vaapi_device', '/dev/dri/renderD128', '-i',
      filename, '-vf', 'format=nv12,hwupload', '-c:v', 'hevc_vaapi',
      '-max_muxing_queue_size', '9999',
  ] + maps + [
      '-c:a', 'copy', '-c:s', 'copy', '-qp', '26', '-n', f'{filedir}/{outfile}']
  print(' '.join(cmd))
  return cmd


def get_filepaths(filename):
  if filename.startswith('/'):
    filedir = '/'.join(filename.split('/')[0:-1])
  else:
    filedir = os.getcwd()
    filename = filedir + '/' + filename
  return (filedir, filename)


def get_frames(filename):
  frames = 0
  (filedir, filename) = get_filepaths(filename)
  cmd = ['docker', 'run', '-it', '--rm', '-v', f'{filedir}:{filedir}:ro',
         '--name', get_container_name(), 'mediainfo:latest', 'mediainfo',
         '--Output=JSON', filename]
  #print(' '.join(cmd))
  streams = json.loads(subprocess.check_output(cmd))['media']['track']
  for stream in streams:
    if stream['@type'] == 'Video':
      frames = int(stream.get('FrameCount', 0))
      break
  if frames:
    print(f'Frames: {frames}')
  else:
    print('Unable to determine frame count.')
  return frames


def get_stream_map(filename):
  stream_map = {}
  print(filename)
  (filedir, filename) = get_filepaths(filename)
  cmd = ['docker', 'run', '-it', '--rm', '-v', f'{filedir}:{filedir}:ro',
         '--name', get_container_name(), '--entrypoint=ffprobe',
         'jrottenberg/ffmpeg:4.1-vaapi', '-v', 'quiet', '-print_format',
         'json', '-show_streams', filename]
  #print(' '.join(cmd))
  try:
    streams = json.loads(subprocess.check_output(cmd))['streams']
  except subprocess.CalledProcessError as error:
    print(' '.join(cmd))
    sys.exit(1)
  for num, stream in enumerate(streams):
    if stream['codec_type'] in ['audio', 'subtitle', 'video']:
      if stream['codec_type'] not in stream_map.values():
        if not stream.get('disposition', {'default': 0}).get('default', 0):
          stream_map[num] = stream['codec_type']
        elif stream.get('tags', {'language': ''}).get('language', 'eng') == 'eng':
          stream_map[num] = stream['codec_type']
        else:
          stream_map[num] = stream['codec_type']
  print('Stream mapping: %s' %
        ', '.join([f'{k}: {v}' for k, v in sorted(stream_map.items())]))
  if any([True for x in ['audio', 'video'] if x not in stream_map.values()]):
    print(' '.join(cmd))
    sys.exit(1)
  return stream_map


def get_outfile(infile):
  outdirs = infile.split('/')[:-1]
  infile = infile.split('/')[-1]
  outfile_split = []
  infile_split = infile.split('.')[:-1]
  infile_split = '_'.join('_'.join(infile_split).split('-')).split('_')
  for x in infile_split:
    if '265' in x or 'HEVC' in x:
      continue
    if '264' in x or 'xvid' in x.lower():
      outfile_split.append('HEVC')
    else:
      outfile_split.append(x)
  if 'HEVC' not in outfile_split:
    outfile_split.append('HEVC')
  outfile = '.'.join(outfile_split + ['mkv'])
  if outfile == infile:
    outfile_split = outfile.split('.')
    outfile = '.'.join(outfile[:-1] + ['new', outfile[-1]])
  if outdirs:
    outfile = '/'.join(outdirs + [outfile])
  if os.path.exists(outfile):
    print(f'{outfile} exists, overwrite? [Y/n] ', end='')
    if input().lower() in ['', 'y']:
      os.remove(outfile)
    else:
      sys.exit(0)
  return outfile


def get_pct(line, frames):
  if line.startswith('frame=') and 'fps=' in line:
    counted_frames = int(line.split('frame=')[-1].split('fps=')[0].strip())
    if frames:
      return float(counted_frames) / frames
  return 0


GREEN_FG = '\x1b[38;5;46m'
RED_FG = '\x1b[38;5;196m'
WHITE_FG = '\x1b[38;5;15m'
BLUE_BG = '\x1b[48;5;21m'
GRAY_BG = '\x1b[48;5;237m'
RESET = '\x1b[0m'

for filenum, infile in enumerate(sys.argv[1:]):
  if filenum <= len(sys.argv[1:]) - 1 and len(sys.argv[1:]) > 1:
    if filenum:
      print()
    sys.stdout.write(f'{filenum + 1}/{len(sys.argv[1:])} ')
  stream_summary = False
  line = ''
  max_line_len = 0
  pct = 0
  outfile = get_outfile(infile)
  stream_map = get_stream_map(infile)
  frames = get_frames(infile)
  start = time.mktime(time.localtime())
  p = subprocess.Popen(get_conv_cmd(infile, outfile, stream_map),
                       stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
  while True:
    if p.poll() is not None:
      break
    c = p.stdout.read(1).decode(errors='ignore')
    line += c
    if not c in '\r\n':
      continue
    if not line.startswith('f'):
      if not line.startswith('\n'):
        if line.lstrip().startswith('Input'):
          stream_summary = True
        if pct: # frame= processing done, output with a newline.
          sys.stdout.write('\n' + line)
          line = ''
          continue
        for heading in ['Input', 'Output', 'Stream']:
          if line.lstrip().startswith(heading):
            if heading != 'Stream':
              line = WHITE_FG + line + RESET
            sys.stdout.write(line)
            line = ''
            continue
        if (not line.startswith(' ') and
            'Press' not in line):
          sys.stdout.write(line + '\n')
          line = ''
    else:
      stream_summary = False
      if not frames:
        sys.stdout.write(line)
        continue
      pct = get_pct(line, frames)
      if pct:
        now = time.mktime(time.localtime())
        delta = now - start
        eta = (delta / pct) - delta
        line = '%s pct=%d%% eta=%02d:%02d%c' % (
            line[:-1].rstrip(), (pct*100), eta / 60, eta % 60, line[-1])
        color_change = int(pct * len(line.strip()))
        sys.stdout.write(BLUE_BG + WHITE_FG + line[:color_change])
        sys.stdout.write(GRAY_BG + line[color_change:-1] + RESET)
        line_len = len(line)
        erase = ''
        if line_len < max_line_len:
          # Erase any artifacts from line length changes.
          sys.stdout.write(' ' * (max_line_len - line_len))
        if line_len > max_line_len:
          max_line_len = line_len
        sys.stdout.write(line[-1])
    line = ''
  p.stdout.close()
  return_code = p.wait()
  if return_code:
    sys.exit(return_code)

  elapsed = time.mktime(time.localtime()) - start
  sys.stdout.write('Elapsed time: %02d:%02d\n' % (elapsed / 60, elapsed % 60))
  diff = get_frames(outfile) - frames
  if abs(diff) > 5 and frames:
    sys.stdout.write(f'{RED_FG}Conversion failed: expected {frames} frames '
                     f'(diff: {diff}).{RESET}\n')
  else:
    sys.stdout.write(f'{GREEN_FG}Conversion successful{RESET}\n')