From 95131b2176bbfcc070aa0d35845ed02987516ea1 Mon Sep 17 00:00:00 2001 From: pukkandan Date: Thu, 20 May 2021 16:55:57 +0530 Subject: [PATCH] [embedthumbnail] Add `flac` support and refactor `mutagen` code https://github.com/ytdl-org/youtube-dl/pull/28894, https://github.com/ytdl-org/youtube-dl/pull/24310 Authored by: tripulse --- yt_dlp/postprocessor/embedthumbnail.py | 70 ++++++++++++++++---------- 1 file changed, 44 insertions(+), 26 deletions(-) diff --git a/yt_dlp/postprocessor/embedthumbnail.py b/yt_dlp/postprocessor/embedthumbnail.py index 55551e92c..d3fd2cafd 100644 --- a/yt_dlp/postprocessor/embedthumbnail.py +++ b/yt_dlp/postprocessor/embedthumbnail.py @@ -1,14 +1,17 @@ # coding: utf-8 from __future__ import unicode_literals -import os -import subprocess -import struct -import re import base64 +import imghdr +import os +import struct +import subprocess +import re try: - import mutagen + from mutagen.oggvorbis import OggVorbis + from mutagen.oggopus import OggOpus + from mutagen.flac import Picture, FLAC has_mutagen = True except ImportError: has_mutagen = False @@ -39,6 +42,23 @@ def __init__(self, downloader=None, already_have_thumbnail=False): FFmpegPostProcessor.__init__(self, downloader) self._already_have_thumbnail = already_have_thumbnail + def _get_thumbnail_resolution(self, filename, thumbnail_dict): + def guess(): + width, height = thumbnail_dict.get('width'), thumbnail_dict.get('height') + if width and height: + return width, height + + try: + size_regex = r',\s*(?P\d+)x(?P\d+)\s*[,\[]' + size_result = self.run_ffmpeg(filename, filename, ['-hide_banner']) + mobj = re.search(size_regex, size_result) + if mobj is None: + return guess() + except PostProcessingError as err: + self.report_warning('unable to find the thumbnail resolution; %s' % error_to_compat_str(err)) + return guess() + return int(mobj.group('w')), int(mobj.group('h')) + def run(self, info): filename = info['filepath'] temp_filename = prepend_extension(filename, 'temp') @@ -135,34 +155,32 @@ def run(self, info): self.report_warning('The file format doesn\'t support embedding a thumbnail') success = False - elif info['ext'] in ['ogg', 'opus']: + elif info['ext'] in ['ogg', 'opus', 'flac']: if not has_mutagen: raise EmbedThumbnailPPError('module mutagen was not found. Please install using `python -m pip install mutagen`') + self.to_screen('Adding thumbnail to "%s"' % filename) - - size_regex = r',\s*(?P\d+)x(?P\d+)\s*[,\[]' - size_result = self.run_ffmpeg(thumbnail_filename, thumbnail_filename, ['-hide_banner']) - mobj = re.search(size_regex, size_result) - width, height = int(mobj.group('w')), int(mobj.group('h')) - mimetype = ('image/%s' % ('png' if thumbnail_ext == 'png' else 'jpeg')).encode('ascii') - - # https://xiph.org/flac/format.html#metadata_block_picture - data = bytearray() - data += struct.pack('>II', 3, len(mimetype)) - data += mimetype - data += struct.pack('>IIIIII', 0, width, height, 8, 0, os.stat(thumbnail_filename).st_size) # 32 if png else 24 - - fin = open(thumbnail_filename, "rb") - data += fin.read() - fin.close() - temp_filename = filename - f = mutagen.File(temp_filename) - f.tags['METADATA_BLOCK_PICTURE'] = base64.b64encode(data).decode('ascii') + f = {'opus': OggOpus, 'flac': FLAC, 'ogg': OggVorbis}[info['ext']](filename) + + pic = Picture() + pic.mime = 'image/%s' % imghdr.what(thumbnail_filename) + with open(thumbnail_filename, 'rb') as thumbfile: + pic.data = thumbfile.read() + pic.type = 3 # front cover + res = self._get_thumbnail_resolution(thumbnail_filename, info['thumbnails'][-1]) + if res is not None: + pic.width, pic.height = res + + if info['ext'] == 'flac': + f.add_picture(pic) + else: + # https://wiki.xiph.org/VorbisComment#METADATA_BLOCK_PICTURE + f['METADATA_BLOCK_PICTURE'] = base64.b64encode(pic.write()).decode('ascii') f.save() else: - raise EmbedThumbnailPPError('Supported filetypes for thumbnail embedding are: mp3, mkv/mka, ogg/opus, m4a/mp4/mov') + raise EmbedThumbnailPPError('Supported filetypes for thumbnail embedding are: mp3, mkv/mka, ogg/opus/flac, m4a/mp4/mov') if success and temp_filename != filename: os.remove(encodeFilename(filename))