particleeditor.py

# -*- coding: utf-8 -*-
"""
Created on Wed Jan 16 12:43:00 2019

@author: brandt
"""

"""
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 numpy as np
import cv2
from PyQt5 import QtWidgets, QtCore
from copy import deepcopy

from .particlePainter import ParticlePainter

def getContourStatsWithPixelScale(contours, pixelscale):
        ##characterize particle
        longellipse, shortellipse = np.nan, np.nan
        
        cnt = contours
        
        if cnt.shape[0] >= 5:       ##at least 5 points required for ellipse fitting...
            ellipse = cv2.fitEllipse(cnt)
            shortellipse, longellipse = ellipse[1]

        rect = cv2.minAreaRect(cnt)
        long, short = rect[1]
        if short>long:
            long, short = short, long
        
        area = cv2.contourArea(cnt)
    
        return long*pixelscale, short*pixelscale, longellipse*pixelscale, shortellipse*pixelscale, area*pixelscale**2


class ParticleContextMenu(QtWidgets.QMenu):
    combineParticlesSignal = QtCore.pyqtSignal(list, str)
    reassignParticlesSignal = QtCore.pyqtSignal(list, str)
    paintParticlesSignal = QtCore.pyqtSignal(list, str)
    deleteParticlesSignal = QtCore.pyqtSignal(list)
    
    def __init__(self, viewparent):
        super(ParticleContextMenu, self).__init__()
        self.viewparent = viewparent
        self.selectedParticleIndices = self.viewparent.selectedParticleIndices
        self.particleContainer = self.viewparent.dataset.particleContainer
        
    def executeAtScreenPos(self, screenPos):
        self.combineActs = []
        self.combineMenu = QtWidgets.QMenu("Combine Particles into")

        self.paintActs = []
        self.paintMenu = QtWidgets.QMenu("Paint Mode, merge into")

        selctedAssignments = []

        for particleIndex in self.selectedParticleIndices:
            try:
                assignment = self.particleContainer.getParticleAssignmentByIndex(particleIndex)
            except:
                return
            selctedAssignments.append(assignment)
        
        for assignment in np.unique(selctedAssignments):
            self.combineActs.append(self.combineMenu.addAction(assignment))
            self.paintActs.append(self.paintMenu.addAction(assignment))
        self.combineActs.append(self.combineMenu.addAction("other"))
        self.paintActs.append(self.paintMenu.addAction("other"))
        
        
        self.reassignActs = []
        self.reassignMenu = QtWidgets.QMenu("Reassign particle(s) into")
        for polymType in self.particleContainer.getUniquePolymers():
            self.reassignActs.append(self.reassignMenu.addAction(polymType))
        self.reassignActs.append(self.reassignMenu.addAction("other"))
        
        numParticles = len(self.selectedParticleIndices)
        if numParticles == 0:
            self.reassignMenu.setDisabled(True)
            self.combineMenu.setDisabled(True)
        elif numParticles == 1:
            self.combineMenu.setDisabled(True)
            
        self.addMenu(self.combineMenu)        
        self.addMenu(self.reassignMenu)
        self.addMenu(self.paintMenu)
        self.deleteAct = self.addAction("Delete particle(s)")
        
        action = self.exec_(screenPos)

        if action:
            newAssignment = self.validifyAssignment(action.text())

            if action in self.combineActs:
                self.combineParticlesSignal.emit(self.selectedParticleIndices, newAssignment)
            elif action in self.reassignActs:
                self.reassignParticlesSignal.emit(self.selectedParticleIndices, newAssignment)
            elif action in self.paintActs:
                self.paintParticlesSignal.emit(self.selectedParticleIndices, newAssignment)
            elif action == self.deleteAct:
                self.deleteParticlesSignal.emit(self.selectedParticleIndices)
    
    def validifyAssignment(self, assignment):
        if assignment == "other":
            assignment = self.getNewEntry()
        return assignment
    
    def getNewEntry(self):
        text, okClicked = QtWidgets.QInputDialog.getText(self.viewparent, "Custom assignment", "Enter new assignment")
        if okClicked and text != '':
            return text


class ParticleEditor(QtCore.QObject):
    particleAssignmentChanged = QtCore.pyqtSignal()
    def __init__(self, viewparent, particleContainer):
        super(ParticleEditor, self).__init__()
        self.particleContainer = particleContainer
        self.viewparent = viewparent #the assigned analysis widget
        self.backupFreq = 3     #save a backup every n actions
        self.neverBackedUp = True
        self.actionCounter = 0
        
        self.storedIndices = []
        self.storedAssignmend = None
    
    def connectToSignals(self, contextMenu):
        contextMenu.combineParticlesSignal.connect(self.combineParticles)
        contextMenu.reassignParticlesSignal.connect(self.reassignParticles)
        contextMenu.paintParticlesSignal.connect(self.paintParticles)
        contextMenu.deleteParticlesSignal.connect(self.deleteParticles)
    
    def createSafetyBackup(self):
        self.actionCounter += 1
        if self.actionCounter == self.backupFreq-1 or self.neverBackedUp:            
            backupname = self.viewparent.dataset.saveBackup()
            print('backing up as', backupname)
            self.neverBackedUp = False
            self.actionCounter = 0        
    
    @QtCore.pyqtSlot(list, str)
    def combineParticles(self, contourIndices, newAssignment):     
        self.createSafetyBackup()
        print(f'Combining particles {contourIndices} into {newAssignment}')
        contours = self.particleContainer.getParticleContoursByIndex(contourIndices)
        try:
            newContour = self.mergeContours(contours) 
        except NotConnectedContoursError:
            return
        
        pixelscale = self.viewparent.dataset.getPixelScale()
        stats = getContourStatsWithPixelScale(newContour, pixelscale)

        self.mergeParticlesInParticleContainerAndSampleView(contourIndices,newContour, stats, newAssignment)
    
    @QtCore.pyqtSlot(list, str)
    def reassignParticles(self, contourindices, newAssignment):
        self.createSafetyBackup()
        print(f'reassigning indices {contourindices} into {newAssignment}')
        for partIndex in contourindices:
            self.particleContainer.reassignParticleToAssignment(partIndex, newAssignment)

        self.particleAssignmentChanged.emit()
     
    @QtCore.pyqtSlot(list, str)
    def paintParticles(self, contourIndices, newAssignment):
        print(f'painting indices {contourIndices} into {newAssignment}')
        self.createSafetyBackup()
        self.storedIndices = contourIndices
        self.storedAssignmend = newAssignment
        
        contours = self.particleContainer.getParticleContoursByIndex(contourIndices)
#        topLeft = self.getTopLeft(contours)
#        img, self.xmin, self.ymin, self.padding = self.contoursToImg(contours, padding=0)
        img, xmin, ymin, self.padding = self.contoursToImg(contours, padding=0)
        topLeft = [ymin, xmin]
        self.particlePainter = ParticlePainter(self, img, topLeft)
        
        self.viewparent.normalSize()
        self.viewparent.particlePainter = self.particlePainter
        self.viewparent.scene().addItem(self.particlePainter)
        self.viewparent.update()
        
    def acceptPaintedResult(self):
        try:
            img = self.particlePainter.img
            xmin = self.particlePainter.topLeft[1]
            ymin = self.particlePainter.topLeft[0]
            newContour = self.imgToCnt(img, xmin, ymin, 0)
        except NotConnectedContoursError:
            self.storedIndices = []
            self.storedAssignmend = None
            self.destroyParticlePainter()
            return
        
        pixelscale = self.viewparent.dataset.getPixelScale()
        stats = getContourStatsWithPixelScale(newContour, pixelscale)
        
        self.mergeParticlesInParticleContainerAndSampleView(self.storedIndices, newContour, stats, self.storedAssignmend)
        
        self.storedIndices = []
        self.storedAssignmend = None
        self.destroyParticlePainter()
    
    def destroyParticlePainter(self):
        if self.particlePainter is not None:
            self.viewparent.particlePainter = None
            self.viewparent.scene().removeItem(self.particlePainter)
            self.viewparent.update()
            self.particlePainter = None
    
    def mergeContours(self, contours):
        img, xmin, ymin, padding = self.contoursToImg(contours)
        return self.imgToCnt(img, xmin, ymin, padding)
    
    def contoursToImg(self, contours, padding=2):
        contourCopy = deepcopy(contours)
        cnt = np.vstack(tuple(contourCopy))  #combine contous
        #draw contours
        xmin, xmax = cnt[:,0,:][:, 0].min(), cnt[:,0,:][:, 0].max()
        ymin, ymax = cnt[:,0,:][:, 1].min(), cnt[:,0,:][:, 1].max()        
        padding = padding   #pixel in each direction
        rangex = int(np.round((xmax-xmin)+2*padding))
        rangey = int(np.round((ymax-ymin)+2*padding))

        img = np.zeros((rangey, rangex))
        for curCnt in contourCopy:
            for i in range(len(curCnt)):
                curCnt[i][0][0] -= xmin-padding
                curCnt[i][0][1] -= ymin-padding
                
            cv2.drawContours(img, [curCnt], -1, 255, -1)
            cv2.drawContours(img, [curCnt], -1, 255, 1)
        
        img = np.uint8(cv2.morphologyEx(img, cv2.MORPH_CLOSE, np.ones((3, 3))))
        return img, xmin, ymin, padding
    
    def imgToCnt(self, img, xmin, ymin, padding):
        if cv2.__version__ > '3.5':
            contours, hierarchy = cv2.findContours(img, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
        else:
            temp, contours, hierarchy = cv2.findContours(img, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)

        if len(contours)>1:
            QtWidgets.QMessageBox.critical(self.viewparent, 'ERROR!', 
                                           'Particle contours are not connected or have holes.\nThat is currently not supported!')
            raise NotConnectedContoursError
        
        newContour = contours[0]
        for i in range(len(newContour)):
            newContour[i][0][0] += xmin-padding
            newContour[i][0][1] += ymin-padding
            
        return newContour
    
    def getTopLeft(self, contours):
        cnt = np.vstack(tuple(contours))  #combine contous
        #draw contours
        xmin = cnt[:,0,:][:, 0].min()
        ymin= cnt[:,0,:][:, 1].min()
        return [ymin, xmin]
    
    def mergeParticlesInParticleContainerAndSampleView(self, indices, newContour, stats, assignment):
        self.viewparent.addParticleContourToIndex(newContour, len(self.viewparent.contourItems)-1)
        self.particleContainer.addMergedParticle(indices, newContour, stats, newAssignment=assignment)
        
        for ind in sorted(indices, reverse=True):
            self.viewparent.removeParticleContour(ind)
            self.particleContainer.removeParticle(ind)
        
        self.viewparent.resetContourIndices()
        self.particleContainer.resetParticleIndices()
        self.particleAssignmentChanged.emit()
        #TODO: INCLUDE SANITY CHECK!!!!!!!!!
    
    @QtCore.pyqtSlot(list)
    def deleteParticles(self, contourIndices):
        reply = QtWidgets.QMessageBox.question(self, f'About to delete {len(contourIndices)} particles.',
                                "Are you sure to permanantly delete these particles?",
                                QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No, QtWidgets.QMessageBox.No)
        if reply == QtWidgets.QMessageBox.Yes:

            for partIndex in sorted(contourIndices, reverse=True):
                self.viewparent.removeParticleContour(partIndex)
                self.viewparent.dataset.particleContainer.removeParticles([partIndex])
                
            self.viewparent.dataset.particleContainer.resetParticleIndices()
            self.viewparent.resetContourIndices()
            self.viewparent.analysiswidget.updateHistogramsAndContours()


class NotConnectedContoursError(Exception):
    pass