aurora-moments/api/apps/library/tasks.py

import io
import json
import os
import subprocess
import tempfile
from datetime import datetime

from celery import shared_task
from django.core.files.base import ContentFile
from PIL import ExifTags, Image

from .models import LibraryItem, MediaRendition

RENDITION_MAX_DIMENSIONS = {
    MediaRendition.Kind.THUMBNAIL: 300,
    MediaRendition.Kind.FEED: 1080,
    MediaRendition.Kind.LARGE: 2048,
}


@shared_task
def process_library_item(library_item_id):
    try:
        item = LibraryItem.objects.get(id=library_item_id)
    except LibraryItem.DoesNotExist:
        return

    item.processing_status = LibraryItem.ProcessingStatus.PROCESSING
    item.save(update_fields=["processing_status"])

    try:
        if item.media_type == LibraryItem.MediaType.PHOTO:
            _process_photo(item)
        else:
            # Full transcoding to web-friendly renditions is a V1.5 feature
            # (see README); the original file is served as-is for playback.
            # We do extract dimensions/duration and a poster-frame thumbnail
            # so videos aren't blank tiles in the UI.
            _process_video(item)
        item.processing_status = LibraryItem.ProcessingStatus.READY
    except Exception:
        item.processing_status = LibraryItem.ProcessingStatus.FAILED
        raise
    finally:
        item.save(update_fields=["processing_status"])


def _process_photo(item):
    with item.file.open("rb") as fh:
        image = Image.open(fh)
        image.load()

        for field, value in _extract_exif(image).items():
            setattr(item, field, value)
        item.width, item.height = image.size
        item.save()

        _write_renditions(item, image)


def _process_video(item):
    suffix = os.path.splitext(item.original_filename)[1] or ".mp4"
    with tempfile.NamedTemporaryFile(suffix=suffix) as src:
        with item.file.open("rb") as fh:
            for chunk in fh.chunks():
                src.write(chunk)
        src.flush()

        probe = _ffprobe(src.name)
        if probe:
            item.width = probe.get("width") or item.width
            item.height = probe.get("height") or item.height
            item.duration_seconds = probe.get("duration")
            item.save()

        # A missing/corrupt frame just means no poster image; the video
        # itself still uploaded fine, so this must not fail the whole job.
        try:
            frame_bytes = _extract_poster_frame(src.name, probe.get("duration") if probe else None)
            if frame_bytes:
                image = Image.open(io.BytesIO(frame_bytes))
                image.load()
                _write_renditions(item, image)
        except Exception:
            pass


def _write_renditions(item, image):
    for kind, max_dimension in RENDITION_MAX_DIMENSIONS.items():
        rendition_image = image.copy()
        rendition_image.thumbnail((max_dimension, max_dimension))
        rgb_image = rendition_image.convert("RGB")
        buffer = io.BytesIO()
        rgb_image.save(buffer, format="JPEG", quality=85)
        MediaRendition.objects.update_or_create(
            library_item=item,
            kind=kind,
            defaults={
                "width": rendition_image.width,
                "height": rendition_image.height,
                "file": ContentFile(buffer.getvalue(), name=f"{item.id}_{kind}.jpg"),
            },
        )


def _ffprobe(path):
    try:
        result = subprocess.run(
            [
                "ffprobe", "-v", "error", "-select_streams", "v:0",
                "-show_entries", "stream=width,height:format=duration",
                "-of", "json", path,
            ],
            capture_output=True, text=True, timeout=30, check=True,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
        return None
    data = json.loads(result.stdout)
    streams = data.get("streams") or [{}]
    duration = (data.get("format") or {}).get("duration")
    return {
        "width": streams[0].get("width"),
        "height": streams[0].get("height"),
        "duration": float(duration) if duration else None,
    }


def _extract_poster_frame(path, duration_seconds):
    timestamp = min(1.0, duration_seconds / 2) if duration_seconds else 0.1
    with tempfile.NamedTemporaryFile(suffix=".jpg") as out:
        try:
            subprocess.run(
                ["ffmpeg", "-y", "-ss", str(timestamp), "-i", path, "-frames:v", "1", "-q:v", "3", out.name],
                capture_output=True, timeout=30, check=True,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
            return None
        return out.read()


def _extract_exif(image):
    result = {}
    raw_exif = image.getexif()
    if not raw_exif:
        return result
    tags = {ExifTags.TAGS.get(tag_id, tag_id): value for tag_id, value in raw_exif.items()}

    if tags.get("Make"):
        result["camera_make"] = str(tags["Make"]).strip()
    if tags.get("Model"):
        result["camera_model"] = str(tags["Model"]).strip()

    date_str = tags.get("DateTimeOriginal") or tags.get("DateTime")
    if date_str:
        try:
            result["captured_at"] = datetime.strptime(str(date_str), "%Y:%m:%d %H:%M:%S")
        except ValueError:
            pass
    return result