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