277 lines
9.8 KiB
Python
277 lines
9.8 KiB
Python
"""
|
|
Celery tasks for thumbnail generation
|
|
"""
|
|
import os
|
|
import asyncio
|
|
from pathlib import Path
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
import json
|
|
|
|
from celery import shared_task
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
import pyvips
|
|
import rawpy
|
|
import imageio
|
|
from PIL import Image
|
|
from pillow_heif import register_heif_opener
|
|
import ffmpeg
|
|
|
|
from app.database import AsyncSessionLocal
|
|
from app.models import Photo
|
|
from app.config import settings
|
|
|
|
# Register HEIF opener with Pillow
|
|
register_heif_opener()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Thumbnail sizes configuration
|
|
THUMB_SIZES = {
|
|
'small': settings.thumbnails.small,
|
|
'medium': settings.thumbnails.medium,
|
|
'large': settings.thumbnails.large
|
|
}
|
|
|
|
def get_thumb_path(photo_id: str, size: str) -> str:
|
|
"""Get the path for a thumbnail file"""
|
|
thumb_dir = f"/data/thumbs/{photo_id}"
|
|
os.makedirs(thumb_dir, exist_ok=True)
|
|
return f"{thumb_dir}/{size}.{settings.thumbnails.format}"
|
|
|
|
def process_standard_image(filepath: str) -> pyvips.Image:
|
|
"""Process standard image formats (JPEG, PNG, etc.)"""
|
|
return pyvips.Image.new_from_file(filepath, access='sequential')
|
|
|
|
def process_raw_image(filepath: str) -> pyvips.Image:
|
|
"""Process RAW image formats"""
|
|
try:
|
|
with rawpy.imread(filepath) as raw:
|
|
# Use half_size for faster processing
|
|
rgb = raw.postprocess(use_camera_wb=True, half_size=True)
|
|
# Convert numpy array to pyvips image
|
|
return pyvips.Image.new_from_array(rgb)
|
|
except Exception as e:
|
|
logger.error(f"Error processing RAW file {filepath}: {e}")
|
|
# Try to extract embedded JPEG preview
|
|
return extract_raw_preview(filepath)
|
|
|
|
def extract_raw_preview(filepath: str) -> Optional[pyvips.Image]:
|
|
"""Extract embedded JPEG preview from RAW file"""
|
|
try:
|
|
# Use exiftool to extract preview
|
|
import subprocess
|
|
import tempfile
|
|
|
|
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
|
|
cmd = ['exiftool', '-b', '-PreviewImage', filepath]
|
|
result = subprocess.run(cmd, capture_output=True)
|
|
|
|
if result.returncode == 0 and result.stdout:
|
|
tmp.write(result.stdout)
|
|
tmp.flush()
|
|
return pyvips.Image.new_from_file(tmp.name, access='sequential')
|
|
except Exception as e:
|
|
logger.error(f"Error extracting RAW preview from {filepath}: {e}")
|
|
|
|
return None
|
|
|
|
def process_heic_image(filepath: str) -> pyvips.Image:
|
|
"""Process HEIC/HEIF image formats"""
|
|
try:
|
|
# Use pillow-heif to open the image
|
|
img = Image.open(filepath)
|
|
# Convert to RGB if needed
|
|
if img.mode != 'RGB':
|
|
img = img.convert('RGB')
|
|
# Save to temp file and load with pyvips
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
|
|
img.save(tmp.name, 'JPEG')
|
|
return pyvips.Image.new_from_file(tmp.name, access='sequential')
|
|
except Exception as e:
|
|
logger.error(f"Error processing HEIC file {filepath}: {e}")
|
|
raise
|
|
|
|
def process_video_thumbnail(filepath: str) -> pyvips.Image:
|
|
"""Extract thumbnail from video file"""
|
|
try:
|
|
# Get video duration
|
|
probe = ffmpeg.probe(filepath)
|
|
duration = float(probe['streams'][0]['duration'])
|
|
|
|
# Extract frame at 10% of duration
|
|
timestamp = duration * 0.1
|
|
|
|
# Extract frame using ffmpeg
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
|
|
stream = ffmpeg.input(filepath, ss=timestamp)
|
|
stream = ffmpeg.output(stream, tmp.name, vframes=1, format='image2', vcodec='mjpeg')
|
|
ffmpeg.run(stream, capture_stdout=True, capture_stderr=True)
|
|
|
|
return pyvips.Image.new_from_file(tmp.name, access='sequential')
|
|
except Exception as e:
|
|
logger.error(f"Error extracting video thumbnail from {filepath}: {e}")
|
|
# Create a placeholder thumbnail
|
|
return create_placeholder_thumbnail('video')
|
|
|
|
def create_placeholder_thumbnail(media_type: str) -> pyvips.Image:
|
|
"""Create a placeholder thumbnail for failed processing"""
|
|
# Create a simple gray placeholder
|
|
placeholder = pyvips.Image.black(640, 480)
|
|
placeholder = placeholder + [128, 128, 128] # Make it gray
|
|
return placeholder
|
|
|
|
def auto_rotate_image(image: pyvips.Image) -> pyvips.Image:
|
|
"""Auto-rotate image based on EXIF orientation"""
|
|
try:
|
|
orientation = image.get('orientation')
|
|
|
|
rotation_map = {
|
|
3: 180,
|
|
6: 90,
|
|
8: 270
|
|
}
|
|
|
|
if orientation in rotation_map:
|
|
image = image.rot(rotation_map[orientation])
|
|
except:
|
|
pass # No orientation data available
|
|
|
|
return image
|
|
|
|
def generate_thumbnail(image: pyvips.Image, size: int, output_path: str):
|
|
"""Generate a thumbnail of the specified size"""
|
|
# Calculate scale to fit within size (longest edge)
|
|
width = image.width
|
|
height = image.height
|
|
|
|
if width > height:
|
|
scale = size / width
|
|
else:
|
|
scale = size / height
|
|
|
|
# Only downscale, never upscale
|
|
if scale < 1:
|
|
image = image.resize(scale)
|
|
|
|
# Save as WebP with specified quality
|
|
image.webpsave(
|
|
output_path,
|
|
Q=settings.thumbnails.quality,
|
|
effort=4 # Balance between speed and compression
|
|
)
|
|
|
|
@shared_task(bind=True, name='generate_thumbnails')
|
|
def generate_thumbnails(self, photo_id: str):
|
|
"""Generate thumbnails for a photo"""
|
|
return asyncio.run(_generate_thumbnails_async(photo_id, self))
|
|
|
|
async def _generate_thumbnails_async(photo_id: str, task):
|
|
"""Async implementation of thumbnail generation"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
# Get photo from database
|
|
result = await session.execute(
|
|
select(Photo).where(Photo.id == photo_id)
|
|
)
|
|
photo = result.scalar_one_or_none()
|
|
|
|
if not photo:
|
|
logger.error(f"Photo not found: {photo_id}")
|
|
return {'status': 'error', 'message': 'Photo not found'}
|
|
|
|
# Check if file exists
|
|
if not os.path.exists(photo.filepath):
|
|
logger.error(f"File not found: {photo.filepath}")
|
|
photo.processing_status = 'failed'
|
|
photo.processing_error = 'File not found'
|
|
await session.commit()
|
|
return {'status': 'error', 'message': 'File not found'}
|
|
|
|
# Update processing status
|
|
photo.processing_status = 'processing'
|
|
await session.commit()
|
|
|
|
# Load and process the image based on type
|
|
image = None
|
|
|
|
if photo.media_type == 'photo':
|
|
image = process_standard_image(photo.filepath)
|
|
elif photo.media_type == 'raw':
|
|
image = process_raw_image(photo.filepath)
|
|
elif photo.media_type == 'heic':
|
|
image = process_heic_image(photo.filepath)
|
|
elif photo.media_type == 'video':
|
|
image = process_video_thumbnail(photo.filepath)
|
|
else:
|
|
logger.error(f"Unsupported media type: {photo.media_type}")
|
|
image = create_placeholder_thumbnail(photo.media_type)
|
|
|
|
if not image:
|
|
raise Exception("Failed to process image")
|
|
|
|
# Auto-rotate based on EXIF
|
|
image = auto_rotate_image(image)
|
|
|
|
# Store original dimensions
|
|
photo.width = image.width
|
|
photo.height = image.height
|
|
|
|
# Generate thumbnails for each size
|
|
for size_name, size_value in THUMB_SIZES.items():
|
|
thumb_path = get_thumb_path(photo_id, size_name)
|
|
generate_thumbnail(image, size_value, thumb_path)
|
|
|
|
# Update database with thumbnail path
|
|
setattr(photo, f'thumb_{size_name}', thumb_path)
|
|
|
|
# Update progress
|
|
task.update_state(
|
|
state='PROGRESS',
|
|
meta={'current_size': size_name, 'photo_id': photo_id}
|
|
)
|
|
|
|
# Update processing status
|
|
photo.processing_status = 'completed'
|
|
photo.processing_error = None
|
|
await session.commit()
|
|
|
|
logger.info(f"Thumbnails generated for photo {photo_id}")
|
|
return {'status': 'success', 'photo_id': photo_id}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating thumbnails for {photo_id}: {e}")
|
|
|
|
# Update error status
|
|
if photo:
|
|
photo.processing_status = 'failed'
|
|
photo.processing_error = str(e)
|
|
await session.commit()
|
|
|
|
return {'status': 'error', 'message': str(e)}
|
|
|
|
@shared_task(name='regenerate_all_thumbnails')
|
|
def regenerate_all_thumbnails():
|
|
"""Regenerate thumbnails for all photos"""
|
|
return asyncio.run(_regenerate_all_thumbnails_async())
|
|
|
|
async def _regenerate_all_thumbnails_async():
|
|
"""Async implementation of regenerating all thumbnails"""
|
|
async with AsyncSessionLocal() as session:
|
|
# Get all photos that need thumbnails
|
|
result = await session.execute(
|
|
select(Photo).where(
|
|
Photo.processing_status.in_(['pending', 'failed'])
|
|
)
|
|
)
|
|
photos = result.scalars().all()
|
|
|
|
logger.info(f"Regenerating thumbnails for {len(photos)} photos")
|
|
|
|
for photo in photos:
|
|
generate_thumbnails.delay(photo.id)
|
|
|
|
return {'status': 'queued', 'count': len(photos)} |