particlePainter.py 6.19 KB
Newer Older
JosefBrandt's avatar
JosefBrandt committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
#!/usr/bin/env python3
# -*- 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/>.
"""
from PyQt5 import QtWidgets, QtCore, QtGui
import numpy as np
JosefBrandt's avatar
 
JosefBrandt committed
24 25
import cv2

JosefBrandt's avatar
JosefBrandt committed
26 27

class ParticlePainter(QtWidgets.QGraphicsItem):
JosefBrandt's avatar
 
JosefBrandt committed
28 29 30 31 32 33 34
    def __init__(self, editorParent, contours, pos=(500,500)):
        super(ParticlePainter, self).__init__()
        self.editorParent = editorParent
        self.viewparent = self.editorParent.viewparent
        self.setZValue(5)
        self.polygons = None
        self.contours = contours
JosefBrandt's avatar
JosefBrandt committed
35
        self.mousePos = None
JosefBrandt's avatar
 
JosefBrandt committed
36
        self.minRadius = 10
JosefBrandt's avatar
JosefBrandt committed
37
        self.maxRadius = 500
JosefBrandt's avatar
 
JosefBrandt committed
38
        self.radius = 50
JosefBrandt's avatar
JosefBrandt committed
39
        self.brect = QtCore.QRectF(0,0,1,1)
JosefBrandt's avatar
 
JosefBrandt committed
40 41 42
        self.getBrectAndPolygon()
        self.painting = False
        self.erasing = False
JosefBrandt's avatar
JosefBrandt committed
43
        
JosefBrandt's avatar
 
JosefBrandt committed
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
    def getBrectAndPolygon(self):
        polygons = []
        x0 = None
        for c in self.contours:
            polygon = QtGui.QPolygonF()
            if x0 is None:
                x0 = c[:,0,0].min()
                x1 = c[:,0,0].max()
                y0 = c[:,0,1].min()
                y1 = c[:,0,1].max()
            else:
                x0 = min(x0, c[:,0,0].min())
                x1 = max(x1, c[:,0,0].max())
                y0 = min(y0, c[:,0,1].min())
                y1 = max(y1, c[:,0,1].max())
            for ci in c:
                polygon.append(QtCore.QPointF(ci[0,0],ci[0,1]))
            polygons.append(polygon)
        if x0 is None:
            self.brect = QtCore.QRectF(0,0,1,1)
        else:
            self.brect.setCoords(x0,y0,x1,y1)
        self.polygons = polygons        
JosefBrandt's avatar
JosefBrandt committed
67 68 69 70
    
    def boundingRect(self):
        return self.brect
    
JosefBrandt's avatar
 
JosefBrandt committed
71 72 73 74 75 76 77 78 79 80 81 82
    def mousePressEvent(self, event):
        self.mousePos = self.viewparent.mapToScene(event.pos())
        self.erasing = self.painting = False
        if event.modifiers()==QtCore.Qt.ControlModifier:
            self.painting = True
        elif event.modifiers()==QtCore.Qt.ShiftModifier:
            self.erasing = True
            
        if self.painting or self.erasing:
            drawPos = self.viewparent.mapToScene(event.pos()) - self.brect.topLeft()
            self.drawParticle(drawPos)
            
JosefBrandt's avatar
JosefBrandt committed
83
    def mouseMoveEvent(self, event):
JosefBrandt's avatar
 
JosefBrandt committed
84 85 86 87 88 89
        self.mousePos = self.viewparent.mapToScene(event.pos())
        self.update()
        if self.painting or self.erasing:
            drawPos = self.viewparent.mapToScene(event.pos()) - self.brect.topLeft()
            self.drawParticle(drawPos)
        
JosefBrandt's avatar
JosefBrandt committed
90
    def wheelEvent(self, event):
JosefBrandt's avatar
 
JosefBrandt committed
91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141
        if event.angleDelta().y() < 0:
            self.radius = int(np.clip(self.radius+self.radius*0.1, self.minRadius, self.maxRadius))
        else:
            self.radius = int(np.clip(self.radius-self.radius*0.1, self.minRadius, self.maxRadius))
            
        self.update()
        
    def mouseReleaseEvent(self, event):
        self.erasing = self.painting = False
            
    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Escape:
            self.editorParent.destroyParticlePainter()
        elif event.key() == QtCore.Qt.Key_Return:
            self.editorParent.acceptPaintedResult()
    
    def drawParticle(self, pos):
        img, xmin, ymin, padding = self.contoursToImg(self.contours)
        center = (int(pos.x()+self.radius), int(pos.y()+self.radius))
        if self.painting:
            cv2.circle(img, center, self.radius, 255, -1)
        elif self.erasing:
            cv2.circle(img, center, self.radius, 0, -1)
            
        img = np.uint8(img)
        self.contours = self.imgToCnt(img, xmin, ymin, padding)
        self.getBrectAndPolygon()
        self.update()
    
    def contoursToImg(self, contours):
        cnt = np.vstack(tuple(contours))  #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 = self.radius+2    #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), dtype = np.uint8)
        for curCnt in contours:
            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)
        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)
JosefBrandt's avatar
JosefBrandt committed
142
        else:
JosefBrandt's avatar
 
JosefBrandt committed
143
            temp, contours, hierarchy = cv2.findContours(img, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
JosefBrandt's avatar
JosefBrandt committed
144

JosefBrandt's avatar
 
JosefBrandt committed
145 146 147 148 149 150 151 152 153 154 155
        for contour in contours:
            for i in range(len(contour)):
                contour[i][0][0] += xmin-padding
                contour[i][0][1] += ymin-padding
            
        return contours
    
    def paint(self, painter, option, widget):   
        painter.setPen(QtCore.Qt.white)
        painter.drawRect(self.brect)
        
JosefBrandt's avatar
JosefBrandt committed
156 157 158 159
        if self.mousePos is not None:
            p = [self.mousePos.x(), self.mousePos.y(), self.radius]
            painter.drawEllipse(p[0]-p[2], p[1]-p[2], 2*p[2], 2*p[2])            
            
JosefBrandt's avatar
 
JosefBrandt committed
160 161 162 163
        if self.polygons is not None:            
            for poly in self.polygons:
                painter.setBrush(QtGui.QColor(200, 200, 255, 128))
                painter.drawPolygon(poly)
164
        
JosefBrandt's avatar
 
JosefBrandt committed
165 166 167