pymca 4.5.0-4 / usr / share / pyshared / PyMca / TiffStack.py

#/*##########################################################################
# Copyright (C) 2004-2012 European Synchrotron Radiation Facility
#
# This file is part of the PyMCA X-ray Fluorescence Toolkit developed at
# the ESRF by the Beamline Instrumentation Software Support (BLISS) group.
#
# This toolkit is free software; you can redistribute it and/or modify it 
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option) 
# any later version.
#
# PyMCA is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# PyMCA; if not, write to the Free Software Foundation, Inc., 59 Temple Place,
# Suite 330, Boston, MA 02111-1307, USA.
#
# PyMCA follows the dual licensing model of Trolltech's Qt and Riverbank's PyQt
# and cannot be used as a free plugin for a non-free program. 
#
# Please contact the ESRF industrial unit (industry@esrf.fr) if this license 
# is a problem for you.
#############################################################################*/

from PyMca import DataObject
from PyMca import TiffIO
import numpy
import os
import sys
if sys.version > '2.9':
    long = int

SOURCE_TYPE = "TiffStack"

class TiffArray(object):
    def __init__(self, filelist, shape, dtype, imagestack=True):
        self.__fileList    = filelist
        self.__shape       = shape
        self.__dtype       = dtype
        self.__imageStack  = imagestack
        if imagestack:
            self.__nImagesPerFile = int(shape[0]/len(filelist))
        else:
            self.__nImagesPerFile = int(shape[-1]/len(filelist))
        self.__oldFileNumber = -1
        
    def __getitem__(self, args0):
        standardSlice = True
        indices = []
        outputShape = []
        scalarArgs = []
        args = []
        for i in range(len(self.__shape)):
            if i < len(args0):
                args.append(args0[i])
            else:
                args.append(slice(None, None, None))
        for i in range(len(args)):
            if isinstance(args[i], slice):
                start = args[i].start
                stop  = args[i].stop
                step  = args[i].step
                if start is None:
                    start = 0
                if stop is None:
                    stop = self.__shape[i]
                if step is None:
                    step = 1
                if step < 1:
                    raise ValueError("Step must be >= 1 (got %d)" % step)
                if start is None:
                    start = 0
                if start < 0:
                    start = self.__shape[i]-start
                if stop < 0:
                    stop = self.__shape[i]-stop
                if stop == start:
                    raise ValueError("Zero-length selections are not allowed")
                indices.append(list(range(start, stop, step)))
            elif type(args[i]) == type([]):
                if len(args[i]):
                    indices.append([int(x) for x in args[i]])
                else:
                    stardardSlice = False
            elif type(args[i]) in [1, long(1)]:
                start = args[i]
                if start < 0:
                    start = self.__shape[i] - start
                stop = start + 1
                step = 1
                start = args[i]
                args[i] = slice(start, stop, step)
                indices.append(list(range(start, stop, step)))
                scalarArgs.append(i)
            else:
                standardSlice = False

        if not standardSlice:
            print("args = ", args)
            raise NotImplemented("__getitem__(self, args) only works on slices")

        if len(indices) < 3:
            print("input args = ", args0)
            print("working args = ", args)
            print("indices = ", indices)
            raise NotImplemented("__getitem__(self, args) only works on slices")
        outputShape = [len(indices[0]), len(indices[1]), len(indices[2])]
        outputArray = numpy.zeros(outputShape, dtype=self.__dtype)
        nbFiles = len(self.__fileList)
        nImagesPerFile = self.__nImagesPerFile
        
        if self.__imageStack:
            i = 0
            rowMin = min(indices[1])
            rowMax = max(indices[1])
            for imageIndex in indices[0]:
                fileNumber = int(imageIndex/nImagesPerFile)
                if fileNumber != self.__oldFileNumber:
                    self.__tmpInstance = TiffIO.TiffIO(self.__fileList[fileNumber],
                                                       mode='rb+')
                    self.__oldFileNumber = fileNumber
                imageNumber = imageIndex % nImagesPerFile
                imageData = self.__tmpInstance.getData(imageNumber,
                                                       rowMin=rowMin,
                                                       rowMax=rowMax)
                try:
                    outputArray[i,:,:] = imageData[args[1],args[2]]
                except:
                    print("outputArray[i,:,:].shape =",outputArray[i,:,:].shape)
                    print("imageData[args[1],args[2]].shape = " , imageData[args[1],args[2]].shape)
                    print("input args = ", args0)
                    print("working args = ", args)
                    print("indices = ", indices)
                    print("scalarArgs = ", scalarArgs)
                    raise
                i += 1
        else:
            i = 0
            rowMin = min(indices[0])
            rowMax = max(indices[0])
            for imageIndex in indices[-1]:
                fileNumber = int(imageIndex/nImagesPerFile)
                if fileNumber != self.__oldFileNumber:
                    self.__tmpInstance = TiffIO.TiffIO(self.__fileList[fileNumber],
                                                       mode='rb+')
                    self.__oldFileNumber = fileNumber
                imageNumber = imageIndex % nImagesPerFile
                imageData = self.__tmpInstance.getData(imageNumber,
                                                       rowMin=rowMin,
                                                       rowMax=rowMax)
                outputArray[:,:, i] = imageData[args[0],args[1]]
                i += 1
        if len(scalarArgs):
            finalShape = []
            for i in range(len(outputShape)):
                if i in scalarArgs:
                    continue
                finalShape.append(outputShape[i])
            outputArray.shape = finalShape
        return outputArray

    def getShape(self):
        return self.__shape
    shape = property(getShape)

    def getDtype(self):
        return self.__dtype
    dtype = property(getDtype)

class TiffStack(DataObject.DataObject):
    def __init__(self, filelist=None, imagestack=None, dtype=None):
        DataObject.DataObject.__init__(self)
        self.sourceType = SOURCE_TYPE
        if imagestack is None:
            self.__imageStack = True
        else:
            self.__imageStack = imagestack
        self.__dtype = dtype
        if filelist is not None:
            if type(filelist) != type([]):
                filelist = [filelist]
            if len(filelist) == 1:
                self.loadIndexedStack(filelist)
            else:
                self.loadFileList(filelist)

    def loadFileList(self, filelist, dynamic=False, fileindex=0):
        if type(filelist) != type([]):
            filelist = [filelist]

        #retain the file list
        self.sourceName = filelist
        
        #the number of files
        nbFiles=len(filelist)

        #the intance to access the first file
        fileInstance = TiffIO.TiffIO(filelist[0])

        #the number of images per file
        nImagesPerFile = fileInstance.getNumberOfImages()

        #get the dimensions from the image itself
        tmpImage = fileInstance.getImage(0)
        if self.__dtype is None:
            self.__dtype = tmpImage.dtype

        nRows, nCols = tmpImage.shape

        #stack shape
        if self.__imageStack:
            shape = (nbFiles * nImagesPerFile, nRows, nCols)
        else:
            shape = (nRows, nCols, nbFiles * nImagesPerFile)

        #we can create the stack
        if not dynamic:
            try:
                data = numpy.zeros(shape,
                                   self.__dtype)
            except (MemoryError, ValueError):
                dynamic = True
        if not dynamic:
            imageIndex = 0
            self.onBegin(nbFiles * nImagesPerFile)
            for i in range(nbFiles):
                tmpInstance =TiffIO.TiffIO(filelist[i])
                for j in range(nImagesPerFile):
                    tmpImage = tmpInstance.getImage(j)
                    if self.__imageStack:
                        data[imageIndex,:,:] = tmpImage
                    else:
                        data[:,:,imageIndex] = tmpImage
                    imageIndex += 1
                    self.incrProgressBar = imageIndex
                    self.onProgress(imageIndex)
            self.onEnd()

        if dynamic:
            data = TiffArray(filelist,
                             shape,
                             self.__dtype,
                             imagestack=self.__imageStack)
        self.info = {}
        self.data = data
        shape = self.data.shape
        for i in range(len(shape)):
            key = 'Dim_%d' % (i+1,)
            self.info[key] = shape[i]

        if self.__imageStack:
            self.info["McaIndex"] = 0
            self.info["FileIndex"] = 1
        else:
            self.info["McaIndex"] = 2
            self.info["FileIndex"] = 0

        self.info["SourceType"] = SOURCE_TYPE
        self.info["SourceName"] = self.sourceName

    def loadIndexedStack(self,filename,begin=None,end=None, skip = None, fileindex=0):
        #if begin is None: begin = 0
        if type(filename) == type([]):
            filename = filename[0]
        if not os.path.exists(filename):
            raise IOError("File %s does not exists" % filename)
        name = os.path.basename(filename)
        n = len(name)
        i = 1
        numbers = ['0', '1', '2', '3', '4', '5',
                   '6', '7', '8','9']
        while (i <= n):
            c = name[n-i:n-i+1]
            if c in ['0', '1', '2',
                                '3', '4', '5',
                                '6', '7', '8',
                                '9']:
                break
            i += 1
        suffix = name[n-i+1:]
        if len(name) == len(suffix):
            #just one file, one should use standard widget
            #and not this one.
            self.loadFileList(filename, fileindex=fileindex)
        else:
            nchain = []
            while (i<=n):
                c = name[n-i:n-i+1]
                if c not in ['0', '1', '2',
                                    '3', '4', '5',
                                    '6', '7', '8',
                                    '9']:
                    break
                else:
                    nchain.append(c)
                i += 1
            number = ""
            nchain.reverse()
            for c in nchain:
                number += c
            fformat = "%" + "0%dd" % len(number)
            if (len(number) + len(suffix)) == len(name):
                prefix = ""
            else:
                prefix = name[0:n-i+1]
            prefix = os.path.join(os.path.dirname(filename),prefix)
            if not os.path.exists(prefix + number + suffix):
                print("Internal error in TIFFStack")
                print("file should exist: %s " % (prefix + number + suffix))
                return
            i = 0
            if begin is None:
                begin = 0
                testname = prefix+fformat % begin+suffix
                while not os.path.exists(prefix+fformat % begin+suffix):
                    begin += 1
                    testname = prefix+fformat % begin+suffix
                    if len(testname) > len(filename):break
                i = begin
            else:
                i = begin
            if not os.path.exists(prefix+fformat % i+suffix):
                raise ValueError("Invalid start index file = %s" % \
                      (prefix+fformat % i+suffix))
            f = prefix+fformat % i+suffix
            filelist = []
            while os.path.exists(f):
                filelist.append(f)
                i += 1
                if end is not None:
                    if i > end:
                        break
                f = prefix+fformat % i+suffix
            self.loadFileList(filelist, fileindex=fileindex)

    def onBegin(self, n):
        pass

    def onProgress(self, n):
        pass

    def onEnd(self):
        pass
            
def test():
    from PyMca import StackBase
    testFileName = "TiffTest.tif"
    nrows = 2000
    ncols = 2000
    #create a dummy stack with 100 images
    nImages = 100
    imagestack = True
    a = numpy.ones((nrows, ncols), numpy.float32)
    if not os.path.exists(testFileName):
        print("Creating test filename %s"  % testFileName)
        tif = TiffIO.TiffIO(testFileName, mode = 'wb+')

        for i in range(nImages):
            data = (a * i).astype(numpy.float32)
            if i == 1:
                tif = TiffIO.TiffIO(testFileName, mode = 'rb+')
            tif.writeImage(data,
                           info={'Title':'Image %d of %d' % (i+1, nImages)})
        tif = None

    stackData = TiffStack(imagestack=imagestack)
    stackData.loadFileList([testFileName], dynamic=True)

    if 0:
        stack = StackBase.StackBase()
        stack.setStack(stackData)
        print("This should be 0 = %f" %  stack.calculateROIImages(0, 0)['ROI'].sum())
        print("This should be %f = %f" %\
              (a.sum(),stack.calculateROIImages(1, 2)['ROI'].sum()))
        if imagestack:
            print("%f should be = %f" %\
                  (stackData.data[0:10,:,:].sum(),
                   stack.calculateROIImages(0, 10)['ROI'].sum()))
            print("Test small ROI 10 should be = %f" %\
                  stackData.data[10:11,[10],11].sum())
            print("Test small ROI 40 should be = %f" %\
                  stackData.data[10:11,[10,12,14,16],11].sum())
        else:
            print("%f should be = %f" %\
                  (stackData.data[:,:, 0:10].sum(),
                   stack.calculateROIImages(0, 10)['ROI'].sum()))
            print("Test small ROI %f" %\
                  stackData.data[10:11,[29],:].sum())
    else:
        from PyMca import PyMcaQt as qt
        from PyMca import QStackWidget
        app = qt.QApplication([])
        w = QStackWidget.QStackWidget()
        print("Setting stack")
        w.setStack(stackData)
        w.show()
        app.exec_()

if __name__ == "__main__":
    test()