165 lines
4.6 KiB
Python
165 lines
4.6 KiB
Python
import io
|
|
import logging
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from django.core.files.base import ContentFile
|
|
from PIL import Image, ImageOps
|
|
|
|
from tasks.models import TaskAttachment
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TASK_COVER_THUMBNAIL_SIZE = (640, 280)
|
|
IMAGE_EXTENSIONS = {
|
|
".avif",
|
|
".bmp",
|
|
".gif",
|
|
".heic",
|
|
".heif",
|
|
".jpeg",
|
|
".jpg",
|
|
".png",
|
|
".webp",
|
|
}
|
|
VIDEO_EXTENSIONS = {
|
|
".m4v",
|
|
".mkv",
|
|
".mov",
|
|
".mp4",
|
|
".mpeg",
|
|
".mpg",
|
|
".ogv",
|
|
".webm",
|
|
}
|
|
|
|
|
|
def task_attachment_media_kind(attachment: TaskAttachment) -> str | None:
|
|
mime_type = (attachment.mime_type or "").lower()
|
|
suffix = Path(attachment.file_name or "").suffix.lower()
|
|
if mime_type.startswith("image/") or suffix in IMAGE_EXTENSIONS:
|
|
return "image"
|
|
if mime_type.startswith("video/") or suffix in VIDEO_EXTENSIONS:
|
|
return "video"
|
|
return None
|
|
|
|
|
|
def _render_image_thumbnail(attachment: TaskAttachment) -> bytes:
|
|
with attachment.file.open("rb") as source:
|
|
image = Image.open(source)
|
|
image.seek(0)
|
|
image = ImageOps.exif_transpose(image)
|
|
if image.mode not in ("RGB", "RGBA"):
|
|
image = image.convert("RGBA" if "transparency" in image.info else "RGB")
|
|
if image.mode == "RGBA":
|
|
background = Image.new("RGB", image.size, (20, 27, 36))
|
|
background.paste(image, mask=image.getchannel("A"))
|
|
image = background
|
|
image = ImageOps.fit(
|
|
image,
|
|
TASK_COVER_THUMBNAIL_SIZE,
|
|
method=Image.Resampling.LANCZOS,
|
|
)
|
|
output = io.BytesIO()
|
|
image.save(output, format="WEBP", quality=84, method=5)
|
|
return output.getvalue()
|
|
|
|
|
|
def _render_video_thumbnail(attachment: TaskAttachment) -> bytes:
|
|
ffmpeg = shutil.which("ffmpeg")
|
|
if not ffmpeg:
|
|
raise RuntimeError("ffmpeg is not installed")
|
|
|
|
input_suffix = Path(attachment.file_name or "").suffix or ".video"
|
|
input_path = ""
|
|
temporary_input = False
|
|
output_path = ""
|
|
try:
|
|
try:
|
|
input_path = attachment.file.path
|
|
except (AttributeError, NotImplementedError):
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=input_suffix, delete=False
|
|
) as source_file:
|
|
input_path = source_file.name
|
|
temporary_input = True
|
|
with attachment.file.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
source_file.write(chunk)
|
|
with tempfile.NamedTemporaryFile(suffix=".webp", delete=False) as output_file:
|
|
output_path = output_file.name
|
|
|
|
filter_value = "scale=640:280:force_original_aspect_ratio=increase,crop=640:280"
|
|
command = [
|
|
ffmpeg,
|
|
"-hide_banner",
|
|
"-loglevel",
|
|
"error",
|
|
"-ss",
|
|
"0.5",
|
|
"-i",
|
|
input_path,
|
|
"-frames:v",
|
|
"1",
|
|
"-vf",
|
|
filter_value,
|
|
"-c:v",
|
|
"libwebp",
|
|
"-quality",
|
|
"84",
|
|
"-y",
|
|
output_path,
|
|
]
|
|
result = subprocess.run(
|
|
command,
|
|
capture_output=True,
|
|
check=False,
|
|
timeout=30,
|
|
)
|
|
if result.returncode != 0:
|
|
command[5] = "0"
|
|
result = subprocess.run(
|
|
command,
|
|
capture_output=True,
|
|
check=False,
|
|
timeout=30,
|
|
)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(result.stderr.decode("utf-8", errors="replace"))
|
|
return Path(output_path).read_bytes()
|
|
finally:
|
|
if temporary_input and input_path:
|
|
Path(input_path).unlink(missing_ok=True)
|
|
if output_path:
|
|
Path(output_path).unlink(missing_ok=True)
|
|
|
|
|
|
def ensure_task_attachment_thumbnail(attachment: TaskAttachment) -> bool:
|
|
if attachment.thumbnail:
|
|
return True
|
|
media_kind = task_attachment_media_kind(attachment)
|
|
if not media_kind:
|
|
return False
|
|
|
|
try:
|
|
thumbnail = (
|
|
_render_image_thumbnail(attachment)
|
|
if media_kind == "image"
|
|
else _render_video_thumbnail(attachment)
|
|
)
|
|
attachment.thumbnail.save(
|
|
f"task-{attachment.task_id}-attachment-{attachment.id}.webp",
|
|
ContentFile(thumbnail),
|
|
save=True,
|
|
)
|
|
return True
|
|
except Exception:
|
|
logger.exception(
|
|
"Could not generate thumbnail for task attachment %s",
|
|
attachment.id,
|
|
)
|
|
return False
|