pygal/views/image.py

from django.conf import settings
import fstools
import io
import logging
import mimetypes
from ..models import get_item_type, TYPE_IMAGE, TYPE_VIDEO
import os
from PIL import Image, ImageEnhance, ExifTags
import platform
import pygal
import subprocess

# Get a logger instance
logger = logging.getLogger("CACHING")


def get_image_class(full_path):
    return {
        TYPE_IMAGE: image,
        TYPE_VIDEO: video,
    }.get(get_item_type(full_path), other)


def get_image_instance(full_path, request):
    return get_image_class(full_path)(full_path, request)


class mm_image(object):
    JOIN_TOP_LEFT = 1
    JOIN_TOP_RIGHT = 2
    JOIN_BOT_LEFT = 3
    JOIN_BOT_RIGHT = 4
    JOIN_CENTER = 5

    def __init__(self, imagepath_handle_image):
        self.imagepath_handle_image = imagepath_handle_image
        self.__image__ = None
        #

    def __init_image__(self):
        if self.__image__ is None:
            if type(self.imagepath_handle_image) is mm_image:
                self.__image__ = self.imagepath_handle_image.__image__
            elif type(self.imagepath_handle_image) is Image.Image:
                self.__image__ = self.imagepath_handle_image
            else:
                self.__image__ = Image.open(self.imagepath_handle_image)

    def orientation(self):
        self.__init_image__()
        exif_tags = dict((v, k) for k, v in ExifTags.TAGS.items())
        try:
            return dict(self.__image__._getexif().items())[exif_tags['Orientation']]
        except AttributeError:
            return 1
        except KeyError:
            return 1

    def copy(self):
        self.__init_image__()
        return mm_image(self.__image__.copy())

    def save(self, *args, **kwargs):
        self.__init_image__()
        im = self.__image__.copy().convert('RGB')
        im.save(*args, **kwargs)

    def resize(self, max_size):
        self.__init_image__()
        #
        # resize
        #
        x, y = self.__image__.size
        xy_max = max(x, y)
        self.__image__ = self.__image__.resize((int(x * float(max_size) / xy_max), int(y * float(max_size) / xy_max)), Image.NEAREST).rotate(0)

    def rotate_by_orientation(self, orientation):
        self.__init_image__()
        #
        # rotate
        #
        angle = {3: 180, 6: 270, 8: 90}.get(orientation)
        if angle is not None:
            self.__image__ = self.__image__.rotate(angle, expand=True)

    def __rgba_copy__(self):
        self.__init_image__()
        if self.__image__.mode != 'RGBA':
            return self.__image__.convert('RGBA')
        else:
            return self.__image__.copy()

    def join(self, image, joint_pos=JOIN_TOP_RIGHT, opacity=0.7):
        """
        This joins another picture to this one.

        :param picture_edit picture: The picture to be joint.
        :param joint_pos: The position of picture in this picture. See also self.JOIN_*
        :param float opacity: The opacity of picture when joint (value between 0 and 1).

        .. note::
          joint_pos makes only sense if picture is smaller than this picture.
        """
        self.__init_image__()
        #
        im2 = image.__rgba_copy__()
        # change opacity of im2
        alpha = im2.split()[3]
        alpha = ImageEnhance.Brightness(alpha).enhance(opacity)
        im2.putalpha(alpha)

        self.__image__ = self.__rgba_copy__()

        # create a transparent layer
        layer = Image.new('RGBA', self.__image__.size, (0, 0, 0, 0))
        # draw im2 in layer
        if joint_pos == self.JOIN_TOP_LEFT:
            layer.paste(im2, (0, 0))
        elif joint_pos == self.JOIN_TOP_RIGHT:
            layer.paste(im2, ((self.__image__.size[0] - im2.size[0]), 0))
        elif joint_pos == self.JOIN_BOT_LEFT:
            layer.paste(im2, (0, (self.__image__.size[1] - im2.size[1])))
        elif joint_pos == self.JOIN_BOT_RIGHT:
            layer.paste(im2, ((self.__image__.size[0] - im2.size[0]), (self.__image__.size[1] - im2.size[1])))
        elif joint_pos == self.JOIN_CENTER:
            layer.paste(im2, (int((self.__image__.size[0] - im2.size[0]) / 2), int((self.__image__.size[1] - im2.size[1]) / 2)))

        self.__image__ = Image.composite(layer, self.__image__, layer)

    def image_data(self):
        self.__init_image__()
        #
        # create return value
        #
        im = self.__image__.copy().convert('RGB')
        output = io.BytesIO()
        im.save(output, format='JPEG')
        return output.getvalue()


class mm_video(object):
    def __init__(self, full_path):
        self.full_path = full_path

    def image(self):
        if platform.system() == 'Linux':
            cmd = 'ffmpeg -ss 0.5 -i "' + self.full_path + '" -vframes 1 -f image2pipe pipe:1 2> /dev/null'
        else:
            cmd = 'ffmpeg -ss 0.5 -i "' + self.full_path + '" -vframes 1 -f image2pipe pipe:1 2> NULL'
        data = subprocess.check_output(cmd, shell=True)
        ffmpeg_handle = io.BytesIO(data)
        im = Image.open(ffmpeg_handle)
        return mm_image(im.copy())


class base_item(object):
    MIME_TYPES = {
        '.ada': 'text/x/adasrc',
        '.hex': 'text/x/hex',
        '.jpg': 'image/x/generic',
        '.jpeg': 'image/x/generic',
        '.jpe': 'image/x/generic',
        '.png': 'image/x/generic',
        '.tif': 'image/x/generic',
        '.tiff': 'image/x/generic',
        '.gif': 'image/x/generic',
        '.avi': 'video/x/generic',
        '.mpg': 'video/x/generic',
        '.mpeg': 'video/x/generic',
        '.mpe': 'video/x/generic',
        '.mov': 'video/x/generic',
        '.qt': 'video/x/generic',
        '.mp4': 'video/x/generic',
        '.webm': 'video/x/generic',
        '.ogv': 'video/x/generic',
        '.flv': 'video/x/generic',
        '.3gp': 'video/x/generic',
    }

    def __init__(self, full_path, request):
        self.full_path = full_path
        self.request = request
        self.rel_path = pygal.get_rel_path(full_path)
        #
        ext = os.path.splitext(self.full_path)[1].lower()
        self.mime_type = self.MIME_TYPES.get(ext, mimetypes.types_map.get(ext, 'unknown'))

    def __cache_image_folder__(self, rel_path):
        return os.path.join(settings.XNAIL_ROOT, rel_path.replace('_', '__').replace('/', '_'))

    def __cache_image_name__(self, max_size):
        filename = '%04d_%02x_%s.jpg' % (max_size, self.XNAIL_VERSION_NUMBER, fstools.uid(self.full_path, None))
        return os.path.join(self.__cache_image_folder__(self.rel_path), filename)

    def __delete_cache_image__(self, max_size):
        for fn in fstools.filelist(self.__cache_image_folder__(self.rel_path), '%04d*' % max_size):
            try:
                os.remove(fn)
            except OSError:
                pass    # possibly file is already removed by another process


class image(base_item, mm_image):
    XNAIL_VERSION_NUMBER = 1

    def __init__(self, *args, **kwargs):
        base_item.__init__(self, *args, **kwargs)
        mm_image.__init__(self, self.full_path)
        self.mime_type_xnails = mimetypes.types_map.get('.jpg', 'unknown')

    def get_resized_image_data(self, max_size):
        cache_filename = self.__cache_image_name__(max_size)
        if not os.path.exists(cache_filename):
            logger.info('Creating xnail-%d for %s', max_size, self.rel_path)
            self.__delete_cache_image__(max_size)
            im = self.copy()
            im.resize(max_size)
            im.rotate_by_orientation(self.orientation())
            #
            # create cache file
            #
            fstools.mkdir(os.path.dirname(cache_filename))
            with open(cache_filename, 'wb') as fh:
                im.save(fh, format='JPEG')
            #
            return im.image_data()
        else:
            return open(cache_filename, 'rb').read()

    def thumbnail_picture(self):
        return self.get_resized_image_data(pygal.get_thumbnail_max_size(self.request))

    def webnail_picture(self):
        return self.get_resized_image_data(pygal.get_webnail_size(self.request))


class video(base_item, mm_video):
    XNAIL_VERSION_NUMBER = 1

    def __init__(self, *args, **kwargs):
        base_item.__init__(self, *args, **kwargs)
        mm_video.__init__(self, self.full_path)
        self.mime_type_xnails = mimetypes.types_map.get('.jpg', 'unknown')

    def get_resized_image_data(self, max_size):
        cache_filename = self.__cache_image_name__(max_size)
        if not os.path.exists(cache_filename):
            logger.info('Creating xnail-%d for %s', max_size, self.rel_path)
            self.__delete_cache_image__(max_size)
            im = self.image()
            im.resize(max_size)
            overlay = mm_image(os.path.join(os.path.dirname(__file__), 'video.png'))
            im.join(overlay)
            #
            # create cache file
            #
            fstools.mkdir(os.path.dirname(cache_filename))
            with open(cache_filename, 'wb') as fh:
                im.save(fh, format='JPEG')
            #
            return im.image_data()
        else:
            return open(cache_filename, 'rb').read()

    def thumbnail_picture(self):
        return image.thumbnail_picture(self)

    def webnail_picture(self):
        return image.webnail_picture(self)


class other(base_item):
    def __init__(self, *args, **kwargs):
        base_item.__init__(self, *args, **kwargs)
        self.mime_type_xnails = mimetypes.types_map.get('.png', 'unknown')

    def thumbnail_picture(self):
        fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'mimetype_icons', '%s.png' % (self.mime_type).replace('/', '-'))
        if not os.path.exists(fn):
            fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'mimetype_icons', 'unknown.png')
        return open(fn, 'rb').read()

    def webnail_picture(self):
        return self.thumbnail_picture()