dataset.py

# -*- coding: utf-8 -*-
"""
GEPARD - Gepard-Enabled PARticle Detection
Copyright (C) 2018  Lars Bittrich and Josef Brandt, Leibniz-Institut für 
Polymerforschung Dresden e. V. <bittrich-lars@ipfdd.de>    

This program 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 3 of the License, or
(at your option) any later version.

This program 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 this program, see COPYING.  
If not, see <https://www.gnu.org/licenses/>.
"""
import os
import pickle
import numpy as np
import sys
import cv2
from copy import copy
from .analysis.particleContainer import ParticleContainer
from .legacyConvert import legacyConversion, currentVersion
from .helperfunctions import cv2imwrite_fix, cv2imread_fix
# for legacy pickle import the old module name dataset must be found 
# (no relative import)
from . import dataset
from . import analysis
sys.modules['dataset'] = dataset
sys.modules['analysis'] = analysis


def loadData(fname):
    retds = None
    with open(fname, "rb") as fp:
        ds = pickle.load(fp)
        ds.fname = fname
        ds.readin = True
        ds.updatePath()
        retds = DataSet(fname)
        retds.version = 0
        retds.__dict__.update(ds.__dict__)
        if retds.version < currentVersion:
            legacyConversion(retds)
        elif retds.zvalimg == "saved":
            retds.loadZvalImg()
        retds.particleContainer.datasetParent = retds #TODO: This is mainly a workaround to update the ref in particleContainer. It probably should be handled differently anyways...
    return retds

def saveData(dataset, fname):
    with open(fname, "wb") as fp:
        # zvalimg is rather large and thus it is saved separately in a tif file 
        # only onces after its creation
        zvalimg = dataset.zvalimg
        if zvalimg is not None:
            dataset.zvalimg = "saved"
        pickle.dump(dataset, fp, protocol=-1)
        dataset.zvalimg = zvalimg

def arrayCompare(a1, a2):
    if a1.shape!=a2.shape:
        return False
    if a1.dtype!=np.float32 and a1.dtype!=np.float64:
        return np.all(a1==a2)
    ind = np.isnan(a1)
    if not np.any(ind):
        return np.all(a1==a2)
    return np.all(a1[~ind]==a2[~ind])

def listCompare(l1, l2):
    if len(l1)!=len(l2):
        return False
    for l1i, l2i in zip(l1, l2):
        if isinstance(l1i, np.ndarray):
            if not isinstance(l2i, np.ndarray) or not arrayCompare(l1i, l2i):
                return False
        elif isinstance(l1i, (list, tuple)):
            if not isinstance(l2i, (list, tuple)) or not listCompare(l1i, l2i):
                return False
        elif l1i!=l2i and ((~np.isnan(l1i)) or (~np.isnan(l2i))):
            return False
    return True

def recursiveDictCompare(d1, d2):
    for key in d1:
        if not key in d2:
            print("key missing in d2:", key, flush=True)
            return False
        a = d1[key]
        b = d2[key]
        print(key, type(a), type(b), flush=True)
        if isinstance(a, np.ndarray):
            if not isinstance(b, np.ndarray) or not arrayCompare(a, b):
                print("data is different!", a, b)
                return False
        elif isinstance(a, dict):
            if not isinstance(b, dict):
                print("data is different!", a, b)
                return False
            if not recursiveDictCompare(a, b):
                return False
        elif isinstance(a, (list, tuple)):
            if not isinstance(b, (list, tuple)) or not listCompare(a, b):
                print("data is different!", a, b)
                return False
        elif a != b:
            if (a is not None) and (b is not None):
                print("data is different!", a, b)
                return False
    return True


class DataSet(object):
    def __init__(self, fname, newProject=False):
        self.fname = fname
        # parameters specifically for optical scan
        self.version = currentVersion
        self.lastpos = None
        self.maxdim = None
        self.pixelscale_df = None  # µm / pixel --> scale of DARK FIELD camera (used for image stitching)
        self.pixelscale_bf = None  # µm / pixel of DARK FIELD camera (set to same as bright field, if both use the same camera)
        self.imagedim_bf = None  # width, height, angle of BRIGHT FIELD camera
        self.imagedim_df = None  # width, height, angle of DARK FIELD camera (set to same as bright field, if both use the same camera)
        self.imagescanMode = 'df'  # was the fullimage acquired in dark- or brightfield?
        self.fitpoints = []   # manually adjusted positions aquired to define the specimen geometry
        self.fitindices = []  # which of the five positions in the ui are already known
        self.boundary = []    # scan boundary computed by a circle around the fitpoints + manual adjustments
        self.grid = []        # scan grid positions for optical scan
        self.zpositions = []  # z-positions for optical scan
        self.heightmap = None
        self.zvalimg = None
        self.coordinatetransform = None  # if imported form extern source coordinate system may be rotated
        self.signx = 1.
        self.signy = -1.

        # tiling parameters
        self.pyramidParams = None
        
        # parameters specifically for raman scan
        self.pshift = None    # shift of spectrometer scan position relative to image center
        self.seedpoints = np.array([])
        self.seeddeletepoints = np.array([])
        self.detectParams = {'points': np.array([[50,0],[100,200],[200,255]]),
                             'contrastcurve': True,
                             'blurRadius': 9,
                             'threshold': 0.2,
                             'maxholebrightness': 0.5,
                             'minparticlearea': 20,
                             'minparticledistance': 20,
                             'measurefrac': 1,
                             'compactness': 0.0,
                             'seedRad': 3}

        self.particleContainer = ParticleContainer(self)
        self.particleDetectionDone = False
        self.specscandone = False

        self.resultParams = {'minHQI': 5}
        self.colorSeed = 'default'
        self.resultsUploadedToSQL = []
        
        self.readin = True    # a value that is always set to True at loadData 
                              # and mark that the coordinate system might be changed in the meantime
        self.mode = "prepare"
        if newProject:
            self.fname = self.newProject(fname)
        self.updatePath()

    @property
    def opticalScanDone(self) -> bool:
        return os.path.exists(self.getZvalImageName())

    def __eq__(self, other):
        return recursiveDictCompare(self.__dict__, other.__dict__)

    def getPyramidParams(self):
        return self.pyramidParams

    def setPyramidParams(self, pyramid_params):
        self.pyramidParams = pyramid_params
        
    def getPixelScale(self, mode=None):
        if mode is None:
            mode = self.imagescanMode
        return (self.pixelscale_df if mode == "df" else self.pixelscale_bf)
    
    def saveZvalImg(self):
        if self.zvalimg is not None:
            cv2imwrite_fix(self.getZvalImageName(), self.zvalimg)
            self.zvalimg = "saved"

    def loadZvalImg(self):
        if os.path.exists(self.getZvalImageName()):
            self.zvalimg = cv2imread_fix(self.getZvalImageName(), cv2.IMREAD_GRAYSCALE)
    
    def getZval(self, pixelpos):
        assert self.zvalimg is not None
        if self.zvalimg == "saved":
            self.loadZvalImg()
        i, j = int(round(pixelpos[1])), int(round(pixelpos[0]))
        if i>=self.zvalimg.shape[0]: 
            print('error in getZval:', self.zvalimg.shape, i, j)
            i = self.zvalimg.shape[0]-1
        if j>=self.zvalimg.shape[1]: 
            print('error in getZval:', self.zvalimg.shape, i, j)
            j = self.zvalimg.shape[1]-1
        zp = self.zvalimg[i,j]
        z0, z1 = self.zpositions[0], self.zpositions[-1]
        return zp/255.*(z1-z0) + z0
    
    def mapHeight(self, x, y):
        assert not self.readin
        assert self.heightmap is not None
        return self.heightmap[0]*x + self.heightmap[1]*y + self.heightmap[2]
        
    def mapToPixel(self, p, mode='df', force=False):
        if not force:
            assert not self.readin
        p0 = copy(self.lastpos)
        
        if self.coordinatetransform is not None:
            z = 0. if len(p) < 3 else p[2]
            T, pc = self.coordinatetransform
            p = (np.dot(np.array([p[0], p[1], z])-pc, T.T))

        assert mode in ['df', 'bf'], f'mapToPixel mode: {mode} not understood'
        pixelscale: float = self.pixelscale_bf if mode == 'bf' else self.pixelscale_df
        p0[0] -= self.signx*self.imagedim_df[0]/2
        p0[1] -= self.signy*self.imagedim_df[1]/2
        x, y = self.signx*(p[0] - p0[0])/pixelscale, self.signy*(p[1] - p0[1])/pixelscale
        return x, y
            
    def mapToLength(self, pixelpos, mode='df', force=False, returnz=False):
        if not force:
            assert not self.readin
        p0 = copy(self.lastpos)
        if mode == 'df':
            p0[0] -= self.signx*self.imagedim_df[0]/2
            p0[1] -= self.signy*self.imagedim_df[1]/2
            x, y = (self.signx*pixelpos[0]*self.pixelscale_df + p0[0]), (p0[1] + self.signy*pixelpos[1]*self.pixelscale_df)
        elif mode == 'bf':
            p0[0] -= self.signx*self.imagedim_bf[0]/2
            p0[1] -= self.signy*self.imagedim_bf[1]/2
            x, y = (self.signx*pixelpos[0]*self.pixelscale_bf + p0[0]), (p0[1] + self.signy*pixelpos[1]*self.pixelscale_bf)
        else:
            raise ValueError(f'mapToLength mode: {mode} not understood')
        
        z = None
        if (returnz and self.zvalimg is not None) or self.coordinatetransform is not None:
            z = self.mapHeight(x, y)
            z += self.getZval(pixelpos)
            
        if self.coordinatetransform is not None:
            T, pc = self.coordinatetransform
            x, y, z = (np.dot(np.array([x, y, z]), T) + pc)
        
        if returnz:
            return x, y, z
        return x, y
    
    def mapToLengthSpectrometer(self, pixelpos, microscopeMode='df', noz=False):
        p0x, p0y, z = self.mapToLength(pixelpos, mode=microscopeMode, returnz=True)
        x, y = p0x + self.pshift[0], p0y + self.pshift[1]
        return x, y, z
        
    def newProject(self, fname):
        path = os.path.split(fname)[0]
        name = os.path.splitext(os.path.basename(fname))[0]
        newpath = os.path.join(path, name)
        fname = os.path.join(newpath, name + ".pkl")
        if not os.path.exists(newpath):
            os.mkdir(newpath)        # for new projects a directory will be created
        elif os.path.exists(fname):  # if this project is already there, load it instead
            self.__dict__.update(loadData(fname).__dict__)
        return fname
    
    def getTilePath(self):
        scandir = os.path.join(self.path, "tiles")
        if not os.path.exists(scandir):
            os.mkdir(scandir)
        return scandir

    def getScanPath(self):
        scandir = os.path.join(self.path, "scanimages")
        if not os.path.exists(scandir):
            os.mkdir(scandir)
        return scandir
        
    def updatePath(self):
        self.path = os.path.split(self.fname)[0]
        self.name = os.path.splitext(os.path.basename(self.fname))[0]
    
    def getSpectraFileName(self):
        return os.path.join(self.path, 'spectra.npy')
    
    def getImageName(self):
        return os.path.join(self.path, 'fullimage.tif')

    def getZvalImageName(self):
        return os.path.join(self.path, "zvalues.tif")
    
    def getLegacyImageName(self):
        return os.path.join(self.path, "fullimage.png")
    
    def getBackgroundImageName(self):
        return os.path.join(self.path, "background.bmp")
    
    def getTmpImageName(self):
        return os.path.join(self.path, "tmp.bmp")
            
    def save(self):
        saveData(self, self.fname)
    
    def saveBackup(self):
        inc = 0
        while True:
            directory = os.path.dirname(self.fname)
            filename = self.name + '_backup_' + str(inc) + '.pkl'
            path = os.path.join(directory, filename)
            if os.path.exists(path):
                inc += 1
            else:
                saveData(self, path)
                return filename