yt_dlp/postprocessor/embedthumbnail.py

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