Skip to content

Instantly share code, notes, and snippets.

@Alives
Last active November 18, 2020 19:24
Show Gist options
  • Save Alives/c12d2918588607eb4ca98af0407aefb3 to your computer and use it in GitHub Desktop.
Save Alives/c12d2918588607eb4ca98af0407aefb3 to your computer and use it in GitHub Desktop.
#!/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')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment