particlePainter.py 6.36 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
class ParticlePainter(QtWidgets.QGraphicsItem):
Josef Brandt's avatar
Josef Brandt committed
27
    def __init__(self, editorParent, img, topLeft):
JosefBrandt's avatar
 
JosefBrandt committed
28 29 30 31
        super(ParticlePainter, self).__init__()
        self.editorParent = editorParent
        self.viewparent = self.editorParent.viewparent
        self.setZValue(5)
Josef Brandt's avatar
Josef Brandt committed
32 33 34 35
        self.topLeft = topLeft
        self.img = img
        self.pixmap = None
        
JosefBrandt's avatar
JosefBrandt committed
36
        self.mousePos = None
Josef Brandt's avatar
Josef Brandt committed
37 38 39
        self.painting = False
        self.erasing = False
        
40
        self.minRadius = 5
JosefBrandt's avatar
JosefBrandt committed
41
        self.maxRadius = 500
Josef Brandt's avatar
Josef Brandt committed
42
        self.radius = 30
JosefBrandt's avatar
JosefBrandt committed
43 44
        self.brect = QtCore.QRectF(0,0,1,1)
        
45
        self.setPixmap()
Josef Brandt's avatar
Josef Brandt committed
46 47
        self.setBrect()

48 49
    def setPixmap(self):
        img = self.img.repeat(3).reshape(self.img.shape[0], self.img.shape[1], 3)
Josef Brandt's avatar
Josef Brandt committed
50 51
        height, width, channel = img.shape
        bytesPerLine = 3 * width
52 53
        self.pixmap = QtGui.QPixmap()
        self.pixmap.convertFromImage(QtGui.QImage(img.data, width, height, bytesPerLine, QtGui.QImage.Format_RGB888))
Josef Brandt's avatar
Josef Brandt committed
54 55 56 57 58 59 60

    def setBrect(self):
        x0 = self.topLeft[1]
        y0 = self.topLeft[0]
        x1 = self.topLeft[1] + self.img.shape[1]
        y1 = self.topLeft[0] + self.img.shape[0]
        self.brect.setCoords(x0,y0,x1,y1)
JosefBrandt's avatar
JosefBrandt committed
61 62 63 64
    
    def boundingRect(self):
        return self.brect
    
JosefBrandt's avatar
 
JosefBrandt committed
65 66 67 68 69 70 71 72 73 74 75 76
    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
77
    def mouseMoveEvent(self, event):
JosefBrandt's avatar
 
JosefBrandt committed
78 79 80 81 82 83
        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
84
    def wheelEvent(self, event):
JosefBrandt's avatar
 
JosefBrandt committed
85 86 87 88 89 90 91 92 93 94 95 96 97
        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()
JosefBrandt's avatar
JosefBrandt committed
98
        elif event.key() in [QtCore.Qt.Key_Return, QtCore.Qt.Key_Enter]:
JosefBrandt's avatar
 
JosefBrandt committed
99 100
            self.editorParent.acceptPaintedResult()
    
101 102 103
    def drawParticle(self, pixelPos):
        x_min, x_max = round(pixelPos.x()-self.radius), round(pixelPos.x()+self.radius)
        y_min, y_max = round(pixelPos.y()-self.radius), round(pixelPos.y()+self.radius)
Josef Brandt's avatar
Josef Brandt committed
104
        
105
        x_shift = y_shift = int(0)
Josef Brandt's avatar
Josef Brandt committed
106 107
        
        if x_min < 0:
108 109 110 111
            x_shift = int(x_min)
            self.topLeft[1] -= abs(x_min)
        elif x_max >= self.img.shape[1]:
            x_shift = int(x_max + 1 - self.img.shape[1])
Josef Brandt's avatar
Josef Brandt committed
112 113
        
        if y_min < 0:
114 115 116 117
            self.topLeft[0] -= abs(y_min)
            y_shift = int(y_min)
        elif y_max >= self.img.shape[0]:
            y_shift = int(y_max + 1 - self.img.shape[0])
Josef Brandt's avatar
Josef Brandt committed
118 119
        
        if x_shift != 0 or y_shift != 0:
120 121
            x_range = int(self.img.shape[1] + abs(x_shift))
            y_range = int(self.img.shape[0] + abs(y_shift))
Josef Brandt's avatar
Josef Brandt committed
122 123
            newImg = np.zeros((y_range, x_range))
            if x_shift < 0:
124 125 126 127 128 129 130 131 132 133 134 135 136 137
                if y_shift < 0:
                    newImg[abs(y_shift):, abs(x_shift):] = self.img
                elif y_shift == 0:
                    newImg[:, abs(x_shift):] = self.img
                elif y_shift > 0:
                    newImg[:self.img.shape[0], abs(x_shift):] = self.img
        
            elif x_shift == 0:
                if y_shift < 0:
                    newImg[abs(y_shift):, :] = self.img
                elif y_shift == 0:
                    newImg[:, :] = self.img
                elif y_shift > 0:
                    newImg[:self.img.shape[0], :] = self.img
JosefBrandt's avatar
 
JosefBrandt committed
138
                
139 140 141 142 143 144
            elif x_shift > 0:
                if y_shift < 0:
                    newImg[abs(y_shift):, :self.img.shape[1]] = self.img
                elif y_shift == 0:
                    newImg[:, :self.img.shape[1]] = self.img
                elif y_shift > 0:
145
                    newImg[:self.img.shape[0], :self.img.shape[1]] = self.img
146 147 148
            
            self.img = np.uint8(newImg)
            self.setBrect()
Josef Brandt's avatar
Josef Brandt committed
149
        
150
        center = (int(pixelPos.x()), int(pixelPos.y()))
Josef Brandt's avatar
Josef Brandt committed
151
        
152 153 154 155 156 157 158 159
        if self.painting:
            cv2.circle(self.img, center, self.radius, 255, -1)
        elif self.erasing:
            cv2.circle(self.img, center, self.radius, 0, -1)
            
        self.setPixmap()
        self.update()

JosefBrandt's avatar
 
JosefBrandt committed
160 161 162 163
    def paint(self, painter, option, widget):   
        painter.setPen(QtCore.Qt.white)
        painter.drawRect(self.brect)
        
JosefBrandt's avatar
JosefBrandt committed
164 165 166
        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])            
Josef Brandt's avatar
Josef Brandt committed
167 168
        
        if self.pixmap is not None:
JosefBrandt's avatar
JosefBrandt committed
169
            painter.setOpacity(0.4)
Josef Brandt's avatar
Josef Brandt committed
170
            painter.drawPixmap(self.topLeft[1], self.topLeft[0], self.pixmap)
171
        
JosefBrandt's avatar
 
JosefBrandt committed
172 173 174