d219dde8c1e4842e8a64a8bf40e42484cdf162e0
[audio.git] / 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)