feat: structure
This commit is contained in:
277
backend/app/tasks/thumbs.py
Normal file
277
backend/app/tasks/thumbs.py
Normal file
@@ -0,0 +1,277 @@
|
||||
"""
|
||||
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)}
|
||||
Reference in New Issue
Block a user