]>
Commit | Line | Data |
---|---|---|
1 | import base64 | |
2 | import os | |
3 | import re | |
4 | import subprocess | |
5 | ||
6 | from .common import PostProcessor | |
7 | from .ffmpeg import FFmpegPostProcessor, FFmpegThumbnailsConvertorPP | |
8 | from ..compat import imghdr | |
9 | from ..dependencies import mutagen | |
10 | from ..utils import ( | |
11 | Popen, | |
12 | PostProcessingError, | |
13 | check_executable, | |
14 | encodeArgument, | |
15 | encodeFilename, | |
16 | error_to_compat_str, | |
17 | prepend_extension, | |
18 | shell_quote, | |
19 | ) | |
20 | ||
21 | if mutagen: | |
22 | from mutagen.flac import FLAC, Picture | |
23 | from mutagen.mp4 import MP4, MP4Cover | |
24 | from mutagen.oggopus import OggOpus | |
25 | from mutagen.oggvorbis import OggVorbis | |
26 | ||
27 | ||
28 | class EmbedThumbnailPPError(PostProcessingError): | |
29 | pass | |
30 | ||
31 | ||
32 | class EmbedThumbnailPP(FFmpegPostProcessor): | |
33 | ||
34 | def __init__(self, downloader=None, already_have_thumbnail=False): | |
35 | FFmpegPostProcessor.__init__(self, downloader) | |
36 | self._already_have_thumbnail = already_have_thumbnail | |
37 | ||
38 | def _get_thumbnail_resolution(self, filename, thumbnail_dict): | |
39 | def guess(): | |
40 | width, height = thumbnail_dict.get('width'), thumbnail_dict.get('height') | |
41 | if width and height: | |
42 | return width, height | |
43 | ||
44 | try: | |
45 | size_regex = r',\s*(?P<w>\d+)x(?P<h>\d+)\s*[,\[]' | |
46 | size_result = self.run_ffmpeg(filename, None, ['-hide_banner'], expected_retcodes=(1,)) | |
47 | mobj = re.search(size_regex, size_result) | |
48 | if mobj is None: | |
49 | return guess() | |
50 | except PostProcessingError as err: | |
51 | self.report_warning('unable to find the thumbnail resolution; %s' % error_to_compat_str(err)) | |
52 | return guess() | |
53 | return int(mobj.group('w')), int(mobj.group('h')) | |
54 | ||
55 | def _report_run(self, exe, filename): | |
56 | self.to_screen(f'{exe}: Adding thumbnail to "{filename}"') | |
57 | ||
58 | @PostProcessor._restrict_to(images=False) | |
59 | def run(self, info): | |
60 | filename = info['filepath'] | |
61 | temp_filename = prepend_extension(filename, 'temp') | |
62 | ||
63 | if not info.get('thumbnails'): | |
64 | self.to_screen('There aren\'t any thumbnails to embed') | |
65 | return [], info | |
66 | ||
67 | idx = next((-i for i, t in enumerate(info['thumbnails'][::-1], 1) if t.get('filepath')), None) | |
68 | if idx is None: | |
69 | self.to_screen('There are no thumbnails on disk') | |
70 | return [], info | |
71 | thumbnail_filename = info['thumbnails'][idx]['filepath'] | |
72 | if not os.path.exists(encodeFilename(thumbnail_filename)): | |
73 | self.report_warning('Skipping embedding the thumbnail because the file is missing.') | |
74 | return [], info | |
75 | ||
76 | # Correct extension for WebP file with wrong extension (see #25687, #25717) | |
77 | convertor = FFmpegThumbnailsConvertorPP(self._downloader) | |
78 | convertor.fixup_webp(info, idx) | |
79 | ||
80 | original_thumbnail = thumbnail_filename = info['thumbnails'][idx]['filepath'] | |
81 | ||
82 | # Convert unsupported thumbnail formats (see #25687, #25717) | |
83 | # PNG is preferred since JPEG is lossy | |
84 | thumbnail_ext = os.path.splitext(thumbnail_filename)[1][1:] | |
85 | if info['ext'] not in ('mkv', 'mka') and thumbnail_ext not in ('jpg', 'jpeg', 'png'): | |
86 | thumbnail_filename = convertor.convert_thumbnail(thumbnail_filename, 'png') | |
87 | thumbnail_ext = 'png' | |
88 | ||
89 | mtime = os.stat(encodeFilename(filename)).st_mtime | |
90 | ||
91 | success = True | |
92 | if info['ext'] == 'mp3': | |
93 | options = [ | |
94 | '-c', 'copy', '-map', '0:0', '-map', '1:0', '-write_id3v1', '1', '-id3v2_version', '3', | |
95 | '-metadata:s:v', 'title="Album cover"', '-metadata:s:v', 'comment=Cover (front)'] | |
96 | ||
97 | self._report_run('ffmpeg', filename) | |
98 | self.run_ffmpeg_multiple_files([filename, thumbnail_filename], temp_filename, options) | |
99 | ||
100 | elif info['ext'] in ['mkv', 'mka']: | |
101 | options = list(self.stream_copy_opts()) | |
102 | ||
103 | mimetype = f'image/{thumbnail_ext.replace("jpg", "jpeg")}' | |
104 | old_stream, new_stream = self.get_stream_number( | |
105 | filename, ('tags', 'mimetype'), mimetype) | |
106 | if old_stream is not None: | |
107 | options.extend(['-map', '-0:%d' % old_stream]) | |
108 | new_stream -= 1 | |
109 | options.extend([ | |
110 | '-attach', self._ffmpeg_filename_argument(thumbnail_filename), | |
111 | '-metadata:s:%d' % new_stream, 'mimetype=%s' % mimetype, | |
112 | '-metadata:s:%d' % new_stream, 'filename=cover.%s' % thumbnail_ext]) | |
113 | ||
114 | self._report_run('ffmpeg', filename) | |
115 | self.run_ffmpeg(filename, temp_filename, options) | |
116 | ||
117 | elif info['ext'] in ['m4a', 'mp4', 'm4v', 'mov']: | |
118 | prefer_atomicparsley = 'embed-thumbnail-atomicparsley' in self.get_param('compat_opts', []) | |
119 | # Method 1: Use mutagen | |
120 | if not mutagen or prefer_atomicparsley: | |
121 | success = False | |
122 | else: | |
123 | try: | |
124 | self._report_run('mutagen', filename) | |
125 | meta = MP4(filename) | |
126 | # NOTE: the 'covr' atom is a non-standard MPEG-4 atom, | |
127 | # Apple iTunes 'M4A' files include the 'moov.udta.meta.ilst' atom. | |
128 | f = {'jpeg': MP4Cover.FORMAT_JPEG, 'png': MP4Cover.FORMAT_PNG}[imghdr.what(thumbnail_filename)] | |
129 | with open(thumbnail_filename, 'rb') as thumbfile: | |
130 | thumb_data = thumbfile.read() | |
131 | meta.tags['covr'] = [MP4Cover(data=thumb_data, imageformat=f)] | |
132 | meta.save() | |
133 | temp_filename = filename | |
134 | except Exception as err: | |
135 | self.report_warning('unable to embed using mutagen; %s' % error_to_compat_str(err)) | |
136 | success = False | |
137 | ||
138 | # Method 2: Use AtomicParsley | |
139 | if not success: | |
140 | success = True | |
141 | atomicparsley = next(( | |
142 | # libatomicparsley.so : See https://github.com/xibr/ytdlp-lazy/issues/1 | |
143 | x for x in ['AtomicParsley', 'atomicparsley', 'libatomicparsley.so'] | |
144 | if check_executable(x, ['-v'])), None) | |
145 | if atomicparsley is None: | |
146 | self.to_screen('Neither mutagen nor AtomicParsley was found. Falling back to ffmpeg') | |
147 | success = False | |
148 | else: | |
149 | if not prefer_atomicparsley: | |
150 | self.to_screen('mutagen was not found. Falling back to AtomicParsley') | |
151 | cmd = [encodeFilename(atomicparsley, True), | |
152 | encodeFilename(filename, True), | |
153 | encodeArgument('--artwork'), | |
154 | encodeFilename(thumbnail_filename, True), | |
155 | encodeArgument('-o'), | |
156 | encodeFilename(temp_filename, True)] | |
157 | cmd += [encodeArgument(o) for o in self._configuration_args('AtomicParsley')] | |
158 | ||
159 | self._report_run('atomicparsley', filename) | |
160 | self.write_debug('AtomicParsley command line: %s' % shell_quote(cmd)) | |
161 | stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) | |
162 | if returncode: | |
163 | self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {stderr.strip()}') | |
164 | # for formats that don't support thumbnails (like 3gp) AtomicParsley | |
165 | # won't create to the temporary file | |
166 | if 'No changes' in stdout: | |
167 | self.report_warning('The file format doesn\'t support embedding a thumbnail') | |
168 | success = False | |
169 | ||
170 | # Method 3: Use ffmpeg+ffprobe | |
171 | # Thumbnails attached using this method doesn't show up as cover in some cases | |
172 | # See https://github.com/yt-dlp/yt-dlp/issues/2125, https://github.com/yt-dlp/yt-dlp/issues/411 | |
173 | if not success: | |
174 | success = True | |
175 | try: | |
176 | options = [*self.stream_copy_opts(), '-map', '1'] | |
177 | ||
178 | old_stream, new_stream = self.get_stream_number( | |
179 | filename, ('disposition', 'attached_pic'), 1) | |
180 | if old_stream is not None: | |
181 | options.extend(['-map', '-0:%d' % old_stream]) | |
182 | new_stream -= 1 | |
183 | options.extend(['-disposition:%s' % new_stream, 'attached_pic']) | |
184 | ||
185 | self._report_run('ffmpeg', filename) | |
186 | self.run_ffmpeg_multiple_files([filename, thumbnail_filename], temp_filename, options) | |
187 | except PostProcessingError as err: | |
188 | success = False | |
189 | raise EmbedThumbnailPPError(f'Unable to embed using ffprobe & ffmpeg; {err}') | |
190 | ||
191 | elif info['ext'] in ['ogg', 'opus', 'flac']: | |
192 | if not mutagen: | |
193 | raise EmbedThumbnailPPError('module mutagen was not found. Please install using `python -m pip install mutagen`') | |
194 | ||
195 | self._report_run('mutagen', filename) | |
196 | f = {'opus': OggOpus, 'flac': FLAC, 'ogg': OggVorbis}[info['ext']](filename) | |
197 | ||
198 | pic = Picture() | |
199 | pic.mime = 'image/%s' % imghdr.what(thumbnail_filename) | |
200 | with open(thumbnail_filename, 'rb') as thumbfile: | |
201 | pic.data = thumbfile.read() | |
202 | pic.type = 3 # front cover | |
203 | res = self._get_thumbnail_resolution(thumbnail_filename, info['thumbnails'][idx]) | |
204 | if res is not None: | |
205 | pic.width, pic.height = res | |
206 | ||
207 | if info['ext'] == 'flac': | |
208 | f.add_picture(pic) | |
209 | else: | |
210 | # https://wiki.xiph.org/VorbisComment#METADATA_BLOCK_PICTURE | |
211 | f['METADATA_BLOCK_PICTURE'] = base64.b64encode(pic.write()).decode('ascii') | |
212 | f.save() | |
213 | temp_filename = filename | |
214 | ||
215 | else: | |
216 | raise EmbedThumbnailPPError('Supported filetypes for thumbnail embedding are: mp3, mkv/mka, ogg/opus/flac, m4a/mp4/m4v/mov') | |
217 | ||
218 | if success and temp_filename != filename: | |
219 | os.replace(temp_filename, filename) | |
220 | ||
221 | self.try_utime(filename, mtime, mtime) | |
222 | converted = original_thumbnail != thumbnail_filename | |
223 | self._delete_downloaded_files( | |
224 | thumbnail_filename if converted or not self._already_have_thumbnail else None, | |
225 | original_thumbnail if converted and not self._already_have_thumbnail else None, | |
226 | info=info) | |
227 | return [], info |