From c2015dd17ceaba4b254cd4fc3ba2f5ebf2063b6a Mon Sep 17 00:00:00 2001 From: shiyu Date: Sun, 3 May 2026 23:48:09 +0800 Subject: [PATCH] feat: enhance thumbnail handling and add native thumbnail support in VirtualFS --- domain/adapters/providers/telegram.py | 50 +++++++++++++++++---------- domain/virtual_fs/listing.py | 8 ++++- domain/virtual_fs/routes.py | 13 +++++-- domain/virtual_fs/thumbnail.py | 24 +++++++------ 4 files changed, 64 insertions(+), 31 deletions(-) diff --git a/domain/adapters/providers/telegram.py b/domain/adapters/providers/telegram.py index 28d243d..dc7f688 100644 --- a/domain/adapters/providers/telegram.py +++ b/domain/adapters/providers/telegram.py @@ -5,7 +5,7 @@ import io import os import struct from models import StorageAdapter -from telethon import TelegramClient +from telethon import TelegramClient, utils from telethon.crypto import AuthKey from telethon.sessions import StringSession from telethon.tl import types @@ -132,29 +132,42 @@ class TelegramAdapter: return None cached = [] - others = [] + downloadable = [] for t in thumbs: if isinstance(t, (types.PhotoCachedSize, types.PhotoStrippedSize)): cached.append(t) elif isinstance(t, (types.PhotoSize, types.PhotoSizeProgressive)): if not isinstance(t, types.PhotoSizeEmpty): - others.append(t) + downloadable.append(t) - if cached: - cached.sort(key=lambda x: len(getattr(x, "bytes", b"") or b"")) - return cached[-1] - - if others: + if downloadable: def _sz(x): if isinstance(x, types.PhotoSizeProgressive): return max(x.sizes or [0]) return int(getattr(x, "size", 0) or 0) - others.sort(key=_sz) - return others[-1] + downloadable.sort(key=_sz) + return downloadable[-1] + + if cached: + cached.sort(key=lambda x: len(getattr(x, "bytes", b"") or b"")) + return cached[-1] return None + @staticmethod + def _get_message_thumbs(message) -> list: + doc = message.document or message.video + if doc and getattr(doc, "thumbs", None): + return list(doc.thumbs or []) + if message.photo and getattr(message.photo, "sizes", None): + return list(message.photo.sizes or []) + return [] + + @classmethod + def _message_has_thumbnail(cls, message) -> bool: + return cls._pick_photo_thumb(cls._get_message_thumbs(message)) is not None + def _build_session(self) -> StringSession: s = (self.session_string or "").strip() if not s: @@ -229,6 +242,7 @@ class TelegramAdapter: "size": size, "mtime": int(message.date.timestamp()), "type": "file", + "has_thumbnail": self._message_has_thumbnail(message), }) finally: if client.is_connected(): @@ -386,17 +400,16 @@ class TelegramAdapter: if not message: return None - doc = message.document or message.video - thumbs = None - if doc and getattr(doc, "thumbs", None): - thumbs = list(doc.thumbs or []) - elif message.photo and getattr(message.photo, "sizes", None): - thumbs = list(message.photo.sizes or []) - - thumb = self._pick_photo_thumb(thumbs) + thumb = self._pick_photo_thumb(self._get_message_thumbs(message)) if not thumb: return None + embedded = getattr(thumb, "bytes", None) + if embedded and isinstance(thumb, types.PhotoCachedSize): + return bytes(embedded) + if embedded and isinstance(thumb, types.PhotoStrippedSize): + return utils.stripped_photo_to_jpg(bytes(embedded)) + result = await client.download_media(message, bytes, thumb=thumb) if isinstance(result, (bytes, bytearray)): return bytes(result) @@ -569,6 +582,7 @@ class TelegramAdapter: "size": size, "mtime": int(message.date.timestamp()), "type": "file", + "has_thumbnail": self._message_has_thumbnail(message), } finally: if client.is_connected(): diff --git a/domain/virtual_fs/listing.py b/domain/virtual_fs/listing.py index 7fd478f..bdf4877 100644 --- a/domain/virtual_fs/listing.py +++ b/domain/virtual_fs/listing.py @@ -89,6 +89,9 @@ class VirtualFSListingMixin(VirtualFSResolverMixin): def annotate_entry(entry: Dict) -> None: if not entry.get("is_dir"): + if entry.get("has_thumbnail") is not None: + entry["has_thumbnail"] = bool(entry.get("has_thumbnail")) + return name = entry.get("name", "") entry["has_thumbnail"] = bool(is_image_filename(name) or is_video_filename(name)) else: @@ -273,7 +276,10 @@ class VirtualFSListingMixin(VirtualFSResolverMixin): is_dir = False rel_name = rel.rstrip("/").split("/")[-1] if rel else path.rstrip("/").split("/")[-1] name_hint = str(info.get("name") or rel_name or "") - info["has_thumbnail"] = bool(not is_dir and (is_image_filename(name_hint) or is_video_filename(name_hint))) + if not is_dir and info.get("has_thumbnail") is not None: + info["has_thumbnail"] = bool(info.get("has_thumbnail")) + else: + info["has_thumbnail"] = bool(not is_dir and (is_image_filename(name_hint) or is_video_filename(name_hint))) if verbose and not is_dir: vector_index = await cls._gather_vector_index(path) if vector_index is not None: diff --git a/domain/virtual_fs/routes.py b/domain/virtual_fs/routes.py index 475e2ff..172d554 100644 --- a/domain/virtual_fs/routes.py +++ b/domain/virtual_fs/routes.py @@ -89,8 +89,17 @@ class VirtualFSRouteMixin(VirtualFSTempLinkMixin): adapter, mount, root, rel = await cls.resolve_adapter_and_rel(full_path) if not rel or rel.endswith("/"): raise HTTPException(400, detail="Not a file") - if not (is_image_filename(rel) or is_video_filename(rel)): - raise HTTPException(404, detail="Not an image or video") + has_native_thumb = False + if callable(getattr(adapter, "get_thumbnail", None)): + stat_file = getattr(adapter, "stat_file", None) + if callable(stat_file): + try: + stat = await stat_file(root, rel) + has_native_thumb = bool(isinstance(stat, dict) and stat.get("has_thumbnail")) + except Exception: + has_native_thumb = False + if not (is_image_filename(rel) or is_video_filename(rel) or has_native_thumb): + raise HTTPException(404, detail="Not an image, video, or native thumbnail file") data, mime, key = await get_or_create_thumb(adapter, mount.id, root, rel, w, h, fit) # type: ignore headers = { "Cache-Control": "public, max-age=3600", diff --git a/domain/virtual_fs/thumbnail.py b/domain/virtual_fs/thumbnail.py index 5b332e8..6538e5b 100644 --- a/domain/virtual_fs/thumbnail.py +++ b/domain/virtual_fs/thumbnail.py @@ -23,6 +23,7 @@ VIDEO_HEAD_FALLBACK_LIMIT = 4 * 1024 * 1024 # 4MB VIDEO_THUMB_SEEK_SECONDS = (15, 10, 5, 3, 1, 0) VIDEO_BLACK_FRAME_MEAN_THRESHOLD = 12.0 CACHE_ROOT = Path('data/.thumb_cache') +THUMB_CACHE_VERSION = "v2" def is_image_filename(name: str) -> bool: @@ -47,7 +48,7 @@ def is_video_filename(name: str) -> bool: def _cache_key(adapter_id: int, rel: str, size: int, mtime: int, w: int, h: int, fit: str) -> str: - raw = f"{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode() + raw = f"{THUMB_CACHE_VERSION}|{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode() return hashlib.sha1(raw).hexdigest() @@ -385,8 +386,11 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: stat = await adapter.stat_file(root, rel) size = int(stat.get('size') or 0) is_video = is_video_filename(rel) - if not is_video and size > MAX_IMAGE_SOURCE_SIZE: - raise HTTPException(400, detail="Image too large for thumbnail") + is_image = is_image_filename(rel) + get_thumb_impl = getattr(adapter, "get_thumbnail", None) + should_try_native_thumb = callable(get_thumb_impl) and ( + is_image or is_video or bool(stat.get("has_thumbnail")) + ) key = _cache_key(adapter_id, rel, size, int( stat.get('mtime', 0)), w, h, fit) @@ -397,8 +401,7 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: _ensure_cache_dir(path) thumb_bytes, mime = None, None - get_thumb_impl = getattr(adapter, "get_thumbnail", None) - if callable(get_thumb_impl): + if should_try_native_thumb: size_str = "large" if w > 400 else "medium" if w > 100 else "small" native_thumb_bytes = await get_thumb_impl(root, rel, size_str) @@ -406,10 +409,7 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: try: from PIL import Image im = Image.open(io.BytesIO(native_thumb_bytes)) - buf = io.BytesIO() - im.save(buf, 'WEBP', quality=85) - thumb_bytes = buf.getvalue() - mime = 'image/webp' + thumb_bytes, mime = _image_to_webp(im, w, h, fit) except Exception as e: print( f"Failed to convert native thumbnail to WebP: {e}, falling back.") @@ -493,7 +493,9 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: thumb_bytes, mime = retry_thumb, retry_mime except Exception: pass - else: + elif is_image: + if size > MAX_IMAGE_SOURCE_SIZE: + raise HTTPException(400, detail="Image too large for thumbnail") read_data = await adapter.read_file(root, rel) try: thumb_bytes, mime = generate_thumb( @@ -502,6 +504,8 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: print(e) raise HTTPException( 500, detail=f"Thumbnail generation failed: {e}") + else: + raise HTTPException(500, detail="Native thumbnail unavailable") if thumb_bytes: path.write_bytes(thumb_bytes)