photolib/photoapp/thumb.py

111 lines
3.7 KiB
Python

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