photolib/photoapp/thumb.py

import os
import sys
import traceback
from time import time
from collections import defaultdict
from multiprocessing import Process
from PIL import Image, ImageOps
import tempfile
from shutil import copyfileobj
from contextlib import closing
import logging


# style tuples: max x, max y, rotate ok)
# rotate ok means x and y maxes can be swapped if it fits the image's aspect ratio better
THUMB_STYLES = {"tiny": (80, 80, False),
                "small": (100, 100, False),
                "feed": (250, 250, False),
                "preview": (1024, 768, True),
                "big": (2048, 1536, True)}


def gen_thumb(src_img, dest_img, width, height, rotation):
    start = time()
    # TODO lock around the dir creation
    os.makedirs(os.path.split(dest_img)[0], exist_ok=True)
    image = Image.open(src_img)
    if image.mode != "RGB":
        image = image.convert("RGB")
    image = image.rotate(90 * rotation, expand=True)
    thumb = ImageOps.fit(image, (width, height), Image.ANTIALIAS)
    thumb.save(dest_img, 'JPEG')
    logging.info("Generated {} in {}s".format(dest_img, round(time() - start, 4)))
    return dest_img


def gen_thumb_process(src_img, dest_img, width, height, rotation):
    try:
        return gen_thumb(src_img, dest_img, width, height, rotation)
    except Exception:
        traceback.print_exc()
        if os.path.exists(dest_img):
            os.unlink(dest_img)
        sys.exit(1)


def thumb_path(style_name, uuid):
    return os.path.join("thumbs", style_name, uuid[0:2], "{}.jpg".format(uuid))


class ThumbGenerator(object):
    def __init__(self, library, storage):
        self.library = library
        self.storage = storage
        self._failed_thumbs_cache = defaultdict(dict)

    def make_photo_thumb(self, photo, style_name):
        """
        Create a thumbnail of the given photo, scaled/cropped to the given named style
        :return: local path to thumbnail file or None if creation failed or was blocked
        """

        dest = thumb_path(style_name, photo.uuid)


        if self.storage.exists(dest):
            return self.storage.open(dest, "rb")


        if photo.uuid in self._failed_thumbs_cache[style_name]:
            return None


        # if photo.width is None:  # todo better detection of images that PIL can't open
        #     return None

        thumb_width, thumb_height, flip_ok = THUMB_STYLES[style_name]
        i_width = photo.width
        i_height = photo.height
        im_is_rotated = photo.orientation % 2 != 0 or i_height > i_width

        if im_is_rotated and flip_ok:
            thumb_width, thumb_height = thumb_height, thumb_width

        thumb_width = min(thumb_width, i_width if i_width > 0 else 999999999)
        thumb_height = min(thumb_height, i_height if i_height > 0 else 999999999)


        # TODO have the subprocess download the file
        # TODO thundering herd
        with tempfile.TemporaryDirectory() as tmpdir:
            fthumblocal = tempfile.NamedTemporaryFile(delete=True)
            fpath = os.path.join(tmpdir, "image")
            with self.library.storage.open(photo.path, 'rb') as fsrc:
                with open(fpath, 'wb') as ftmpdest:
                    copyfileobj(fsrc, ftmpdest)

            p = Process(target=gen_thumb_process,
                        args=(fpath, fthumblocal.name, thumb_width, thumb_height, photo.orientation))
            p.start()
            p.join()
            if p.exitcode != 0:
                self._failed_thumbs_cache[style_name][photo.uuid] = True  # dont retry failed generations
                return None

            with closing(self.storage.open(dest, 'wb')) as fdest:
                copyfileobj(fthumblocal, fdest)

        fthumblocal.seek(0)
        return fthumblocal