feat: enhance video thumbnail generation

domain/virtual_fs/thumbnail.py

@@ -13,8 +13,12 @@ ALLOWED_EXT = {"jpg", "jpeg", "png", "webp", "gif", "bmp",
 RAW_EXT = {"arw", "cr2", "cr3", "nef", "rw2", "orf", "pef", "dng"}
 VIDEO_EXT = {"mp4", "mov", "m4v", "avi", "mkv", "wmv", "flv", "webm", "mpg", "mpeg", "3gp"}
 MAX_IMAGE_SOURCE_SIZE = 200 * 1024 * 1024
-VIDEO_RANGE_LIMIT = 16 * 1024 * 1024  # 16MB
-VIDEO_INITIAL_CHUNK = 4 * 1024 * 1024
+VIDEO_TAIL_LIMIT = 2 * 1024 * 1024  # 2MB
+VIDEO_TAIL_FALLBACK_LIMIT = 4 * 1024 * 1024  # 4MB
+VIDEO_HEAD_LIMIT = 2 * 1024 * 1024  # 2MB
+VIDEO_HEAD_FALLBACK_LIMIT = 4 * 1024 * 1024  # 4MB
+VIDEO_THUMB_SEEK_SECONDS = (20, 10, 5, 3, 1, 0)
+VIDEO_BLACK_FRAME_MEAN_THRESHOLD = 12.0
 CACHE_ROOT = Path('data/.thumb_cache')
 
 
@@ -176,42 +180,54 @@ async def _read_range_slice(adapter, root: str, rel: str, start: int, end: int)
     return b""
 
 
-async def _read_video_prefix(adapter, root: str, rel: str, size: int, limit: int = VIDEO_RANGE_LIMIT) -> bytes:
-    chunk_size = min(VIDEO_INITIAL_CHUNK, limit)
-    offset = 0
-    collected = bytearray()
-
-    while len(collected) < limit:
-        end = offset + chunk_size - 1
-        data = await _read_range_slice(adapter, root, rel, offset, end)
-        if not data:
-            break
-        collected.extend(data)
-        if len(data) < chunk_size:
-            break
-        offset += len(data)
-        remaining = limit - len(collected)
-        if remaining <= 0:
-            break
-        chunk_size = min(chunk_size * 2, remaining)
-
-    if not collected and size <= limit:
-        read_file = getattr(adapter, "read_file", None)
-        if callable(read_file):
-            blob = await read_file(root, rel)
-            if blob:
-                return bytes(blob[:limit])
-
-    return bytes(collected[:limit])
+async def _read_video_head(adapter, root: str, rel: str, size: int, limit: int = VIDEO_HEAD_LIMIT) -> bytes:
+    end = limit - 1
+    if size > 0:
+        end = min(end, size - 1)
+    if end < 0:
+        return b""
+    return await _read_range_slice(adapter, root, rel, 0, end)
 
 
-async def _run_ffmpeg_extract_frame(src_path: str, dst_path: str):
+async def _read_video_tail(adapter, root: str, rel: str, size: int, limit: int) -> Tuple[bytes, int]:
+    if size <= 0:
+        return b"", 0
+    start = max(0, size - limit)
+    end = size - 1
+    data = await _read_range_slice(adapter, root, rel, start, end)
+    return data, start
+
+
+def _write_video_sparse_file(rel: str, head_bytes: bytes, tail_bytes: bytes, tail_offset: int) -> str:
+    suffix = Path(rel).suffix or ".mp4"
+    src_tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
+    src_path = src_tmp.name
+    try:
+        if head_bytes:
+            src_tmp.write(head_bytes)
+            src_tmp.flush()
+    finally:
+        src_tmp.close()
+
+    if tail_bytes:
+        with open(src_path, "r+b") as f:
+            f.seek(max(0, int(tail_offset)))
+            f.write(tail_bytes)
+            f.flush()
+    return src_path
+
+
+async def _run_ffmpeg_extract_frame(src_path: str, dst_path: str, *, seek_seconds: float | None = None):
     cmd = [
         "ffmpeg",
         "-y",
         "-hide_banner",
         "-loglevel", "error",
         "-i", src_path,
+    ]
+    if seek_seconds is not None:
+        cmd += ["-ss", str(seek_seconds)]
+    cmd += [
         "-frames:v", "1",
         dst_path,
     ]
@@ -230,32 +246,72 @@ async def _run_ffmpeg_extract_frame(src_path: str, dst_path: str):
         raise RuntimeError(message)
 
 
-async def _generate_video_thumb(video_bytes: bytes, rel: str, w: int, h: int, fit: str) -> Tuple[bytes, str]:
-    from PIL import Image
+def _frame_mean_luma(im) -> float:
+    from PIL import ImageStat
+    gray = im.convert("L").resize((64, 64))
+    return float(ImageStat.Stat(gray).mean[0])
 
-    suffix = Path(rel).suffix or ".mp4"
-    src_tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
-    src_path = src_tmp.name
-    try:
-        src_tmp.write(video_bytes)
-        src_tmp.flush()
-    finally:
-        src_tmp.close()
+
+def _is_black_image_bytes(image_bytes: bytes) -> bool:
+    from PIL import Image
+    with Image.open(io.BytesIO(image_bytes)) as im:
+        im.load()
+        return _frame_mean_luma(im) < VIDEO_BLACK_FRAME_MEAN_THRESHOLD
+
+
+async def _generate_video_thumb_from_src_path(src_path: str, w: int, h: int, fit: str) -> Tuple[bytes, str]:
+    from PIL import Image
 
     dst_tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
     dst_path = dst_tmp.name
     dst_tmp.close()
 
+    best: tuple[float, bytes, str] | None = None
+    last_error: Exception | None = None
     try:
-        await _run_ffmpeg_extract_frame(src_path, dst_path)
-        with Image.open(dst_path) as im:
-            im.load()
-            return _image_to_webp(im, w, h, fit)
+        for seek_seconds in VIDEO_THUMB_SEEK_SECONDS:
+            try:
+                with suppress(FileNotFoundError):
+                    Path(dst_path).unlink()
+                await _run_ffmpeg_extract_frame(src_path, dst_path, seek_seconds=seek_seconds)
+                with Image.open(dst_path) as im:
+                    im.load()
+                    mean = _frame_mean_luma(im)
+                    webp_bytes, mime = _image_to_webp(im, w, h, fit)
+
+                if best is None or mean > best[0]:
+                    best = (mean, webp_bytes, mime)
+                if mean >= VIDEO_BLACK_FRAME_MEAN_THRESHOLD:
+                    return webp_bytes, mime
+            except Exception as e:
+                last_error = e
+                continue
+
+        if best is not None:
+            return best[1], best[2]
+        if last_error is not None:
+            raise last_error
+        raise RuntimeError("ffmpeg 截帧失败")
+    finally:
+        with suppress(FileNotFoundError):
+            Path(dst_path).unlink()
+
+
+async def _generate_video_thumb_from_segments(
+    head_bytes: bytes,
+    tail_bytes: bytes,
+    tail_offset: int,
+    rel: str,
+    w: int,
+    h: int,
+    fit: str,
+) -> Tuple[bytes, str]:
+    src_path = _write_video_sparse_file(rel, head_bytes, tail_bytes, tail_offset)
+    try:
+        return await _generate_video_thumb_from_src_path(src_path, w, h, fit)
     finally:
         with suppress(FileNotFoundError):
             Path(src_path).unlink()
-        with suppress(FileNotFoundError):
-            Path(dst_path).unlink()
 
 
 async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: int, h: int, fit: str = 'cover'):
@@ -294,23 +350,55 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w:
 
     if not thumb_bytes:
         if is_video:
-            try:
-                video_bytes = await _read_video_prefix(adapter, root, rel, size)
-            except HTTPException:
-                raise
-            except Exception as e:
-                print(f"Video prefix read failed: {e}")
-                raise HTTPException(500, detail=f"Video read failed: {e}")
+            async def _read_head(limit: int) -> bytes:
+                try:
+                    return await _read_video_head(adapter, root, rel, size, limit=limit)
+                except HTTPException:
+                    raise
+                except Exception as e:
+                    print(f"Video head read failed: {e}")
+                    raise HTTPException(500, detail=f"Video read failed: {e}")
 
-            if not video_bytes:
+            async def _read_tail(limit: int) -> Tuple[bytes, int]:
+                try:
+                    return await _read_video_tail(adapter, root, rel, size, limit=limit)
+                except HTTPException:
+                    raise
+                except Exception as e:
+                    print(f"Video tail read failed: {e}")
+                    raise HTTPException(500, detail=f"Video read failed: {e}")
+
+            head_bytes = await _read_head(VIDEO_HEAD_LIMIT)
+            tail_bytes, tail_offset = await _read_tail(VIDEO_TAIL_LIMIT)
+            if not head_bytes and not tail_bytes:
                 raise HTTPException(500, detail="Unable to read video data for thumbnail")
 
             try:
-                thumb_bytes, mime = await _generate_video_thumb(video_bytes, rel, w, h, fit)
-            except Exception as e:
-                print(f"Video thumbnail generation failed: {e}")
-                raise HTTPException(
-                    500, detail=f"Video thumbnail generation failed: {e}")
+                thumb_bytes, mime = await _generate_video_thumb_from_segments(
+                    head_bytes, tail_bytes, tail_offset, rel, w, h, fit
+                )
+            except Exception:
+                try:
+                    tail_bytes, tail_offset = await _read_tail(VIDEO_TAIL_FALLBACK_LIMIT)
+                    thumb_bytes, mime = await _generate_video_thumb_from_segments(
+                        head_bytes, tail_bytes, tail_offset, rel, w, h, fit
+                    )
+                except HTTPException:
+                    raise
+                except Exception as e2:
+                    print(f"Video thumbnail generation failed: {e2}")
+                    raise HTTPException(500, detail=f"Video thumbnail generation failed: {e2}")
+
+            if thumb_bytes and _is_black_image_bytes(thumb_bytes):
+                try:
+                    head_bytes = await _read_head(VIDEO_HEAD_FALLBACK_LIMIT)
+                    retry_thumb, retry_mime = await _generate_video_thumb_from_segments(
+                        head_bytes, tail_bytes, tail_offset, rel, w, h, fit
+                    )
+                    if retry_thumb and not _is_black_image_bytes(retry_thumb):
+                        thumb_bytes, mime = retry_thumb, retry_mime
+                except Exception:
+                    pass
         else:
             read_data = await adapter.read_file(root, rel)
             try: