hydrus/include/HydrusImageHandling.py

import cStringIO
import numpy.core.multiarray # important this comes before cv!
import cv
import HydrusConstants as HC
from PIL import Image as PILImage
import struct
import threading
import time
import traceback
import wx

#LINEAR_SCALE_PALETTE = [ 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14, 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19, 20, 20, 20, 21, 21, 21, 22, 22, 22, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 29, 29, 29, 30, 30, 30, 31, 31, 31, 32, 32, 32, 33, 33, 33, 34, 34, 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39, 40, 40, 40, 41, 41, 41, 42, 42, 42, 43, 43, 43, 44, 44, 44, 45, 45, 45, 46, 46, 46, 47, 47, 47, 48, 48, 48, 49, 49, 49, 50, 50, 50, 51, 51, 51, 52, 52, 52, 53, 53, 53, 54, 54, 54, 55, 55, 55, 56, 56, 56, 57, 57, 57, 58, 58, 58, 59, 59, 59, 60, 60, 60, 61, 61, 61, 62, 62, 62, 63, 63, 63, 64, 64, 64, 65, 65, 65, 66, 66, 66, 67, 67, 67, 68, 68, 68, 69, 69, 69, 70, 70, 70, 71, 71, 71, 72, 72, 72, 73, 73, 73, 74, 74, 74, 75, 75, 75, 76, 76, 76, 77, 77, 77, 78, 78, 78, 79, 79, 79, 80, 80, 80, 81, 81, 81, 82, 82, 82, 83, 83, 83, 84, 84, 84, 85, 85, 85, 86, 86, 86, 87, 87, 87, 88, 88, 88, 89, 89, 89, 90, 90, 90, 91, 91, 91, 92, 92, 92, 93, 93, 93, 94, 94, 94, 95, 95, 95, 96, 96, 96, 97, 97, 97, 98, 98, 98, 99, 99, 99, 100, 100, 100, 101, 101, 101, 102, 102, 102, 103, 103, 103, 104, 104, 104, 105, 105, 105, 106, 106, 106, 107, 107, 107, 108, 108, 108, 109, 109, 109, 110, 110, 110, 111, 111, 111, 112, 112, 112, 113, 113, 113, 114, 114, 114, 115, 115, 115, 116, 116, 116, 117, 117, 117, 118, 118, 118, 119, 119, 119, 120, 120, 120, 121, 121, 121, 122, 122, 122, 123, 123, 123, 124, 124, 124, 125, 125, 125, 126, 126, 126, 127, 127, 127, 128, 128, 128, 129, 129, 129, 130, 130, 130, 131, 131, 131, 132, 132, 132, 133, 133, 133, 134, 134, 134, 135, 135, 135, 136, 136, 136, 137, 137, 137, 138, 138, 138, 139, 139, 139, 140, 140, 140, 141, 141, 141, 142, 142, 142, 143, 143, 143, 144, 144, 144, 145, 145, 145, 146, 146, 146, 147, 147, 147, 148, 148, 148, 149, 149, 149, 150, 150, 150, 151, 151, 151, 152, 152, 152, 153, 153, 153, 154, 154, 154, 155, 155, 155, 156, 156, 156, 157, 157, 157, 158, 158, 158, 159, 159, 159, 160, 160, 160, 161, 161, 161, 162, 162, 162, 163, 163, 163, 164, 164, 164, 165, 165, 165, 166, 166, 166, 167, 167, 167, 168, 168, 168, 169, 169, 169, 170, 170, 170, 171, 171, 171, 172, 172, 172, 173, 173, 173, 174, 174, 174, 175, 175, 175, 176, 176, 176, 177, 177, 177, 178, 178, 178, 179, 179, 179, 180, 180, 180, 181, 181, 181, 182, 182, 182, 183, 183, 183, 184, 184, 184, 185, 185, 185, 186, 186, 186, 187, 187, 187, 188, 188, 188, 189, 189, 189, 190, 190, 190, 191, 191, 191, 192, 192, 192, 193, 193, 193, 194, 194, 194, 195, 195, 195, 196, 196, 196, 197, 197, 197, 198, 198, 198, 199, 199, 199, 200, 200, 200, 201, 201, 201, 202, 202, 202, 203, 203, 203, 204, 204, 204, 205, 205, 205, 206, 206, 206, 207, 207, 207, 208, 208, 208, 209, 209, 209, 210, 210, 210, 211, 211, 211, 212, 212, 212, 213, 213, 213, 214, 214, 214, 215, 215, 215, 216, 216, 216, 217, 217, 217, 218, 218, 218, 219, 219, 219, 220, 220, 220, 221, 221, 221, 222, 222, 222, 223, 223, 223, 224, 224, 224, 225, 225, 225, 226, 226, 226, 227, 227, 227, 228, 228, 228, 229, 229, 229, 230, 230, 230, 231, 231, 231, 232, 232, 232, 233, 233, 233, 234, 234, 234, 235, 235, 235, 236, 236, 236, 237, 237, 237, 238, 238, 238, 239, 239, 239, 240, 240, 240, 241, 241, 241, 242, 242, 242, 243, 243, 243, 244, 244, 244, 245, 245, 245, 246, 246, 246, 247, 247, 247, 248, 248, 248, 249, 249, 249, 250, 250, 250, 251, 251, 251, 252, 252, 252, 253, 253, 253, 254, 254, 254, 255, 255, 255 ]

def ConvertToPngIfBmp( file ):
    
    if file.startswith( 'BM' ):
        
        pil_image = GeneratePILImageFromFile( file )
        
        f = cStringIO.StringIO()
        
        pil_image = pil_image.convert( 'P' )
        
        pil_image.save( f, 'PNG' )
        
        f.seek( 0 )
        
        file = f.read()
        
        f.close()
        
    
    return file
    
def EfficientlyResizeImage( pil_image, ( x, y ) ):
    
    ( im_x, im_y ) = pil_image.size
    
    if x >= im_x and y >= im_y: return pil_image
    
    if pil_image.mode == 'RGB': # low quality resize screws up alpha channel!
        
        if im_x > 2 * x and im_y > 2 * y: pil_image.thumbnail( ( 2 * x, 2 * y ), PILImage.NEAREST )
        
    
    return pil_image.resize( ( x, y ), PILImage.ANTIALIAS )
    
def EfficientlyThumbnailImage( pil_image, ( x, y ) ):
    
    ( im_x, im_y ) = pil_image.size
    
    if pil_image.mode == 'RGB': # low quality resize screws up alpha channel!
        
        if im_x > 2 * x or im_y > 2 * y: pil_image.thumbnail( ( 2 * x, 2 * y ), PILImage.NEAREST )
        
    
    pil_image.thumbnail( ( x, y ), PILImage.ANTIALIAS )
    
def GenerateAnimatedFrame( pil_image, target_resolution, canvas ):
    
    if 'duration' not in pil_image.info: duration = 40 # 25 fps default when duration is missing or too funky to extract. most stuff looks ok at this.
    else:
        
        duration = pil_image.info[ 'duration' ]
        
        if duration == 0: duration = 40
        
    
    current_frame = EfficientlyResizeImage( pil_image, target_resolution )
    
    if pil_image.mode == 'P' and 'transparency' in pil_image.info:
        
        # I think gif problems are around here somewhere; the transparency info is not converted to RGBA properly, so it starts drawing colours when it should draw nothing
        
        current_frame = current_frame.convert( 'RGBA' )
        
        if canvas is None: canvas = current_frame
        else: canvas.paste( current_frame, None, current_frame ) # yeah, use the rgba image as its own mask, wut.
        
    else: canvas = current_frame
    
    return ( canvas, duration )
    
def GenerateHydrusBitmapFromFile( file ):
    
    pil_image = GeneratePILImageFromFile( file )
    
    return GenerateHydrusBitmapFromPILImage( pil_image )
    
def GenerateHydrusBitmapFromPILImage( pil_image ):
    
    if pil_image.mode == 'RGBA' or ( pil_image.mode == 'P' and pil_image.info.has_key( 'transparency' ) ):
        
        if pil_image.mode == 'P': pil_image = pil_image.convert( 'RGBA' )
        
        return HydrusBitmap( pil_image.tostring(), wx.BitmapBufferFormat_RGBA, pil_image.size )
        
    else:
        
        if pil_image.mode != 'RGB': pil_image = pil_image.convert( 'RGB' )
        
        return HydrusBitmap( pil_image.tostring(), wx.BitmapBufferFormat_RGB, pil_image.size )
        
    
def GeneratePerceptualHash( thumbnail_file ):
    
    thumbnail = GeneratePILImageFromFile( thumbnail_file )
    
    # convert to 32 x 32 greyscale
    
    if thumbnail.mode == 'RGBA':
        
        # this is some code i picked up somewhere
        # another great example of PIL failing; it turns all alpha to pure black on a RGBA->RGB
        
        thumbnail.load()
        
        canvas = PILImage.new( 'RGB', thumbnail.size, ( 255, 255, 255 ) )
        
        canvas.paste( thumbnail, mask = thumbnail.split()[3] )
        
        thumbnail = canvas
        
    
    thumbnail = thumbnail.convert( 'L' )
    
    thumbnail = thumbnail.resize( ( 32, 32 ), PILImage.ANTIALIAS )
    
    # convert to mat
    
    cv_thumbnail_8 = cv.CreateMatHeader( 32, 32, cv.CV_8UC1 )
    
    cv.SetData( cv_thumbnail_8, thumbnail.tostring() )
    
    cv_thumbnail_32 = cv.CreateMat( 32, 32, cv.CV_32FC1 )
    
    cv.Convert( cv_thumbnail_8, cv_thumbnail_32 )
    
    # compute dct
    
    dct = cv.CreateMat( 32, 32, cv.CV_32FC1 )
    
    cv.DCT( cv_thumbnail_32, dct, cv.CV_DXT_FORWARD )
    
    # take top left 8x8 of dct
    
    dct = cv.GetSubRect( dct, ( 0, 0, 8, 8 ) )
    
    # get mean of dct, excluding [0,0]
    
    mask = cv.CreateMat( 8, 8, cv.CV_8U )
    
    cv.Set( mask, 1 )
    
    mask[0,0] = 0
    
    channel_averages = cv.Avg( dct, mask )
    
    average = channel_averages[0]
    
    # make a monochromatic, 64-bit hash of whether the entry is above or below the mean
    
    bytes = []
    
    for i in range( 8 ):
        
        byte = 0
        
        for j in range( 8 ):
            
            byte <<= 1 # shift byte one left
            
            value = dct[i,j]
            
            if value > average: byte |= 1
            
        
        bytes.append( byte )
        
    
    answer = str( bytearray( bytes ) )
    
    # we good
    
    return answer
    
def GeneratePILImageFromFile( file ): return PILImage.open( cStringIO.StringIO( file ) )

def GenerateResolutionAndNumFramesFromFile( file ):
    
    pil_image = GeneratePILImageFromFile( file )
    
    ( x, y ) = pil_image.size
    
    try:
        
        pil_image.seek( 1 )
        pil_image.seek( 0 )
        
        num_frames = 1
        
        while True:
            
            try:
                
                pil_image.seek( pil_image.tell() + 1 )
                num_frames += 1
                
            except: break
            
        
    except: num_frames = 1
    
    return ( ( x, y ), num_frames )
    
def GenerateThumbnailFileFromFile( file, dimensions ):
    
    pil_image = GeneratePILImageFromFile( file )
    
    return GenerateThumbnailFileFromImage( pil_image, dimensions )
    
def GenerateThumbnailFileFromImage( pil_image, dimensions = None ):
    
    if dimensions is None: dimensions = HC.UNSCALED_THUMBNAIL_DIMENSIONS
    
    EfficientlyThumbnailImage( pil_image, dimensions )
    
    f_t = cStringIO.StringIO()
    
    if pil_image.mode == 'P' and pil_image.info.has_key( 'transparency' ):
        
        pil_image.save( f_t, 'PNG', transparency = pil_image.info[ 'transparency' ] )
        
    elif pil_image.mode == 'RGBA': pil_image.save( f_t, 'PNG' )
    else:
        
        pil_image = pil_image.convert( 'RGB' )
        
        pil_image.save( f_t, 'JPEG', quality=92 )
        
    
    f_t.seek( 0 )
    
    thumbnail = f_t.read()
    
    f_t.close()
    
    return thumbnail
    
def GetHammingDistance( phash1, phash2 ):
    
    distance = 0
    
    phash1 = bytearray( phash1 )
    phash2 = bytearray( phash2 )
    
    for i in range( len( phash1 ) ):
        
        xor = phash1[i] ^ phash2[i]
        
        while xor > 0:
            
            distance += 1
            xor &= xor - 1
            
        
    
    return distance
    
def RenderImageFrames( image_container, file, target_resolution, synchronous ):
    
    if image_container.IsAnimated(): renderer = AnimatedFrameRenderer( image_container, file, target_resolution )
    else: renderer = StaticFrameRenderer( image_container, file, target_resolution )
    
    if synchronous: renderer.Render()
    else: threading.Thread( target = renderer.RenderCallAfter ).start()
    
def RenderImageFromFile( file, hash, target_resolution = None, synchronous = True ):
    
    try:
        
        ( original_resolution, num_frames ) = GenerateResolutionAndNumFramesFromFile( file )
        
        if target_resolution is None: target_resolution = original_resolution
        
        image_container = RenderedImageContainer( hash, original_resolution, target_resolution, num_frames )
        
        RenderImageFrames( image_container, file, target_resolution, synchronous )
        
        return image_container
        
    except: raise Exception( 'Attempted to render the image, but it was either formatted slightly incorrectly or PIL could not handle it; look up PIL in the hydrus help for more info.' )
    
class FrameRenderer():
    
    def __init__( self, image_container, file, target_resolution ):
        
        self._image_container = image_container
        self._pil_image = GeneratePILImageFromFile( file )
        self._target_resolution = target_resolution
        
    
class AnimatedFrameRenderer( FrameRenderer ):
    
    def GetFrames( self ):
        
        canvas = None
        
        global_palette = self._pil_image.palette
        
        dirty = self._pil_image.palette.dirty
        mode = self._pil_image.palette.mode
        rawmode = self._pil_image.palette.rawmode
        
        # believe it or not, doing this actually fixed a couple of gifs!
        self._pil_image.seek( 1 )
        self._pil_image.seek( 0 )
        
        while True:
            
            ( canvas, duration ) = GenerateAnimatedFrame( self._pil_image, self._target_resolution, canvas )
            
            yield ( GenerateHydrusBitmapFromPILImage( canvas ), duration )
            
            try:
                
                self._pil_image.seek( self._pil_image.tell() + 1 )
                
                if self._pil_image.palette == global_palette: # for some reason, when we fall back to global palette (no local-frame palette), we reset bunch of important variables!
                    
                    self._pil_image.palette.dirty = dirty
                    self._pil_image.palette.mode = mode
                    self._pil_image.palette.rawmode = rawmode
                    
                
            except: break
            
        
    
    def Render( self ):
        
        for ( frame, duration ) in self.GetFrames(): self._image_container.AddFrame( frame, duration )
        
    
    def RenderCallAfter( self ):
        
        time.sleep( 0 ) # thread yield
        
        for ( frame, duration ) in self.GetFrames(): wx.CallAfter( self._image_container.AddFrame, frame, duration )
        
        HC.pubsub.pub( 'finished_rendering', self._image_container.GetKey() )
        
    
class StaticFrameRenderer( FrameRenderer ):
    
    def GetFrame( self ): return GenerateHydrusBitmapFromPILImage( EfficientlyResizeImage( self._pil_image, self._target_resolution ) )
    
    def Render( self ): self._image_container.AddFrame( self.GetFrame() )
    
    def RenderCallAfter( self ):
        
        time.sleep( 0 ) # thread yield
        
        wx.CallAfter( self._image_container.AddFrame, self.GetFrame() )
        
        HC.pubsub.pub( 'finished_rendering', self._image_container.GetKey() )
        
    
class HydrusBitmap():
    
    def __init__( self, data, format, size ):
        
        self._data = data
        self._format = format
        self._size = size
        
    
    def CreateWxBmp( self ):
        
        ( width, height ) = self._size
        
        if self._format == wx.BitmapBufferFormat_RGB: return wx.BitmapFromBuffer( width, height, self._data )
        else: return wx.BitmapFromBufferRGBA( width, height, self._data )
        
    
    def GetEstimatedMemoryFootprint( self ): return len( self._data )
    
    def GetSize( self ): return self._size
    
class RenderedImageContainer():
    
    def __init__( self, hash, original_resolution, my_resolution, num_frames ):
        
        self._hash = hash
        self._original_resolution = original_resolution
        self._my_resolution = my_resolution
        self._num_frames = num_frames
        
        ( original_width, original_height ) = original_resolution
        
        ( my_width, my_height ) = my_resolution
        
        width_zoom = my_width / float( original_width )
        height_zoom = my_height / float( original_height )
        
        self._zoom = min( ( width_zoom, height_zoom ) )
        
        if self._zoom > 1.0: self._zoom = 1.0
        
        self._frames = []
        self._durations = []
        
        self._finished_rendering = False
        
    
    def AddFrame( self, frame, duration = None ):
        
        self._frames.append( frame )
        
        if duration is not None: self._durations.append( duration )
        
    
    def GetDuration( self, index ): return self._durations[ index ]
    
    def GetEstimatedMemoryFootprint( self ): return sum( [ frame.GetEstimatedMemoryFootprint() for frame in self._frames ] )
    
    def GetFrame( self, index = None ):
        
        if index is None: return self._frames[ 0 ]
        else: return self._frames[ index ]
        
    
    def GetHash( self ): return self._hash
    
    def GetKey( self ): return ( self._hash, self._my_resolution )
    
    def GetNumFrames( self ): return self._num_frames
    
    def GetResolution( self ): return self._original_resolution
    
    def GetSize( self ): return self._my_resolution
    
    def GetTotalDuration( self ): return sum( self._durations )
    
    def GetZoom( self ): return self._zoom
    
    def HasFrame( self, index = None ):
        
        if index is None: index = 0
        
        return len( self._frames ) > index
        
    
    def IsAnimated( self ): return self._num_frames > 1
    
    def IsFinishedRendering( self ): return len( self._frames ) == self._num_frames
    
    def IsScaled( self ): return self._zoom != 1.0