134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
import os
|
|
import io
|
|
import hashlib
|
|
import subprocess
|
|
|
|
from django.db import transaction
|
|
|
|
from celery import shared_task
|
|
from celery.exceptions import Retry
|
|
from PIL import Image as PillowImage
|
|
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):
|
|
try:
|
|
with transaction.atomic():
|
|
img = PostFileModel.objects.select_for_update().get(id=file_id)
|
|
image_data = io.BytesIO(img.file.read())
|
|
pil_img = PillowImage.open(image_data)
|
|
|
|
blurhash_string = blurhash.encode(pil_img, 4, 3)
|
|
|
|
img.refresh_from_db()
|
|
img.blur_hash = blurhash_string
|
|
img.save()
|
|
return f"Successfully generated blur hash for file {file_id}" # Success message
|
|
except Exception as e:
|
|
error_message = f"Error generating blur hash for file {file_id}: {e}"
|
|
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(autoretry_for=(Exception,), retry_backoff=True, max_retries=5)
|
|
def generate_md5_hash_PostFile(file_id):
|
|
try:
|
|
with transaction.atomic():
|
|
pstfile = PostFileModel.objects.select_for_update().get(id=file_id)
|
|
hash_md5 = hashlib.md5()
|
|
with open(pstfile.file.path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
hash_md5.update(chunk)
|
|
md5_hash = hash_md5.hexdigest()
|
|
|
|
pstfile.refresh_from_db()
|
|
pstfile.hash_md5 = md5_hash
|
|
pstfile.save()
|
|
return f"Successfully generated MD5 hash for file {file_id}" # Success message
|
|
except Exception as e:
|
|
error_message = f"Error generating MD5 hash for file {file_id}: {e}"
|
|
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 = "/tmp/thumbgen/"
|
|
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)
|
|
|
|
thumbnail_hash_blake3 = compute_file_hash_blake3(thumbnail_file_path)
|
|
|
|
# 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 = thumbnail_hash_blake3
|
|
|
|
pstfile.thumbnail_blur_hash = compute_blur_hash(thumbnail_file_path)
|
|
|
|
pstfile.save()
|
|
|
|
os.remove(thumbnail_file_path)
|
|
|
|
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)}"
|