_miam/_miamG/miam/image/Image.py

# import
# ------------------------------------------------------------------------------------------
from . import Exif, ColorSpace, imageType, channel
import miam.processing.TMO_CCTF
import miam.processing.ColorSpaceTransform as MCST
import rawpy, colour, os, skimage.transform, copy
import numpy as np
import json
import matplotlib.pyplot as plt
# ------------------------------------------------------------------------------------------
# MIAM project 2020
# ------------------------------------------------------------------------------------------
# author: remi.cozot@univ-littoral.fr
# ------------------------------------------------------------------------------------------

class Image(object):
    """ 
    class image:
        attribute(s):
            name: image file name (str)
            colorData:  array of pixels (np.ndarray)
            type: type of image (image.imageType)
            linear: image encoding is linear (boolean)
            colorSpace: colorspace (colour.models.RGB_COLOURSPACES)
            scalingFactor: scaling factor to [0..1] space (float)                      
    """

    def __init__(self, colorData, name, type, linear, colorspace, scalingFactor):
        """ constructor """

        self.name           = name
        self.colorData      = colorData
        self.shape          = self.colorData.shape
        self.type           = type
        self.linear         = linear
        self.colorSpace     = colorspace
        self.scalingFactor  = scalingFactor

        # ----------------------------------------------------
        # update May 2020
        # ----------------------------------------------------
        # mask mask shape is image shape with only one channel 
        self.mask           = None  
        # ----------------------------------------------------
        # ----------------------------------------------------


    def isHDR(self): return self.type == imageType.imageType.HDR

    def __repr__(self):

        res =   " Image{ name:"                 + self.name                     + "\n"  + \
                "         shape: "              + str(self.shape)               + "\n"  + \
                "         type: "               + str(self.type)                + "\n"  + \
                "         color space: "        + self.colorSpace.__repr__()    + "\n"  + \
                "         linear: "             + str(self.linear)              + "\n"  + \
                "         scaling factor: "     + str(self.scalingFactor)       + "\n }"  
        return res

    def __str__(self) : 

        res =   " Image{ name:"                 + self.name                     + "\n"  + \
                "         shape: "              + str(self.shape)               + "\n"  + \
                "         type: "               + str(self.type)                + "\n"  + \
                "         color space: "        + self.colorSpace.name          + "\n"  + \
                "         linear: "             + str(self.linear)              + "\n"  + \
                "         scaling factor: "     + str(self.scalingFactor)       + "\n }"  
        return res

    def writeImage(self, filename):
        """
        Write Image: two output files are generated: filename+".json" contains additional data, filename+".jpg|hdr" contains colorData
        @params:
            filename   - Required  : filename WITHOUT extension [for example "../images/image00"] (Str)
        """
        # outfile depend on image type
        ext, bitDepth = None, None
        if self.type == imageType.imageType.SDR:
            ext, bitDepth = "jpg", "uint8"
        elif self.type == imageType.imageType.HDR:
            ext, bitDepth = "hdr", "float32"
        else:
            print("WARNING[miam.image.Image.writeImage(",filename,"):: writing RAW image not yet supported!]")

        (fpath, fname) = os.path.split(filename)
        output = {  'name':         self.name,
                    'filename':     fname+"."+ext,
                    'shape':        str(self.shape),
                    'type':         str(self.type),
                    'colorspace':   self.colorSpace.name,
                    'linear':       str(self.linear),
                    'scaling':      str(self.scalingFactor) 
                  }
        with open(filename+".json", 'w') as outfile:
            json.dump(output, outfile, indent=4)

        colour.write_image(self.colorData, 
                           filename+"."+ext,
                           bit_depth=bitDepth, 
                           method='Imageio')

    def getChannelVector(self, channel):
        """ get channel vector: works only for sR|sG|sB, X|Y|Z and L|a|b """

        # take into account colorSpace
        destColor = channel.colorSpace()
        image = MCST.ColorSpaceTransform().compute(self,dest=destColor)

        if channel.getValue() <3:
            imgVector = Image.array2vector(image.colorData)
            return imgVector[:,channel.getValue()]
        else:
            return None

    def getChannel(self, channel):
        """ get channel : works only for sR|sG|sB, X|Y|Z and L|a|b """

        # take into account colorSpace
        destColor = channel.colorSpace()
        image = MCST.ColorSpaceTransform().compute(self,dest=destColor)

        if channel.getValue() <3:
            return image.colorData[:,:,channel.getValue()]
        else:
            return None

    def getMinMaxPerChannel(self):

        img = self.colorData
        R, G, B = img[:,:,0], img[:,:,1], img[:,:,2]

        minR, minG, minB = np.amin(R), np.amin(G), np.amin(B)
        maxR, maxG, maxB = np.amax(R), np.amax(G), np.amax(B)

        return ((minR,maxR),(minG,maxG),(minB,maxB))

    def getDynamicRange(self, percentile = None, mode=None):
        """ return dynamic range of image
          @params:
            percentile  - Optional  : percentile if None: just remove zero values   (Float)
            mode        - Optional  : "maxmin" (default)| "f-stops"                 (Str)      
        """

        Y_min,Y_max = None, None

        Y = self.getChannelVector(channel.channel.Y)

        if percentile == None :
            # remove zeros
            Y = vY[vY>0]
            # use min and max
            Y_min = np.amin(Y)
            Y_max = np.amax(Y)
        else:
            Y_min = np.percentile(Y,percentile)
            Y_max = np.percentile(Y,100-percentile)

        # take mode into account
        if not mode: mode = 'maxmin'

        if mode == "f-stops":
            return np.log2(Y_max)-np.log2(Y_min)

        if mode=="maxmin":
            return (Y_max,Y_min)

    # functionnal methods
    def removeZeros(self, minPencitile=None):

        # copy gimage
        res = copy.deepcopy(self)
        img = res.colorData        
        
        # get channels
        R, G, B = img[:,:,0], img[:,:,1], img[:,:,2]

        # remove zero value per channel
        notZeroR, notZeroG, notZeroB = R[R>0], G[G>0],B[B>0]

        if minPencitile==None:
            # take none zero min of channels
            Rmin, Gmin, Bmin = np.amin(notZeroR), np.amin(notZeroG), np.amin(notZeroB)

            # replace zero values by min (per channel)
            R[R==0] = Rmin 
            G[G==0] = Gmin 
            B[B==0] = Bmin
        else:
            # replace zeros by min Pecentile
            Rmin = np.percentile(R,minPencitile)
            Gmin = np.percentile(G,minPencitile)
            Bmin = np.percentile(B,minPencitile)

            R[R<Rmin] = Rmin
            G[G<Gmin] = Gmin
            B[B<Bmin] = Bmin

        # combine channels
        img[:,:,0], img[:,:,1], img[:,:,2] = R, G, B
        res.colorData = img

        return res

    def clip(self, min=0.0, max=1.0):
        res = copy.deepcopy(self)
        res.colorData[res.colorData>max] = max
        res.colorData[res.colorData<min] = min
        return res

    def smartResize(self,maxSize=400,anti_aliasing=False):
        res = copy.deepcopy(self)
        x, y, c =  tuple(res.colorData.shape)
        maxImage = max(x,y)
        factor = maxImage/maxSize if maxImage > maxSize else 1
        res.colorData = skimage.transform.resize(res.colorData, (x // factor, y // factor), anti_aliasing)
        res.shape = res.colorData.shape
        return res

    def resize(self,size=(None,None),anti_aliasing=False):
        res = copy.deepcopy(self)
        x, y, c =  tuple(res.colorData.shape)
        ny,nx = size
        if nx and (not ny): 
            factor = nx/x
            res.colorData = skimage.transform.resize(res.colorData, (nx, int(y * factor)), anti_aliasing)
            res.shape = res.colorData.shape
        elif (not nx) and ny:
            factor = ny/y
            res.colorData = skimage.transform.resize(res.colorData, (int(x * factor),ny), anti_aliasing)
            res.shape = res.colorData.shape
        elif nx and ny:
            res.colorData = skimage.transform.resize(res.colorData, (nx,ny), anti_aliasing)
            res.shape = res.colorData.shape
        return res

    def process(self, process, **kwargs):
        """ process image: 
            (Object) process must have compute(image, **kwargs) method
            **kwargs optionnal parameters packed as dict
        """
        return process.compute(self,**kwargs)
    
    def buildMaskByValue(self, channel, min, max, zero = 0, one=1):

        # convert into channel color space
        ch = MCST.ColorSpaceTransform().compute(self,dest=channel.colorSpace()).colorData[:,:,channel.getValue()]

        # mask
        ch[ch<=min] = zero
        ch[ch>=max] = zero
        ch[ch!=zero] = one

        res = Image.newImage(self.shape, colorSpaceName='sRGB',name='mask')
        res = copy.deepcopy(self)
        res.colorData[:,:,0] = ch
        res.colorData[:,:,1] = ch
        res.colorData[:,:,2] = ch

        return res

    # ----------------------------------------------------
    # update May 2020
    # ----------------------------------------------------
    def addMask(self, one=True):
        [height, width, channel] = self.shape
        self.mask = np.ones((height, width)) if one else np.zeros((height, width))

    def removeMask(self) : self.mask = None

    def hasMask(self): return isinstance(self.mask,np.ndarray)

    def binaryMask(self): return (self.mask[self.mask!=1] ==0).all()

    def isMaskOne(self): return (self.mask ==1).all()
    # ----------------------------------------------------
    # ----------------------------------------------------

    # plot
    def plot(self, ax, shortName=True, title=True):
        if not (self.type == imageType.imageType.HDR):
            ax.imshow(self.colorData)
            sep = '/' if ('/' in self.name) else '\\'
            name = self.name.split(sep)[-1] if shortName else self.name
            if title: ax.set_title(name+"("+self.colorSpace.name +"/"+ str(self.type)+")")

        else: # HDR
            ax.imshow(self.process(miam.processing.TMO_CCTF.TMO_CCTF(), function='sRGB').colorData)
            sep = '/' if ('/' in self.name) else '\\'
            name = self.name.split(sep)[-1] if shortName else self.name
            if title : ax.set_title(name+"("+self.colorSpace.name +"/"+ str(self.type)+"/"+ "sRGB_CCTF" +")")

        ax.axis("off")

    # magic operators
    def __add__(self, other):
        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other,type(self)): # both are images
            if self.shape == other.shape: # both share same size

                if (not self.hasMask()) and (not other.hasMask()): # none has mask
                    res.colorData = self.colorData + other.colorData

                else: # self or other has a mask

                    if not self.hasMask(): self.addMask()
                    if not other.hasMask(): other.addMask()

                    # both have a mask
                    res.colorData = self.colorData*self.mask[...,np.newaxis] \
                        + other.colorData*other.mask[...,np.newaxis]

                    # mask
                    res.mask = self.mask + other.mask

                    # regularization
                    res.colorData[res.mask>1] = res.colorData[res.mask>1] / res.mask[res.mask>1,np.newaxis] 

                    res.mask[res.mask>1] = 1
            else:
                print("WARNING[Image.__add__(self,other): both image must have the same shape ! return a copy of ",self,"]")
            
            # hdr or sdr
            res.type = imageType.imageType.HDR if (self.isHDR() or other.isHDR()) else imageType.imageType.SDR

        elif isinstance(other, (int, float)): # sum with int or float
            res.colorData = self.colorData + other

        if not self.isHDR(): # if not HDR clip value
            res.colorData[res.colorData>1] = 1
            res.colorData[res.colorData<0] = 0

        return res

    def __radd__(self, other):  return self.__add__(other)

    def __sub__(self, other):
        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other,type(self)): # both are image
            if self.shape == self.other: # same size

                if (not self.hasMask()) and (not other.hasMask()): # mask management
                    res.colorData = self.colorData - other.colorData

                else: # self or other has a mask
                    if not self.hasMask(): self.addMask()
                    if not other.hasMask(): other.addMask()
                    # both have a mask
                    res.colorData = self.colorData*self.mask[...,np.newaxis] \
                        - other.colorData*other.mask[...,np.newaxis]
                    # mask
                    res.mask = self.mask + other.mask[...,np.newaxis]
                    res.mask[res.mask>1] = 1            
            else:
                print("WARNING[Image.__sub__(self,other): both image must have the same shape ! return a copy of ",self,"]")

        elif isinstance(other, (int, float)):
            res.colorData = self.colorData + other

        return res

    def __rsub__(self, other):  
        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other,type(self)):
            if self.shape == self.other:
                res.colorData = other.colorData - self.colorData
            else:
                print("WARNING[Image.__sub__(self,other): both image must have the same shape ! return a copy of ",self,"]")

        elif isinstance(other, (int, float)):
            res.colorData = other - self.colorData

        return res

    def __mul__ (self, other):
        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other,type(self)):
            if self.shape == self.other:
                res.colorData = self.colorData * other.colorData
            else:
                print("WARNING[Image.__mul__(self,other): both image must have the same shape ! return a copy of ",self,"]")

        elif isinstance(other, (int, float)):
            res.colorData = self.colorData * other

        return res

    def __rmul__ (self, other): return self.__mul__(other)

    def __pow__(self,other):
        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other, (int, float)):
            res.colorData = self.colorData**other

        return res

    def __rpow__(self,other): return self.__pow__(other)

    # class methods
    def readImage(filename,readExif=True):

        # default values
        scalingFactor = 1.0
        type = None
        linear = None

        # image name
        path, name = os.path.split(filename)

        if readExif:
            # reading metadata then build  exposure and colour space from exif
            exif        = Exif.Exif.buildFromFileName(filename)
            colorspace = ColorSpace.ColorSpace.buildFromExif(exif)
        else:
            colorspace = ColorSpace.ColorSpace.build('sRGB')

        # extension sensitive
        splits = filename.split('.')
        ext = splits[-1].lower()

        # load raw file using rawpy
        if ext=="arw" or ext=="dng":
            outBit = 16
            raw = rawpy.imread(filename)
            ppParams = rawpy.Params(demosaic_algorithm=None, half_size=False, 
                                            four_color_rgb=False, dcb_iterations=0, 
                                            dcb_enhance=False, fbdd_noise_reduction=rawpy.FBDDNoiseReductionMode.Off, 
                                            noise_thr=None, median_filter_passes=0, 
                                            use_camera_wb=True,                                 # default False
                                            use_auto_wb=False, 
                                            user_wb=None, 
                                            output_color=rawpy.ColorSpace.sRGB,                 # output in SRGB
                                            output_bps=outBit,                                  # default 8
                                            user_flip=None, user_black=None, 
                                            user_sat=None, no_auto_bright=False, 
                                            auto_bright_thr=None, adjust_maximum_thr=0.75, 
                                            bright=1.0, highlight_mode=rawpy.HighlightMode.Clip, 
                                            exp_shift=None, exp_preserve_highlights=0.0, 
                                            no_auto_scale=False,
                                            gamma=None,                                         # linear output
                                            chromatic_aberration=None, bad_pixels_path=None)
            imgDouble = colour.utilities.as_float_array(raw.postprocess(ppParams))/(pow(2,16)-1)
            raw.close()
            type = imageType.RAW
            linear = True

        # load jpg, tiff, hdr file using colour
        else:
            imgDouble = colour.read_image(filename, bit_depth='float32', method='Imageio')
            imgDouble = Image.forceColorData3(imgDouble)
            type = imageType.imageType.SDR
            linear = False

        # post processing for HDR scaling to [ ,1]
        if ext =="hdr":
            imgDouble, scalingFactor = Image.scaleMaxOne(imgDouble)
            type = imageType.imageType.HDR
            linear = True

        #return Image(imgDouble, filename, type, linear, colorspace, scalingFactor) # long name
        return Image(imgDouble, name, type, linear, colorspace, scalingFactor)      # short name

    def read(filename,exif=True): return Image.readImage(filename,readExif=exif)

    def newImage(shape, colorSpaceName=None, color=None, type=None,name=None):
        """ description """

        # default values
        if not colorSpaceName : colorSpaceName = 'sRGB'
        if not color : color = np.asarray([0.0,0.0,0.0]) 
        if not type : type : imageType.imageType.SDR
        linear = False if ((type==imageType.imageType.SDR) and (colorSpaceName=='sRGB')) else True
        if not name : name = "no name["+colorSpaceName+"]" 
        scalingFactor = 1.0

        # colorSpace
        colorSpace = ColorSpace.ColorSpace.build(colorSpaceName) 

        # colorData
        if len(shape) == 2 :
            (h,w) = shape
            newShape = (h,w,3)
        else:
            newShape = shape

        colorData = np.ones(newShape)
        colorData[:,:,0] = colorData[:,:,0]*color[0]
        colorData[:,:,1] = colorData[:,:,1]*color[1]
        colorData[:,:,2] = colorData[:,:,2]*color[2]

        return Image(colorData, name, type, linear, colorSpace, scalingFactor)

    def forceColorData3(imgDouble):
        """ force color data to hace 3 channels """

        # force image to 3 channels
        if len(imgDouble.shape) == 2:
            # single channel image
            h,w = imgDouble.shape
            img3 = np.ones([h,w,3])
            img3[:,:,0] = imgDouble
            img3[:,:,1] = imgDouble
            img3[:,:,2] = imgDouble
        else:
            h,w,c = imgDouble.shape
            if c==4:
                # remove alpha channel
                img3 = np.ones([h,w,3])
                img3[:,:,0] = imgDouble[:,:,0]
                img3[:,:,1] = imgDouble[:,:,1]
                img3[:,:,2] = imgDouble[:,:,2]
            else:
                img3 = imgDouble
        return img3 

    def array2vector(img):
        """ transform 2D array of color data to vector """

        if len(img.shape) ==2 :
            x,y = img.shape
            c = 1
        else:
            x,y,c = img.shape
        return np.reshape(img, (x * y, c))

    def scaleMaxOne(img):
        """ scale image colorData in [0, 1] space """

        imgVector = Image.array2vector(img)
        R, G, B = imgVector[:,0], imgVector[:,1], imgVector[:,2]
        maxRGB = max([np.amax(R), np.amax(G), np.amax(B)]) 

        return img/maxRGB, 1.0/maxRGB