src/youtube/utils.py

   1 import hashlib
   2 import os
   3 import shutil
   4 import subprocess
   5 from tempfile import NamedTemporaryFile
   6 from django.conf import settings
   7
   8
   9 FILE_CACHE = getattr(settings, 'FILE_CACHE', 'file_cache/')
  10
  11
  12 def link_or_copy(src, dst):
  13     dstdir = os.path.dirname(dst)
  14     if not os.path.exists(dstdir):
  15         os.makedirs(dstdir)
  16     if os.path.exists(dst):
  17         os.unlink(dst)
  18         # FIXME: tiny window here when the temp path is not taken.
  19     try:
  20         os.link(src, dst)
  21     except OSError:
  22         shutil.copyfile(src, dst)
  23
  24
  25 def process_to_file(cmdline, prefix='', suffix='', cache_key=None, output_path=None):
  26     if not output_path:
  27         tmp = NamedTemporaryFile(
  28             prefix=prefix, suffix=suffix, delete=False,
  29             dir=settings.FILE_UPLOAD_TEMP_DIR
  30         )
  31         tmp.close()
  32         output_path = tmp.name
  33
  34     if cache_key:
  35         cache_path = cache_key.replace('/', '__')
  36         if len(cache_path) > 255:
  37             parts = cache_path.rsplit('.', 1)
  38             limit = 255 - 9
  39             if len(parts) > 1:
  40                 limit -= len(parts[1]) + 1
  41             cache_path = parts[0][:limit] + '.' + hashlib.sha1(cache_key.encode('utf-8')).hexdigest()[:8]
  42             if len(parts) > 1:
  43                 cache_path += '.' + parts[1]
  44         cache_path = FILE_CACHE + cache_path
  45
  46     if cache_key and os.path.exists(cache_path):
  47         link_or_copy(cache_path, output_path)
  48     else:
  49         # Actually run the processing.
  50         subprocess.run(cmdline + [output_path], check=True)
  51         if cache_key:
  52             link_or_copy(output_path, cache_path)
  53
  54     return output_path
  55
  56
  57 def video_from_image(img_path, duration, fps=25, cache=True):
  58     return process_to_file(
  59         ['ffmpeg', '-y', '-loop', '1', '-t', str(duration), '-i', img_path, '-c:v', 'libx264', '-vf', f'fps={fps},format=yuv420p'],
  60         'image-',
  61         '.mkv',
  62         f'video_from_image:{img_path}:{duration}:{fps}.mkv' if cache else None
  63     )
  64
  65
  66 def cut_video(video_path, duration):
  67     return process_to_file(
  68         ['ffmpeg', '-y', '-i', video_path, '-t', str(duration), '-c', 'copy'],
  69         'cut-',
  70         '.mkv'
  71     )
  72
  73
  74 def ffmpeg_concat(paths, suffix, copy=False):
  75     filelist = NamedTemporaryFile(
  76         prefix='concat-', suffix='.txt',
  77         dir=settings.FILE_UPLOAD_TEMP_DIR
  78     )
  79     for path in paths:
  80         filelist.write(f"file '{path}'\n".encode('utf-8'))
  81     filelist.flush()
  82
  83     args = ['ffmpeg', '-y', '-safe', '0', '-f', 'concat', '-i', filelist.name]
  84     if copy:
  85         args += ['-c', 'copy']
  86     outname = process_to_file(args, 'concat-', suffix)
  87
  88     filelist.close()
  89     return outname
  90
  91
  92 def concat_videos(paths):
  93     return ffmpeg_concat(paths, '.mkv', copy=True)
  94
  95
  96 def concat_audio(paths):
  97     return ffmpeg_concat(paths, '.flac')
  98
  99
 100 def standardize_audio(p, cache=True):
 101     return process_to_file(
 102         ['ffmpeg', '-y', '-i', p, '-sample_fmt', 's16', '-acodec', 'flac', '-ac', '2', '-ar', '44100'],
 103         'standardize-', '.flac',
 104         f'standardize_audio:{p}.flac' if cache else None
 105     )
 106
 107
 108 def standardize_video(p, cache=True):
 109     return process_to_file(
 110         ['ffmpeg', '-y', '-i', p],
 111         'standardize-', '.mkv',
 112         f'standardize_video:{p}.mkv' if cache else None
 113     )
 114
 115
 116 def mux(channels, output_path=None):
 117     args = ['ffmpeg', '-y']
 118     for c in channels:
 119         args.extend(['-i', c])
 120     args.extend(['-c', 'copy'])
 121     return process_to_file(args, 'mux-', '.mkv', output_path=output_path)
 122
 123
 124 def get_duration(path):
 125     return float(
 126         subprocess.run(
 127             [
 128                 "ffprobe",
 129                 "-i",
 130                 path,
 131                 "-show_entries",
 132                 "format=duration",
 133                 "-v",
 134                 "quiet",
 135                 "-of",
 136                 "csv=p=0",
 137             ],
 138             capture_output=True,
 139             text=True,
 140             check=True,
 141         ).stdout
 142     )
 143
 144
 145 def get_framerate(path):
 146     rates = subprocess.run(
 147             [
 148                 "ffprobe",
 149                 "-i",
 150                 path,
 151                 "-show_entries",
 152                 "stream=r_frame_rate",
 153                 "-v",
 154                 "quiet",
 155                 "-of",
 156                 "csv=p=0",
 157             ],
 158             capture_output=True,
 159             text=True,
 160             check=True,
 161         ).stdout.strip().split('\n')
 162     for rate in rates:
 163         a, b = rate.split('/')
 164         if b == '1':
 165             return int(a)