Add moviepy as optional dep for HQ thumbnails, make Pillow optional

[db updated]

mautrix_telegram/util/file_transfer.py

@@ -19,11 +19,22 @@ import time
 import logging
 
 import magic
-from PIL import Image
 from sqlalchemy.exc import IntegrityError, InvalidRequestError
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+try:
+    from moviepy.editor import VideoFileClip
+    import random
+    import string
+    import os
+    import mimetypes
+except ImportError:
+    VideoFileClip = random = string = os = mimetypes = None
 
 from telethon_aio.tl.types import (Document, FileLocation, InputFileLocation,
-                                   InputDocumentFileLocation, PhotoCachedSize)
+                                   InputDocumentFileLocation, PhotoSize, PhotoCachedSize)
 from telethon_aio.errors import LocationInvalidError
 
 from ..db import TelegramFile as DBTelegramFile
@@ -32,24 +43,86 @@ log = logging.getLogger("mau.util")
 
 
 def _convert_webp(file, to="png"):
+    if not Image:
+        return "image/webp", file
     try:
         image = Image.open(BytesIO(file)).convert("RGBA")
         new_file = BytesIO()
         image.save(new_file, to)
-        return f"image/{to}", new_file.getvalue()
+        w, h = image.size
+        return f"image/{to}", new_file.getvalue(), w, h
     except Exception:
         log.exception(f"Failed to convert webp to {to}")
         return "image/webp", file
 
 
-async def transfer_file_to_matrix(db, client, intent, location):
+def _temp_file_name(ext):
+    return ("/tmp/mxtg-video-"
+            + "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10))
+            + ext)
+
+
+def _read_video_thumbnail(data, video_ext="mp4", frame_ext="png", max_size=(1024, 720)):
+    # We don't have any way to read the video from memory, so save it to disk.
+    temp_file = _temp_file_name(video_ext)
+    with open(temp_file, "wb") as file:
+        file.write(data)
+
+    # Read temp file and get frame
+    clip = VideoFileClip(temp_file)
+    frame = clip.get_frame(0)
+
+    # Convert to png and save to BytesIO
+    image = Image.fromarray(frame).convert("RGBA")
+    thumbnail_file = BytesIO()
+    if max_size:
+        image.thumbnail(max_size, Image.ANTIALIAS)
+    image.save(thumbnail_file, frame_ext)
+
+    os.remove(temp_file)
+
+    w, h = image.size
+    return thumbnail_file.getvalue(), w, h
+
+
+def _location_to_id(location):
     if isinstance(location, (Document, InputDocumentFileLocation)):
-        id = f"{location.id}-{location.version}"
+        return f"{location.id}-{location.version}"
     elif isinstance(location, (FileLocation, InputFileLocation)):
-        id = f"{location.volume_id}-{location.local_id}"
+        return f"{location.volume_id}-{location.local_id}"
     else:
         return None
 
+
+async def transfer_thumbnail_to_matrix(client, intent, thumbnail_loc, video, mime):
+    if not Image or not VideoFileClip:
+        return None
+
+    id = _location_to_id(thumbnail_loc)
+    if not id:
+        return None
+
+    video_ext = mimetypes.guess_extension(mime)
+    if VideoFileClip and video_ext:
+        file, width, height = _read_video_thumbnail(video, video_ext, frame_ext="png")
+        mime_type = "image/png"
+    else:
+        file = await client.download_file_bytes(thumbnail_loc)
+        width, height = None, None
+        mime_type = magic.from_buffer(file, mime=True)
+
+    uploaded = await intent.upload_file(file, mime_type)
+
+    return DBTelegramFile(id=id, mxc=uploaded["content_uri"], mime_type=mime_type,
+                          was_converted=False, timestamp=int(time.time()), size=len(file),
+                          width=width, height=height)
+
+
+async def transfer_file_to_matrix(db, client, intent, location, thumbnail=None):
+    id = _location_to_id(location)
+    if not id:
+        return None
+
     db_file = DBTelegramFile.query.get(id)
     if db_file:
         return db_file
@@ -58,18 +131,26 @@ async def transfer_file_to_matrix(db, client, intent, location):
         file = await client.download_file_bytes(location)
     except LocationInvalidError:
         return None
+    width, height = None, None
     mime_type = magic.from_buffer(file, mime=True)
 
     image_converted = False
     if mime_type == "image/webp":
-        mime_type, file = _convert_webp(file, to="png")
+        mime_type, file, width, height = _convert_webp(file, to="png")
         image_converted = True
 
     uploaded = await intent.upload_file(file, mime_type)
 
     db_file = DBTelegramFile(id=id, mxc=uploaded["content_uri"],
                              mime_type=mime_type, was_converted=image_converted,
-                             timestamp=int(time.time()))
+                             timestamp=int(time.time()), size=len(file),
+                             width=width, height=height)
+    if thumbnail:
+        if isinstance(thumbnail, (PhotoSize, PhotoCachedSize)):
+            thumbnail = thumbnail.location
+        db_file.thumbnail = await transfer_thumbnail_to_matrix(client, intent, thumbnail, file,
+                                                               mime_type)
+
     try:
         db.add(db_file)
         db.commit()