126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
"""
|
||
Thumbnail generation for uploaded media files.
|
||
|
||
Supports:
|
||
- Images (via Pillow): JPEG thumbnail at 300×300 max
|
||
- Videos (via ffmpeg subprocess): extract first frame as JPEG
|
||
- PDFs (via pdf2image + Poppler): render first page as JPEG
|
||
|
||
Returns None if the type is unsupported or if generation fails.
|
||
"""
|
||
import io
|
||
import logging
|
||
import subprocess
|
||
from pathlib import Path
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
THUMB_SIZE = (220, 220) # small enough for gallery tiles; keeps files ~4-6 KB
|
||
|
||
|
||
def _thumb_from_image(content: bytes) -> bytes | None:
|
||
try:
|
||
from PIL import Image, ImageOps
|
||
img = Image.open(io.BytesIO(content))
|
||
img = ImageOps.exif_transpose(img) # honour EXIF Orientation tag before resizing
|
||
img = img.convert("RGB")
|
||
img.thumbnail(THUMB_SIZE, Image.LANCZOS)
|
||
out = io.BytesIO()
|
||
# quality=55 + optimize=True + progressive encoding → ~4-6 KB for typical photos
|
||
img.save(out, format="JPEG", quality=65, optimize=True, progressive=True)
|
||
return out.getvalue()
|
||
except Exception as e:
|
||
logger.warning("Image thumbnail failed: %s", e)
|
||
return None
|
||
|
||
|
||
def _thumb_from_video(content: bytes) -> bytes | None:
|
||
"""
|
||
Extract the first frame of a video as a JPEG thumbnail.
|
||
|
||
We write the video to a temp file instead of piping it to ffmpeg because
|
||
most video containers (MP4, MOV, MKV …) store their index (moov atom) at
|
||
an arbitrary offset and ffmpeg cannot seek on a pipe — causing rc≠0 with
|
||
"moov atom not found" or similar errors when stdin is used.
|
||
"""
|
||
import tempfile
|
||
import os
|
||
try:
|
||
# Write to a temp file so ffmpeg can seek freely
|
||
with tempfile.NamedTemporaryFile(suffix=".video", delete=False) as tmp_in:
|
||
tmp_in.write(content)
|
||
tmp_in_path = tmp_in.name
|
||
|
||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_out:
|
||
tmp_out_path = tmp_out.name
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"ffmpeg", "-y",
|
||
"-i", tmp_in_path,
|
||
"-vframes", "1",
|
||
"-vf", f"scale={THUMB_SIZE[0]}:-2",
|
||
"-q:v", "4", # JPEG quality 1-31 (lower = better); 4 ≈ ~80% quality
|
||
tmp_out_path,
|
||
],
|
||
capture_output=True,
|
||
timeout=60,
|
||
)
|
||
if result.returncode == 0 and os.path.getsize(tmp_out_path) > 0:
|
||
with open(tmp_out_path, "rb") as f:
|
||
return f.read()
|
||
logger.warning(
|
||
"ffmpeg video thumb failed (rc=%s): %s",
|
||
result.returncode,
|
||
result.stderr[-400:].decode(errors="replace") if result.stderr else "",
|
||
)
|
||
return None
|
||
finally:
|
||
os.unlink(tmp_in_path)
|
||
try:
|
||
os.unlink(tmp_out_path)
|
||
except FileNotFoundError:
|
||
pass
|
||
except FileNotFoundError:
|
||
logger.warning("ffmpeg not found — video thumbnails unavailable")
|
||
return None
|
||
except Exception as e:
|
||
logger.warning("Video thumbnail failed: %s", e)
|
||
return None
|
||
|
||
|
||
def _thumb_from_pdf(content: bytes) -> bytes | None:
|
||
try:
|
||
from pdf2image import convert_from_bytes
|
||
pages = convert_from_bytes(content, first_page=1, last_page=1, size=THUMB_SIZE)
|
||
if not pages:
|
||
return None
|
||
out = io.BytesIO()
|
||
pages[0].save(out, format="JPEG", quality=55, optimize=True, progressive=True)
|
||
return out.getvalue()
|
||
except ImportError:
|
||
logger.warning("pdf2image not installed — PDF thumbnails unavailable")
|
||
return None
|
||
except Exception as e:
|
||
logger.warning("PDF thumbnail failed: %s", e)
|
||
return None
|
||
|
||
|
||
def generate_thumbnail(content: bytes, mime_type: str, filename: str) -> bytes | None:
|
||
"""
|
||
Generate a small JPEG thumbnail for the given file content.
|
||
Returns JPEG bytes or None if unsupported / generation fails.
|
||
"""
|
||
mt = (mime_type or "").lower()
|
||
fn = (filename or "").lower()
|
||
|
||
if mt.startswith("image/"):
|
||
return _thumb_from_image(content)
|
||
if mt.startswith("video/"):
|
||
return _thumb_from_video(content)
|
||
if mt == "application/pdf" or fn.endswith(".pdf"):
|
||
return _thumb_from_pdf(content)
|
||
|
||
return None
|