Add: celery task for generating video thumbnails

This commit is contained in:
Aroy-Art 2025-02-24 22:33:59 +01:00
parent e67f5f2dd7
commit 0e6146c6e8
Signed by: Aroy
GPG key ID: 583642324A1D2070

View file

@ -1,5 +1,7 @@
import os
import io
import hashlib
import subprocess
from django.db import transaction
@ -10,6 +12,8 @@ import blurhash
from .models import PostFileModel
from utils.hash import compute_file_hash_blake3, compute_blur_hash
@shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5)
def generate_blur_hash_PostFile(file_id):
@ -52,3 +56,78 @@ def generate_md5_hash_PostFile(file_id):
print(error_message)
raise Retry(exc=e) # Retry on exception
return error_message # This ensures the error message is stored in the results
@shared_task(name="generate_video_thumbnail")
def generate_video_thumbnail(
file_id: int, size: int = 0, timestamp=None, movie_strip: bool = False
):
"""
Generate video thumbnails using ffmpegthumbnailer and update the PostFileModel instance.
Args:
file_id (int): ID of the PostFileModel instance
size (int): Desired thumbnail width or height defulats to video size
timestamp (float): Timestamp(s) in seconds where the thumbnail should be extracted
movie_strip (bool): Create a movie strip overlay
Returns:
str: Success message or error message
"""
try:
with transaction.atomic():
# Retrieve the PostFileModel instance with a lock
pstfile = PostFileModel.objects.select_for_update().get(id=file_id)
if not pstfile.file:
return "Error: Video file not found for the given file_id."
video_path = pstfile.file.path
# Create output directory if it doesn't exist
output_dir = os.path.join("/tmp/thumbgen/", str(file_id))
os.makedirs(output_dir, exist_ok=True)
thumbnail_filename = f"thumbnail_{pstfile.hash_blake3}.png"
thumbnail_file_path = os.path.join(output_dir, thumbnail_filename)
cmd = [
"ffmpegthumbnailer",
"-i",
video_path,
"-o",
thumbnail_file_path,
"-s",
str(size),
"-m",
]
if movie_strip:
cmd.extend(["-f"])
# Generate thumbnail at specified timestamps
if timestamp is not None:
cmd.extend(["-t", f"{timestamp}"])
subprocess.run(cmd, check=True)
# Update the PostFileModel's thumbnail field with the new file
with open(thumbnail_file_path, "rb") as file:
pstfile.thumbnail.save(thumbnail_filename, file)
pstfile.thumbnail_hash_blake3 = compute_file_hash_blake3(
thumbnail_file_path
)
print(compute_blur_hash(thumbnail_file_path))
pstfile.thumbnail_blur_hash = compute_blur_hash(thumbnail_file_path)
pstfile.save()
return f"Thumbnail generated and saved to {thumbnail_file_path}"
except subprocess.CalledProcessError as e:
return f"Error generating thumbnail: {str(e)}"
except Exception as e:
return f"Unexpected error: {str(e)}"