class MandelbrotWidget(QWidget):
def __init__(self, parent=None):
super(MandelbrotWidget, self).__init__(parent)
self.thread = RenderThread()
self.pixmap = QPixmap()
self.pixmapOffset = QPoint()
self.lastDragPos = QPoint()
self.centerX = DefaultCenterX
self.centerY = DefaultCenterY
self.pixmapScale = DefaultScale
self.curScale = DefaultScale
self.thread.renderedImage.connect(self.updatePixmap)
self.setWindowTitle("Mandelbrot")
self.setCursor(Qt.CrossCursor)
self.resize(550, 400)
def paintEvent(self, event):
painter = QPainter(self)
painter.fillRect(self.rect(), Qt.black)
if self.pixmap.isNull():
painter.setPen(Qt.white)
painter.drawText(self.rect(), Qt.AlignCenter,
"Rendering initial image, please wait...")
return
if self.curScale == self.pixmapScale:
painter.drawPixmap(self.pixmapOffset, self.pixmap)
else:
scaleFactor = self.pixmapScale / self.curScale
newWidth = int(self.pixmap.width() * scaleFactor)
newHeight = int(self.pixmap.height() * scaleFactor)
newX = self.pixmapOffset.x() + (self.pixmap.width() - newWidth) / 2
newY = self.pixmapOffset.y() + (self.pixmap.height() - newHeight) / 2
painter.save()
painter.translate(newX, newY)
painter.scale(scaleFactor, scaleFactor)
exposed, _ = painter.matrix().inverted()
exposed = exposed.mapRect(self.rect()).adjusted(-1, -1, 1, 1)
painter.drawPixmap(exposed, self.pixmap, exposed)
painter.restore()
text = "Use mouse wheel or the '+' and '-' keys to zoom. Press and " \
"hold left mouse button to scroll."
metrics = painter.fontMetrics()
textWidth = metrics.width(text)
painter.setPen(Qt.NoPen)
painter.setBrush(QColor(0, 0, 0, 127))
painter.drawRect((self.width() - textWidth) / 2 - 5, 0, textWidth + 10,
metrics.lineSpacing() + 5)
painter.setPen(Qt.white)
painter.drawText((self.width() - textWidth) / 2,
metrics.leading() + metrics.ascent(), text)
def resizeEvent(self, event):
self.thread.render(self.centerX, self.centerY, self.curScale, self.size())
def keyPressEvent(self, event):
if event.key() == Qt.Key_Plus:
self.zoom(ZoomInFactor)
elif event.key() == Qt.Key_Minus:
self.zoom(ZoomOutFactor)
elif event.key() == Qt.Key_Left:
self.scroll(-ScrollStep, 0)
elif event.key() == Qt.Key_Right:
self.scroll(+ScrollStep, 0)
elif event.key() == Qt.Key_Down:
self.scroll(0, -ScrollStep)
elif event.key() == Qt.Key_Up:
self.scroll(0, +ScrollStep)
else:
super(MandelbrotWidget, self).keyPressEvent(event)
def wheelEvent(self, event):
numDegrees = event.angleDelta().y() / 8
numSteps = numDegrees / 15.0
self.zoom(pow(ZoomInFactor, numSteps))
def mousePressEvent(self, event):
if event.buttons() == Qt.LeftButton:
self.lastDragPos = QPoint(event.pos())
def mouseMoveEvent(self, event):
if event.buttons() & Qt.LeftButton:
self.pixmapOffset += event.pos() - self.lastDragPos
self.lastDragPos = QPoint(event.pos())
self.update()
def mouseReleaseEvent(self, event):
if event.button() == Qt.LeftButton:
self.pixmapOffset += event.pos() - self.lastDragPos
self.lastDragPos = QPoint()
deltaX = (self.width() - self.pixmap.width()) / 2 - self.pixmapOffset.x()
deltaY = (self.height() - self.pixmap.height()) / 2 - self.pixmapOffset.y()
self.scroll(deltaX, deltaY)
def updatePixmap(self, image, scaleFactor):
if not self.lastDragPos.isNull():
return
self.pixmap = QPixmap.fromImage(image)
self.pixmapOffset = QPoint()
self.lastDragPosition = QPoint()
self.pixmapScale = scaleFactor
self.update()
def zoom(self, zoomFactor):
self.curScale *= zoomFactor
self.update()
self.thread.render(self.centerX, self.centerY, self.curScale,
self.size())
def scroll(self, deltaX, deltaY):
self.centerX += deltaX * self.curScale
self.centerY += deltaY * self.curScale
self.update()
self.thread.render(self.centerX, self.centerY, self.curScale,
self.size())
class GLWidget(QOpenGLWidget, QOpenGLFunctions):
xRotationChanged = Signal(int)
yRotationChanged = Signal(int)
zRotationChanged = Signal(int)
def __init__(self, parent=None):
QOpenGLWidget.__init__(self, parent)
QOpenGLFunctions.__init__(self)
self.core = "--coreprofile" in QCoreApplication.arguments()
self.xRot = 0
self.yRot = 0
self.zRot = 0
self.lastPos = 0
self.logo = Logo()
self.vao = QOpenGLVertexArrayObject()
self.logoVbo = QOpenGLBuffer()
self.program = QOpenGLShaderProgram()
self.projMatrixLoc = 0
self.mvMatrixLoc = 0
self.normalMatrixLoc = 0
self.lightPosLoc = 0
self.proj = QMatrix4x4()
self.camera = QMatrix4x4()
self.world = QMatrix4x4()
self.transparent = "--transparent" in QCoreApplication.arguments()
if self.transparent:
fmt = self.format()
fmt.setAlphaBufferSize(8)
self.setFormat(fmt)
def xRotation(self):
return self.xRot
def yRotation(self):
return self.yRot
def zRotation(self):
return self.zRot
def minimumSizeHint(self):
return QSize(50, 50)
def sizeHint(self):
return QSize(400, 400)
def normalizeAngle(self, angle):
while angle < 0:
angle += 360 * 16
while angle > 360 * 16:
angle -= 360 * 16
return angle
def setXRotation(self, angle):
angle = self.normalizeAngle(angle)
if angle != self.xRot:
self.xRot = angle
self.emit(SIGNAL("xRotationChanged(int)"), angle)
self.update()
def setYRotation(self, angle):
angle = self.normalizeAngle(angle)
if angle != self.yRot:
self.yRot = angle
self.emit(SIGNAL("yRotationChanged(int)"), angle)
self.update()
def setZRotation(self, angle):
angle = self.normalizeAngle(angle)
if angle != self.zRot:
self.zRot = angle
self.emit(SIGNAL("zRotationChanged(int)"), angle)
self.update()
def cleanup(self):
self.makeCurrent()
self.logoVbo.destroy()
del self.program
self.program = None
self.doneCurrent()
def vertexShaderSourceCore(self):
return """#version 150
in vec4 vertex;
in vec3 normal;
out vec3 vert;
out vec3 vertNormal;
uniform mat4 projMatrix;
uniform mat4 mvMatrix;
uniform mat3 normalMatrix;
void main() {
vert = vertex.xyz;
vertNormal = normalMatrix * normal;
gl_Position = projMatrix * mvMatrix * vertex;
}"""
def fragmentShaderSourceCore(self):
return """#version 150
in highp vec3 vert;
in highp vec3 vertNormal;
out highp vec4 fragColor;
uniform highp vec3 lightPos;
void main() {
highp vec3 L = normalize(lightPos - vert);
highp float NL = max(dot(normalize(vertNormal), L), 0.0);
highp vec3 color = vec3(0.39, 1.0, 0.0);
highp vec3 col = clamp(color * 0.2 + color * 0.8 * NL, 0.0, 1.0);
fragColor = vec4(col, 1.0);
}"""
def vertexShaderSource(self):
return """attribute vec4 vertex;
attribute vec3 normal;
varying vec3 vert;
varying vec3 vertNormal;
uniform mat4 projMatrix;
uniform mat4 mvMatrix;
uniform mat3 normalMatrix;
void main() {
vert = vertex.xyz;
vertNormal = normalMatrix * normal;
gl_Position = projMatrix * mvMatrix * vertex;
}"""
def fragmentShaderSource(self):
return """varying highp vec3 vert;
varying highp vec3 vertNormal;
uniform highp vec3 lightPos;
void main() {
highp vec3 L = normalize(lightPos - vert);
highp float NL = max(dot(normalize(vertNormal), L), 0.0);
highp vec3 color = vec3(0.39, 1.0, 0.0);
highp vec3 col = clamp(color * 0.2 + color * 0.8 * NL, 0.0, 1.0);
gl_FragColor = vec4(col, 1.0);
}"""
def initializeGL(self):
self.context().aboutToBeDestroyed.connect(self.cleanup)
self.initializeOpenGLFunctions()
self.glClearColor(0, 0, 0, 1)
self.program = QOpenGLShaderProgram()
if self.core:
self.vertexShader = self.vertexShaderSourceCore()
self.fragmentShader = self.fragmentShaderSourceCore()
else:
self.vertexShader = self.vertexShaderSource()
self.fragmentShader = self.fragmentShaderSource()
self.program.addShaderFromSourceCode(QOpenGLShader.Vertex, self.vertexShader)
self.program.addShaderFromSourceCode(QOpenGLShader.Fragment, self.fragmentShader)
self.program.bindAttributeLocation("vertex", 0)
self.program.bindAttributeLocation("normal", 1)
self.program.link()
self.program.bind()
self.projMatrixLoc = self.program.uniformLocation("projMatrix")
self.mvMatrixLoc = self.program.uniformLocation("mvMatrix")
self.normalMatrixLoc = self.program.uniformLocation("normalMatrix")
self.lightPosLoc = self.program.uniformLocation("lightPos")
self.vao.create()
vaoBinder = QOpenGLVertexArrayObject.Binder(self.vao)
self.logoVbo.create()
self.logoVbo.bind()
float_size = ctypes.sizeof(ctypes.c_float)
self.logoVbo.allocate(self.logo.constData(), self.logo.count() * float_size)
self.setupVertexAttribs()
self.camera.setToIdentity()
self.camera.translate(0, 0, -1)
self.program.setUniformValue(self.lightPosLoc, QVector3D(0, 0, 70))
self.program.release()
vaoBinder = None
def setupVertexAttribs(self):
self.logoVbo.bind()
f = QOpenGLContext.currentContext().functions()
f.glEnableVertexAttribArray(0)
f.glEnableVertexAttribArray(1)
float_size = ctypes.sizeof(ctypes.c_float)
null = VoidPtr(0)
pointer = VoidPtr(3 * float_size)
f.glVertexAttribPointer(0, 3, int(GL.GL_FLOAT), int(GL.GL_FALSE), 6 * float_size, null)
f.glVertexAttribPointer(1, 3, int(GL.GL_FLOAT), int(GL.GL_FALSE), 6 * float_size, pointer)
self.logoVbo.release()
def paintGL(self):
self.glClear(GL.GL_COLOR_BUFFER_BIT | GL.GL_DEPTH_BUFFER_BIT)
self.glEnable(GL.GL_DEPTH_TEST)
self.glEnable(GL.GL_CULL_FACE)
self.world.setToIdentity()
self.world.rotate(180 - (self.xRot / 16), 1, 0, 0)
self.world.rotate(self.yRot / 16, 0, 1, 0)
self.world.rotate(self.zRot / 16, 0, 0, 1)
vaoBinder = QOpenGLVertexArrayObject.Binder(self.vao)
self.program.bind()
self.program.setUniformValue(self.projMatrixLoc, self.proj)
self.program.setUniformValue(self.mvMatrixLoc, self.camera * self.world)
normalMatrix = self.world.normalMatrix()
self.program.setUniformValue(self.normalMatrixLoc, normalMatrix)
self.glDrawArrays(GL.GL_TRIANGLES, 0, self.logo.vertexCount())
self.program.release()
vaoBinder = None
def resizeGL(self, width, height):
self.proj.setToIdentity()
self.proj.perspective(45, width / height, 0.01, 100)
def mousePressEvent(self, event):
self.lastPos = QPoint(event.pos())
def mouseMoveEvent(self, event):
dx = event.x() - self.lastPos.x()
dy = event.y() - self.lastPos.y()
if event.buttons() & Qt.LeftButton:
self.setXRotation(self.xRot + 8 * dy)
self.setYRotation(self.yRot + 8 * dx)
elif event.buttons() & Qt.RightButton:
self.setXRotation(self.xRot + 8 * dy)
self.setZRotation(self.zRot + 8 * dx)
self.lastPos = QPoint(event.pos())
class ElectrodesSynchroWidget(QWidget):
def __init__(self):
super().__init__()
self.setBaseSize(800, 850)
self.setMaximumWidth(800)
self.drawConnection = False
#Function which load electrodes position, all are based on the up-left-quarter electrodes positions
def LoadElectrodesInfos(self):
self.electrodesPos = {}
self.electrodesPos["FP1"] = [
QPoint(self.center.x() - (self.rx * 1 / 10),
self.center.y() - (self.ry * 9 / 10)), False
]
self.electrodesPos["F3"] = [
QPoint(self.center.x() - (self.rx * 2 / 10),
self.center.y() - (self.ry * 5 / 10)), False
]
self.electrodesPos["FZ"] = [
QPoint(self.center.x(),
self.center.y() - (self.ry * 9 / 20)), False
]
self.electrodesPos["F7"] = [
QPoint(self.center.x() - (self.rx * 8 / 20),
self.center.y() - (self.ry * 16 / 20)), False
]
self.electrodesPos["FC7"] = [
QPoint(self.center.x(),
self.center.y() - (self.ry * 2 / 10)), False
]
self.electrodesPos["FT7"] = [
QPoint(self.center.x() - (self.rx * 15 / 20),
self.center.y() - (self.ry * 10 / 20)), False
]
self.electrodesPos["T3"] = [
QPoint(self.center.x() - (self.rx * 9 / 10), self.center.y()),
False
]
self.electrodesPos["C3"] = [
QPoint(self.center.x() - (self.rx * 5 / 10), self.center.y()),
False
]
self.electrodesPos["FC3"] = [
QPoint(self.center.x() - (self.rx * 3 / 10),
self.center.y() - (self.ry * 5 / 20)), False
]
self.electrodesPos["FP2"] = [
QPoint(-self.electrodesPos["FP1"][0].x() + 2 * self.center.x(),
self.electrodesPos["FP1"][0].y()), False
]
self.electrodesPos["F4"] = [
QPoint(-self.electrodesPos["F3"][0].x() + 2 * self.center.x(),
self.electrodesPos["F3"][0].y()), False
]
self.electrodesPos["F8"] = [
QPoint(-self.electrodesPos["F7"][0].x() + 2 * self.center.x(),
self.electrodesPos["F7"][0].y()), False
]
self.electrodesPos["CZ"] = [self.center, False]
self.electrodesPos["C4"] = [
QPoint(-self.electrodesPos["C3"][0].x() + 2 * self.center.x(),
self.electrodesPos["C3"][0].y()), False
]
self.electrodesPos["T4"] = [
QPoint(-self.electrodesPos["T3"][0].x() + 2 * self.center.x(),
self.electrodesPos["T3"][0].y()), False
]
self.electrodesPos["FT8"] = [
QPoint(-self.electrodesPos["FT7"][0].x() + 2 * self.center.x(),
self.electrodesPos["FT7"][0].y()), False
]
self.electrodesPos["FC4"] = [
QPoint(-self.electrodesPos["FC3"][0].x() + 2 * self.center.x(),
self.electrodesPos["FC3"][0].y()), False
]
self.electrodesPos["T5"] = [
QPoint(self.electrodesPos["F7"][0].x(),
-self.electrodesPos["F7"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["P3"] = [
QPoint(self.electrodesPos["F3"][0].x(),
-self.electrodesPos["F3"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["PZ"] = [
QPoint(self.electrodesPos["FZ"][0].x(),
-self.electrodesPos["FZ"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["CPZ"] = [
QPoint(self.electrodesPos["FC7"][0].x(),
-self.electrodesPos["FC7"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["TP8"] = [
QPoint(self.electrodesPos["FT8"][0].x(),
-self.electrodesPos["FT8"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["CP4"] = [
QPoint(-self.electrodesPos["FC3"][0].x() + 2 * self.center.x(),
-self.electrodesPos["FC3"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["P4"] = [
QPoint(self.electrodesPos["F4"][0].x(),
-self.electrodesPos["F4"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["T6"] = [
QPoint(self.electrodesPos["F8"][0].x(),
-self.electrodesPos["F8"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["O1"] = [
QPoint(self.electrodesPos["FP1"][0].x(),
-self.electrodesPos["FP1"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["O2"] = [
QPoint(self.electrodesPos["FP2"][0].x(),
-self.electrodesPos["FP2"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["TP7"] = [
QPoint(-self.electrodesPos["TP8"][0].x() + 2 * self.center.x(),
self.electrodesPos["TP8"][0].y()), False
]
self.electrodesPos["CP3"] = [
QPoint(self.electrodesPos["FC3"][0].x(),
-self.electrodesPos["FC3"][0].y() + 2 * self.center.y()),
False
]
self.electrodesPos["OZ"] = [
QPoint(self.center.x(),
self.center.y() + (self.ry * 19 / 20)), False
]
self.electrodesList = getElectrodesList()
#Function which load the connection from the PSC Matrix Size : Nelec x Nelec
def LoadConnection(self, mat, limitdownvalue, limitupvalue):
self.mat = mat
self.drawConnection = True
self.limitUpValue = limitupvalue
self.limitDownValue = limitdownvalue
self.update()
#Stuff to draw head, connection, electrodes using QPainter
def paintEvent(self, e):
qp = QPainter()
qp.begin(self)
self.drawHead(qp)
self.LoadElectrodesInfos()
if self.drawConnection:
self.drawConnections(qp)
self.drawElectrode(qp,
QPoint(self.center.x() - self.rx - 20,
self.center.y()), "A1") #left ear
self.drawElectrode(qp,
QPoint(self.center.x() + self.rx + 20,
self.center.y()), "A2") #right ear
for key in self.electrodesPos.keys():
self.drawElectrode(qp, self.electrodesPos[key][0], key)
qp.end()
def drawConnections(self, qp):
already_draw = []
for i in range(len(self.mat)):
for j in range(len(self.mat[0])):
if ((i, j) in already_draw or (j, i) in already_draw):
continue
if self.mat[i][j] < self.limitUpValue:
continue
if self.mat[i][j] > self.limitDownValue:
continue
if (self.electrodesPos[self.electrodesList[i]][1] == False):
self.drawElectrode(
qp, self.electrodesPos[self.electrodesList[i]][0],
self.electrodesList[i])
self.electrodesPos[self.electrodesList[i]][1] = True
if (self.electrodesPos[self.electrodesList[j]][1] == False):
self.drawElectrode(
qp, self.electrodesPos[self.electrodesList[j]][0],
self.electrodesList[j])
self.electrodesPos[self.electrodesList[j]][1] = True
self.drawLine(qp,
self.electrodesPos[self.electrodesList[i]][0],
self.electrodesPos[self.electrodesList[j]][0],
round(self.mat[i][j], 2))
already_draw.append((i, j))
def drawHead(self, qp):
pen = QPen(Qt.black, 1, Qt.SolidLine)
qp.setPen(pen)
self.center = QPoint(self.size().width() / 2, self.size().height() / 2)
self.rx = (self.size().width() - 150) / 2
self.ry = (self.size().height() - 100) / 2
qp.drawEllipse(self.center, self.rx, self.ry)
pen.setStyle(Qt.DashLine)
qp.setPen(pen)
qp.drawEllipse(self.center, self.rx - 1 / 10 * self.rx,
self.ry - 1 / 10 * self.ry)
qp.drawLine(self.center.x() - self.rx, self.center.y(),
self.center.x() + self.rx, self.center.y())
qp.drawLine(self.center.x(),
self.center.y() - self.ry, self.center.x(),
self.center.y() + self.ry)
def drawElectrode(self, qp, center, name):
pen = QPen(Qt.black, 1, Qt.SolidLine)
qp.setPen(pen)
qp.setBrush(Qt.black)
qp.drawEllipse(center, 5, 5)
qp.setFont(QFont('Arial', 13))
qp.setBrush(Qt.black)
qp.drawText(center.x() - 25,
center.y() + 5, 50, 40, Qt.AlignCenter, name)
def drawLine(self, qp, start, end, width):
color = self.floatRgb(width, 0, 1)
pen = QPen(QColor(color[0], color[1], color[2], 255), 3, Qt.SolidLine)
qp.setPen(pen)
qp.drawLine(start, end)
def floatRgb(self, n, cmin, cmax):
R, G, B = 1.0, 1.0, 1.0
if (n < cmin):
n = cmin
if (n > cmax):
n = cmax
dv = cmax - cmin
if (n < (cmin + 0.25 * dv)):
R = 0
G = 4 * (n - cmin) / dv
elif (n < (cmin + 0.5 * dv)):
R = 0
B = 1 + 4 * (cmin + 0.25 * dv - n) / dv
elif (n < (cmin + 0.75 * dv)):
R = 4 * (n - cmin - 0.5 * dv) / dv
B = 0
else:
G = 1 + 4 * (cmin + 0.75 * dv - n) / dv
B = 0
return R * 255, G * 255, B * 255
def mapGridToIndex(pos: QPoint) -> int: """Map grid coordinates to a list index.""" return pos.y() * GRID_SIZE + pos.x()
class AppWindow(QMainWindow):
def __init__(self, parent=None):
super(AppWindow, self).__init__(parent)
#Load UI
self.ui = Ui_MainWindow()
self.ui.setupUi(self)
self.setWindowTitle("Annotation Tool")
self.annotation_format = None
#Initialization
self.wiring_signal_slot()
self.init_params()
self.init_rects()
self.init_ui()
self.init_color_dic()
def init_rects(self):
# Init rects which contains rects during labeling.
self.saved_rects = []
self.unsaved_rects = []
self.selected_rect = []
def init_params(self):
# Init some params
self.current_category_index = -1
self.current_category_id = -1
self.current_category_name = ""
self.category_color_dic = {}
self.id_str = ""
self.status_content = ""
self.keep_loading = False
self.frame_index = 0
self.frame_step = 1
self.imgarea_width = 0
self.imgarea_height = 0
self.imgarea_pos = QPoint(0, 0)
self.lefttop_corner = QPoint(0, 0)
self.rightbottom_corner = QPoint(0, 0)
self.cursor_pos = QPoint(0, 0)
self.is_mouse_moving = False
def init_color_dic(self):
#Contains default colors for labeling if user doesn't set any color.
self.category_defaultcolor_dic = {
0: QColor(255, 0, 0),
1: QColor(0, 255, 0),
2: QColor(0, 0, 255),
3: QColor(0, 255, 204),
4: QColor(102, 153, 0),
5: QColor(153, 102, 0),
6: QColor(255, 255, 0),
7: QColor(204, 204, 255),
8: QColor(102, 0, 255),
9: QColor(255, 0, 255),
10: QColor(128, 0, 0),
11: QColor(0, 51, 0),
12: QColor(255, 153, 102),
13: QColor(102, 153, 153),
14: QColor(204, 204, 0),
15: QColor(0, 0, 0)
}
def init_ui(self):
#Add addtional UI setting
self.ui.tableWidget_categories.setColumnCount(2)
self.ui.tableWidget_categories.setColumnWidth(0, 30)
self.ui.tableWidget_categories.setColumnWidth(1, 98)
self.ui.tableWidget_categories.insertRow(0)
self.ui.tableWidget_categories.setSpan(0, 0, 1, 2)
first_row = QTableWidgetItem(" Category Name")
first_row.setFlags(False)
self.ui.tableWidget_categories.setItem(0, 0, first_row)
self.ui.tableWidget_categories.setVerticalHeaderItem(
0, QTableWidgetItem(""))
def wiring_signal_slot(self):
# Wiring signals and slots
self.ui.btn_next.clicked.connect(self.next_frame)
self.ui.btn_back.clicked.connect(self.prev_frame)
self.ui.btn_save.clicked.connect(partial(self.save_label, 0))
self.ui.btn_deleteSelected.clicked.connect(self.delete_label)
self.ui.btn_deleteUnsaved.clicked.connect(self.delete_unsaved)
self.ui.btn_deleteAll.clicked.connect(self.clear_label)
self.ui.btn_applyCategory.clicked.connect(self.apply_category)
self.ui.btn_chooseColor.clicked.connect(self.choose_color)
self.ui.btn_addCategory.clicked.connect(self.add_category)
self.ui.btn_deleteCategory.clicked.connect(self.delete_category)
self.ui.btn_refresh.clicked.connect(self.refresh_image)
self.ui.btn_saveCategory.clicked.connect(self.save_category_name)
self.ui.btn_applySetting.clicked.connect(self.apply_frame_setting)
self.ui.tableWidget_categories.cellClicked.connect(self.cell_clicked)
self.ui.act_openfile.triggered.connect(self.choose_file)
self.ui.act_opendir.triggered.connect(self.choose_dir)
self.ui.act_openUsage.triggered.connect(self.sec_window)
def sec_window(self):
self.sec_window = SecWindow()
self.sec_window.setWindowTitle("Annotation Tool Usage")
self.sec_window.show()
def mousePressEvent(self, event):
if self.ui.img_area.pixmap():
if event.button() == Qt.LeftButton:
print("left")
print(self.ui.img_area.pos())
print(event.pos())
self.is_mouse_moving = True
self.lefttop_corner = event.pos()
if event.button() == Qt.RightButton:
self.cursor_pos = event.pos()
print("right: ", self.cursor_pos.x(), self.cursor_pos.y())
def mouseMoveEvent(self, event):
if self.ui.img_area.pixmap():
self.rightbottom_corner = event.pos()
self.scale_img()
self.calculate_bbox_pos()
self.draw_img()
if self.lefttop_corner != self.rightbottom_corner:
self.unsaved_rects.pop()
def mouseReleaseEvent(self, event):
self.is_mouse_moving = False
if self.ui.img_area.pixmap():
if event.button() == Qt.LeftButton:
self.rightbottom_corner = event.pos()
if self.lefttop_corner != self.rightbottom_corner:
self.selected_rect = []
self.scale_img()
self.calculate_bbox_pos()
try:
if self.unsaved_rects[-1][0] > -1:
self.selected_unsaved_rect_index = self.newly_draw_rect_index
self.selected_rect.append(self.unsaved_rects[
self.selected_unsaved_rect_index])
#self.selected_saved_rect_index, self.selected_unsaved_rect_index = self.find_selected_bbox(event.x(), event.y())
else:
self.unsaved_rects.pop()
self.draw_img()
except:
pass
if event.button() == Qt.RightButton:
self.selected_rect = []
self.selected_saved_rect_index, self.selected_unsaved_rect_index = self.find_selected_bbox(
self.cursor_pos.x(), self.cursor_pos.y())
self.draw_img()
self.ui.tableWidget_categories.clearFocus()
def resizeEvent(self, event):
if self.ui.img_area.pixmap():
self.scale_img()
#self.calculate_bbox_pos()
self.draw_img()
def keyPressEvent(self, event):
keyboard_dic = {
Qt.Key_0: "0",
Qt.Key_1: "1",
Qt.Key_2: "2",
Qt.Key_3: "3",
Qt.Key_4: "4",
Qt.Key_5: "5",
Qt.Key_6: "6",
Qt.Key_7: "7",
Qt.Key_8: "8",
Qt.Key_9: "9"
}
if event.key() in keyboard_dic and self.selected_rect:
self.id_str += keyboard_dic[event.key()]
elif event.key() == Qt.Key_B:
self.prev_frame()
elif event.key() == Qt.Key_N:
self.next_frame()
elif event.key() == Qt.Key_S:
self.save_label(0)
elif event.key() == Qt.Key_D:
self.delete_label()
elif event.key() == Qt.Key_E and self.selected_rect:
if self.id_str != "":
if self.selected_saved_rect_index != -1:
self.saved_rects[self.selected_saved_rect_index][0] = int(
self.id_str)
self.save_label(0)
elif self.selected_unsaved_rect_index != -1:
self.unsaved_rects[
self.selected_unsaved_rect_index][0] = int(self.id_str)
self.draw_img()
self.id_str = ""
else:
print("please input id first")
else:
self.id_str = ""
def show_info(self, target_widget, new_info):
"""
Print info on UI
"""
if isinstance(target_widget, QTextEdit):
target_widget.append(new_info)
elif isinstance(target_widget, QLabel):
target_widget.setText(new_info)
else:
pass
@Slot()
def next_frame(self):
try:
self.wait = False
self.unsaved_rects = []
self.saved_rects = []
self.selected_rect = []
if (self.frame_index + self.frame_step) > self.total_frames:
return
else:
self.frame_index += self.frame_step
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Frame ndex or frame step is empty!"
)
def prev_frame(self):
try:
self.selected_rect = []
self.wait = False
#self.keep_loading = False
if self.frame_index == 0:
return
self.frame_index -= self.frame_step
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> No previous frame!")
def save_label(self, flag):
"""
If flag == 0, save all rects(including unsaved rects);
If flag == 1, update saved rects only(keep unsaved rects unsaved).
"""
try:
self.keep_labeling = False
if flag == 0:
for item in self.unsaved_rects:
self.saved_rects.append(item)
self.unsaved_rects = []
self.draw_img()
self.annotation_format.write_label(self.saved_rects,
self.label_path)
if flag == 0:
self.show_info(self.ui.textEdit_statusBar,
self.img_path + " Saved successfully!")
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Nothing to save!")
def delete_label(self):
try:
if self.selected_saved_rect_index != -1:
# if delete a saved rect, revise label, too.
del self.saved_rects[self.selected_saved_rect_index]
self.save_label(1)
if self.selected_unsaved_rect_index != -1:
# if delete a unsaved rect, no need to revise label file.
del self.unsaved_rects[self.selected_unsaved_rect_index]
self.selected_saved_rect_index = self.selected_unsaved_rect_index = -1
self.selected_rect = []
self.draw_img()
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Nothing to delete!")
def delete_unsaved(self):
try:
self.unsaved_rects = []
self.selected_rect = []
self.draw_img()
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Nothing to delete!")
def clear_label(self):
try:
self.saved_rects = []
self.unsaved_rects = []
self.selected_rect = []
self.save_label(0)
self.draw_img()
except:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Nothing to clear!")
def apply_category(self):
try:
self.current_category_id = int(
self.ui.tableWidget_categories.verticalHeaderItem(
self.current_category_index).text())
except:
print("no text")
current_item = self.ui.tableWidget_categories.currentItem()
try:
self.current_category_name = current_item.text()
except:
self.show_info(self.ui.textEdit_statusBar, "Name is empty!")
if self.current_category_id != -1:
if self.selected_saved_rect_index != -1:
self.saved_rects[self.selected_saved_rect_index][0] = int(
self.current_category_id)
self.save_label(0)
elif self.selected_unsaved_rect_index != -1:
self.unsaved_rects[self.selected_unsaved_rect_index][0] = int(
self.current_category_id)
self.draw_img()
def choose_color(self):
try:
color = QColorDialog.getColor()
except:
print("No color selected!")
if color.isValid() and self.current_category_index != -1:
self.category_color_dic[self.current_category_id] = color
item = self.ui.tableWidget_categories.item(
self.current_category_index, 0)
item.setBackground(
QBrush(self.category_color_dic[self.current_category_id]))
def add_category(self):
row_numbers = self.ui.tableWidget_categories.rowCount()
self.ui.tableWidget_categories.insertRow(row_numbers)
first_item = QTableWidgetItem()
second_item = QTableWidgetItem()
self.ui.tableWidget_categories.setItem(row_numbers, 0, first_item)
self.ui.tableWidget_categories.setItem(row_numbers, 1, second_item)
if row_numbers == 1:
new_header_text = '0'
else:
header = self.ui.tableWidget_categories.verticalHeaderItem(
row_numbers - 1)
last_header_text = header.text()
new_header_text = str(int(last_header_text) + 1)
try:
item = self.ui.tableWidget_categories.item(1, 0)
item.setBackground(
QBrush(self.category_color_dic[self.current_category_id]))
except:
print("no item")
self.ui.tableWidget_categories.setVerticalHeaderItem(
row_numbers, QTableWidgetItem(new_header_text))
def cell_clicked(self):
self.current_category_index = self.ui.tableWidget_categories.currentRow(
)
self.current_category_id = int(
self.ui.tableWidget_categories.verticalHeaderItem(
self.current_category_index).text())
self.ui.tableWidget_categories.setFocus()
print(self.current_category_index)
def delete_category(self):
self.ui.tableWidget_categories.removeRow(self.current_category_index)
def refresh_image(self):
self.draw_img()
def apply_frame_setting(self):
if not self.keep_loading:
return
start_frame = self.ui.plainTextEdit_startFrame.toPlainText()
frame_step = self.ui.plainTextEdit_frameStep.toPlainText()
if start_frame and start_frame != "N.A.":
try:
self.frame_index = int(start_frame)
except:
print("wrong start frame!")
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Wrong start frame!")
return
if frame_step and frame_step != "N.A.":
try:
self.frame_step = int(frame_step)
except:
print("wrong frame step!")
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Wrong frame step!")
return
self.frame_index -= self.frame_step
self.next_frame()
def clear_category_area(self):
row_no = self.ui.tableWidget_categories.rowCount()
for i in reversed(range(1, row_no)):
self.ui.tableWidget_categories.removeRow(i)
def find_selected_bbox(self, x, y):
saved_rect_index = unsaved_rect_index = -1
min_saved_bbox_size = min_unsaved_bbox_size = 0
for i in range(len(self.saved_rects)):
if x >= (self.res_width * self.saved_rects[i][1][0] + self.res_pos_x) and \
x <= (self.res_width * self.saved_rects[i][1][0] + self.res_width * self.saved_rects[i][1][2] + self.res_pos_x) and \
y >= (self.res_height * self.saved_rects[i][1][1] + self.res_pos_y) and \
y <= (self.res_height * self.saved_rects[i][1][1] + self.res_height * self.saved_rects[i][1][3] + self.res_pos_y):
bbox_size = self.saved_rects[i][1][2] * self.saved_rects[i][1][
3]
if min_saved_bbox_size == 0:
min_saved_bbox_size = bbox_size
saved_rect_index = i
elif min_saved_bbox_size > bbox_size:
min_saved_bbox_size = bbox_size
saved_rect_index = i
else:
continue
for j in range(len(self.unsaved_rects)):
if x >= (self.res_width * self.unsaved_rects[j][1][0] + self.res_pos_x) and \
x <= (self.res_width * self.unsaved_rects[j][1][0] + self.res_width * self.unsaved_rects[j][1][2] + self.res_pos_x) and \
y >= (self.res_height * self.unsaved_rects[j][1][1] + + self.res_pos_y) and \
y <= (self.res_height * self.unsaved_rects[j][1][1] + self.res_height * self.unsaved_rects[j][1][3] + self.res_pos_y):
bbox_size = self.unsaved_rects[j][1][2] * self.unsaved_rects[
j][1][3]
if min_unsaved_bbox_size == 0:
min_unsaved_bbox_size = bbox_size
unsaved_rect_index = j
elif min_unsaved_bbox_size > bbox_size:
min_unsaved_bbox_size = bbox_size
unsaved_rect_index = j
else:
continue
if saved_rect_index != -1 and unsaved_rect_index == -1:
self.selected_rect.append(self.saved_rects[saved_rect_index])
elif unsaved_rect_index != -1 and saved_rect_index == -1:
self.selected_rect.append(self.unsaved_rects[unsaved_rect_index])
elif saved_rect_index == -1 and unsaved_rect_index == -1:
self.selected_rect = []
elif saved_rect_index != -1 and unsaved_rect_index != -1:
if min_saved_bbox_size > min_unsaved_bbox_size:
self.selected_rect.append(
self.unsaved_rects[unsaved_rect_index])
saved_rect_index = -1
else:
self.selected_rect.append(self.saved_rects[saved_rect_index])
unsaved_rect_index = -1
return saved_rect_index, unsaved_rect_index
def close_window(self):
sys.exit(app.exec_())
def open_dir(self):
try:
selected_dir = QFileDialog.getExistingDirectory()
return selected_dir
except:
return None
def choose_file(self):
_file_path = self.open_file()
if _file_path:
self.file_path = _file_path[0]
if self.file_path.endswith('.mp4') or self.file_path.endswith(
'.avi') or self.file_path.endswith('.mkv'):
self.set_frame_setting(0)
self.load_video()
#if self.file_path.endswith('.jpg') or self.file_path.endswith('.png'):
# self.load_image(self.file_path)
else:
print("no file is selected!")
def choose_dir(self):
_file_dir = self.open_dir()
parent_dir = os.path.dirname(_file_dir)
if _file_dir:
self.file_dir = _file_dir
self.file_names = []
for filename in os.listdir(_file_dir):
if os.path.splitext(
filename)[-1] == ".jpg" or os.path.splitext(
filename)[-1] == ".png":
self.file_names.append(filename)
txt_name = os.path.splitext(filename)[0] + ".txt"
label_path = os.path.join(self.file_dir,
txt_name).replace('\\', '/')
if not os.path.isfile(label_path):
f = open(label_path, "w")
f.close()
if self.file_names:
self.category_name_file_path = parent_dir + "/category.name"
self.load_category_name(self.category_name_file_path)
self.set_frame_setting(1)
self.clear_history()
self.keep_loading = True
while self.keep_loading:
self.read_image()
self.wait = True
while self.wait:
qApp.processEvents()
time.sleep(0.05)
else:
print("The directory is empty")
def open_file(self):
dialog = QFileDialog(self)
dialog.setNameFilter(
self.tr("Images or Videos (*.jpg *.png *.mp4 *.avi *.mkv)"))
dialog.setViewMode(QFileDialog.Detail)
if dialog.exec_():
try:
filename = dialog.selectedFiles()
return filename
except:
return None
def load_category_name(self, path):
categoty_list = []
if os.path.isfile(path):
with open(path, 'r') as f:
lines = f.readlines()
for i in range(len(lines)):
categoty_list.append(lines[i].strip('\n'))
else:
print("No existing categories!")
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> No existing categories!")
if categoty_list:
self.clear_category_area()
self.show_category_name(categoty_list)
else:
print("Name file is empty!")
def show_category_name(self, lst):
name_no = len(lst)
row_no = self.ui.tableWidget_categories.rowCount()
for i in range(1, row_no):
self.ui.tableWidget_categories.removeRow(i)
for j in range(name_no):
self.ui.tableWidget_categories.insertRow(j + 1)
self.ui.tableWidget_categories.setItem(j + 1, 0,
QTableWidgetItem())
self.ui.tableWidget_categories.setItem(j + 1, 1,
QTableWidgetItem(lst[j]))
self.ui.tableWidget_categories.setVerticalHeaderItem(
j + 1, QTableWidgetItem(str(j)))
def save_category_name(self):
name_list = []
row_no = self.ui.tableWidget_categories.rowCount()
for i in range(row_no - 1):
name_list.append(
self.ui.tableWidget_categories.item(i + 1, 1).text())
with open(self.category_name_file_path, 'w') as f:
lines = []
for j in range(len(name_list)):
if j == len(name_list) - 1:
lines.append(name_list[j])
else:
lines.append(name_list[j] + '\n')
for line in lines:
f.write(line)
self.show_info(
self.ui.textEdit_statusBar, self.category_name_file_path +
" Category name saved successfully!")
def clear_history(self):
"""
Clear the old data of last labeling.
It is no necessary to reset all params, but I do so for safety's sake.
"""
self.frame_index = 0
self.real_frame_index = 0
self.saved_rects = []
self.unsaved_rects = []
self.selected_rect = []
self.selected_saved_rect_index = -1
self.selected_unsaved_rect_index = -1
self.current_category_index = -1
self.current_category_id = -1
self.current_category_name = ""
self.category_color_dic = {}
self.id_str = ""
self.keep_loading = False
self.frame_index = 0
self.frame_step = 1
self.imgarea_width = 0
self.imgarea_height = 0
self.imgarea_pos = QPoint(0, 0)
self.lefttop_corner = QPoint(0, 0)
self.rightbottom_corner = QPoint(0, 0)
self.cursor_pos = QPoint(0, 0)
self.is_mouse_moving = False
def load_video(self):
self.file_path_without_ext = os.path.splitext(self.file_path)[0]
self.file_name = os.path.split(self.file_path_without_ext)[-1]
#Delete old data if it exists
if os.path.isdir(self.file_path_without_ext) == True:
shutil.rmtree(self.file_path_without_ext)
try:
os.mkdir(self.file_path_without_ext)
except:
print("close folder first!")
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Close folder first!")
#TODO: let user to decide weather to keep old data
parent_dir = os.path.dirname(self.file_path)
self.category_name_file_path = parent_dir + "/category.name"
self.clear_category_area()
self.load_category_name(self.category_name_file_path)
self.clear_history()
self.keep_labeling = True # set keep_labeling as False after finished labeling work
while self.keep_labeling:
self.open_cap()
def open_cap(self):
self.cap = cv2.VideoCapture(self.file_path)
self.total_frames = self.cap.get(7)
self.show_info(self.ui.label_totalFrames_2,
str(int(self.total_frames)))
self.num_digits = len(str(int(self.total_frames)))
self.count = 0 #TODO: count should start from current frame index
if self.ui.plainTextEdit_frameStep.toPlainText() != "":
frame_step = self.ui.plainTextEdit_frameStep.toPlainText()
else:
frame_step = self.cap.get(5)
frame_step = int(frame_step)
self.keep_loading = True
while self.keep_loading:
self.read_frame_from_video()
self.wait = True
while self.wait:
qApp.processEvents()
time.sleep(0.05)
self.cap.release()
print("unlocked")
def load_bbox():
pass
def read_frame_from_video(self):
if self.frame_index < 0:
print("Error: index < 0")
return
self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.frame_index)
self.real_frame_index = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
#for i in range(self.step):
#self.cap.grab()
#self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.frame_index)
#f = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
ret, cur_img = self.cap.read()
#ret, cur_img = self.cap.retrieve()
if self.frame_index != self.real_frame_index:
print("Error: Read wrong frame!!! Expect: " +
str(self.frame_index) + " Real: " +
str(int(self.real_frame_index)))
self.show_info(
self.ui.textEdit_statusBar, self.img_path +
"<span style=\" color:#ff0000;\"> Error : Read wrong frame!")
else:
self.show_info(self.ui.label_currentFrame_2, str(self.frame_index))
print("Read correctly Expect: " + str(self.frame_index) +
" Real: " + str(int(self.real_frame_index)))
if ret == True:
img_name = self.file_name + "_frame{}.jpg".format(
str(self.frame_index).zfill(self.num_digits))
label_name = self.file_name + "_frame{}.txt".format(
str(self.frame_index).zfill(self.num_digits))
self.img_path = os.path.join(self.file_path_without_ext,
img_name).replace('\\', '/')
self.label_path = os.path.join(self.file_path_without_ext,
label_name).replace('\\', '/')
if not os.path.isfile(self.label_path):
f = open(self.label_path, "w")
f.close()
if os.path.isfile(self.img_path) == False:
cv2.imwrite(self.img_path, cur_img)
self.setWindowTitle("Tagging: " + self.img_path)
self.raw_pixmap = QPixmap(self.img_path)
#self.raw_pixmap_copy = self.raw_pixmap.copy()
self.scale_img()
self.calculate_bbox_pos()
self.draw_img()
def read_image(self):
self.total_frames = len(self.file_names)
if self.frame_index < self.total_frames:
cur_img_name = self.file_names[self.frame_index]
self.img_path = os.path.join(self.file_dir,
cur_img_name).replace('\\', '/')
self.setWindowTitle("Tagging: " + self.img_path)
# TODO: change title to default at somewhere
image = QPixmap(self.img_path)
self.raw_pixmap = image.copy()
txt_name = os.path.splitext(cur_img_name)[0] + ".txt"
self.label_path = os.path.join(self.file_dir,
txt_name).replace('\\', '/')
self.saved_rects = self.annotation_format.read_label(
self.label_path)
self.scale_img()
self.calculate_bbox_pos()
self.draw_img()
def scale_img(self):
self.imgarea_width = self.ui.img_area.width()
self.imgarea_height = self.ui.img_area.height()
self.imgarea_pos = self.ui.img_area.pos()
self.res_pixmap = self.raw_pixmap.scaled(self.ui.img_area.width(),
self.ui.img_area.height(),
Qt.KeepAspectRatio)
self.ui.img_area.setAlignment(Qt.AlignCenter)
self.res_width = self.res_pixmap.width()
self.res_height = self.res_pixmap.height()
if self.res_width == self.imgarea_width:
self.centralWidget_pos = self.ui.centralWidget.pos()
self.res_pos_x = self.centralWidget_pos.x() + self.imgarea_pos.x()
self.res_pos_y = self.centralWidget_pos.y() + self.imgarea_pos.y(
) + self.imgarea_height / 2 - self.res_height / 2
print("width", self.res_pos_x, self.res_pos_y)
elif self.res_height == self.imgarea_height:
self.centralWidget_pos = self.ui.centralWidget.pos()
self.res_pos_x = self.imgarea_pos.x(
) + self.imgarea_width / 2 - self.res_width / 2
self.res_pos_y = self.centralWidget_pos.y() + self.imgarea_pos.y()
print("height", self.res_pos_x, self.res_pos_y)
def calculate_bbox_pos(self):
self.bbox_pos_x = self.lefttop_corner.x() - self.res_pos_x
self.bbox_pos_y = self.lefttop_corner.y() - self.res_pos_y
self.bbox_width = self.rightbottom_corner.x() - self.lefttop_corner.x()
self.bbox_height = self.rightbottom_corner.y() - self.lefttop_corner.y(
)
if self.is_mouse_moving == False:
self.lefttop_corner = QPoint(0, 0)
self.rightbottom_corner = QPoint(0, 0)
self.calculate_bbox_ratio()
self.save_bbox_ratio()
def calculate_bbox_ratio(self):
self.bbox_pos_x_ratio = self.bbox_pos_x / self.res_width
self.bbox_pos_y_ratio = self.bbox_pos_y / self.res_height
self.bbox_width_ratio = self.bbox_width / self.res_width
self.bbox_height_ratio = self.bbox_height / self.res_height
def save_bbox_ratio(self):
if self.bbox_pos_x_ratio >= 0 and self.bbox_pos_y_ratio >= 0:
lst = []
_lst = []
if self.current_category_id != -1:
lst.append(self.current_category_id)
print("current category: ", self.current_category_id)
else:
lst.append(-1)
_lst.append(self.bbox_pos_x_ratio)
_lst.append(self.bbox_pos_y_ratio)
_lst.append(self.bbox_width_ratio)
_lst.append(self.bbox_height_ratio)
lst.append(_lst)
self.unsaved_rects.append(lst)
self.newly_draw_rect_index = len(self.unsaved_rects) - 1
def draw_img(self):
image = self.res_pixmap.copy()
painter = QPainter(image)
if not self.category_color_dic:
self.category_color_dic = self.category_defaultcolor_dic
defaultcolor_no = len(self.category_defaultcolor_dic)
for i in range(defaultcolor_no - 1):
last_row = self.ui.tableWidget_categories.rowCount() - 1
if i + 1 <= last_row:
item2 = self.ui.tableWidget_categories.item(i + 1, 0)
item2.setBackground(QBrush(self.category_defaultcolor_dic[i]))
elif i + 1 > last_row:
self.add_category()
item2 = self.ui.tableWidget_categories.item(i + 1, 0)
item2.setBackground(QBrush(self.category_defaultcolor_dic[i]))
if self.saved_rects:
for item in self.saved_rects:
if item[0] < len(self.category_defaultcolor_dic) - 1:
painter.setPen(
QPen(self.category_color_dic[int(item[0])], 2))
else:
painter.setPen(QPen(
self.category_color_dic[15],
2)) # TODO: Warn user to assign colors by himself.
_bbox_pos_x = self.res_width * item[1][0]
_bbox_pos_y = self.res_height * item[1][1]
_bbox_width = self.res_width * item[1][2]
_bbox_height = self.res_height * item[1][3]
self.rect = QRectF(_bbox_pos_x, _bbox_pos_y, _bbox_width,
_bbox_height)
painter.drawRect(self.rect)
self.put_txt(item[0], painter)
if self.unsaved_rects:
for item in self.unsaved_rects:
if int(item[0]) > -1:
painter.setPen(
QPen(self.category_color_dic[int(item[0])], 1,
Qt.DashDotDotLine))
_bbox_pos_x = self.res_width * item[1][0]
_bbox_pos_y = self.res_height * item[1][1]
_bbox_width = self.res_width * item[1][2]
_bbox_height = self.res_height * item[1][3]
self.rect = QRectF(_bbox_pos_x, _bbox_pos_y, _bbox_width,
_bbox_height)
painter.drawRect(self.rect)
self.put_txt(item[0], painter)
else:
self.show_info(
self.ui.textEdit_statusBar,
"<span style=\" color:#ff0000;\"> Choose a color first!"
)
# TODO: for now, if dont apply color, still draw rect
self.selected_rect = []
if self.selected_rect:
for item in self.selected_rect:
painter.setPen(QPen(Qt.gray, 2, Qt.DotLine))
_bbox_pos_x = self.res_width * item[1][0]
_bbox_pos_y = self.res_height * item[1][1]
_bbox_width = self.res_width * item[1][2]
_bbox_height = self.res_height * item[1][3]
self.rect = QRectF(_bbox_pos_x, _bbox_pos_y, _bbox_width,
_bbox_height)
painter.drawRect(self.rect)
self.put_txt(item[0], painter)
self.ui.img_area.setPixmap(image)
painter.end()
def put_txt(self, content, painter):
if content != -1:
bbox_text = " " + str(content)
painter.setFont(QFont("Arial", 15))
painter.drawText(int(self.rect.x()), int(self.rect.y()), 40, 40, 1,
bbox_text)
else:
painter.setFont(QFont("Arial", 15))
painter.drawText(int(self.rect.x()), int(self.rect.y()), 40, 40, 1,
" ?")
def unlock_loop(self):
self.wait = False
def set_frame_setting(self, flag):
if flag == 0:
self.ui.plainTextEdit_startFrame.setPlainText("0")
self.ui.plainTextEdit_frameStep.setPlainText("1")
else:
self.ui.plainTextEdit_startFrame.setPlainText("N.A.")
self.ui.plainTextEdit_frameStep.setPlainText("N.A.")
self.ui.label_totalFrames_2.setText("N.A.")
self.ui.label_currentFrame_2.setText("N.A.")
class Screenshot(QGraphicsView):
""" Main Class """
screen_shot_grabed = Signal(QImage)
screen_shot_pos_grabed = Signal(QRect)
widget_closed = Signal()
def __init__(self, flags=constant.DEFAULT, parent=None):
"""
flags: binary flags. see the flags in the constant.py
"""
super().__init__(parent)
# Init
self.penColorNow = QColor(PENCOLOR)
self.penSizeNow = PENSIZE
self.fontNow = QFont('Sans')
self.clipboard = QApplication.clipboard()
self.drawListResult = [
] # draw list that sure to be drew, [action, coord]
self.drawListProcess = None # the process to the result
self.selected_area = QRect(
) # a QRect instance which stands for the selected area
self.selectedAreaRaw = QRect()
self.mousePosition = MousePosition.OUTSIDE_AREA # mouse position
self.screenPixel = None
self.textRect = None
self.mousePressed = False
self.action = ACTION_SELECT
self.mousePoint = self.cursor().pos()
self.startX, self.startY = 0, 0 # the point where you start
self.endX, self.endY = 0, 0 # the point where you end
self.pointPath = QPainterPath(
) # the point mouse passes, used by draw free line
self.items_to_remove = [
] # the items that should not draw on screenshot picture
self.textPosition = None
# result
self.target_img = None
self.target_img_pos = None
# Init window
self.getscreenshot()
self.setWindowFlags(Qt.WindowStaysOnTopHint | Qt.FramelessWindowHint)
self.setMouseTracking(True)
self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
self.setContentsMargins(0, 0, 0, 0)
self.setStyleSheet("QGraphicsView { border-style: none; }")
self.tooBar = MyToolBar(flags, self)
self.tooBar.trigger.connect(self.changeAction)
self.penSetBar = None
if flags & constant.RECT or flags & constant.ELLIPSE or flags & constant.LINE or flags & constant.FREEPEN \
or flags & constant.ARROW or flags & constant.TEXT:
self.penSetBar = PenSetWidget(self)
self.penSetBar.penSizeTrigger.connect(self.changePenSize)
self.penSetBar.penColorTrigger.connect(self.changePenColor)
self.penSetBar.fontChangeTrigger.connect(self.changeFont)
self.textInput = TextInput(self)
self.textInput.inputChanged.connect(self.textChange)
self.textInput.cancelPressed.connect(self.cancelInput)
self.textInput.okPressed.connect(self.okInput)
self.graphics_scene = QGraphicsScene(0, 0, self.screenPixel.width(),
self.screenPixel.height())
self.show()
self.setScene(self.graphics_scene)
self.windowHandle().setScreen(QGuiApplication.screenAt(QCursor.pos()))
self.scale = self.get_scale()
# self.setFixedSize(self.screenPixel.width(), self.screenPixel.height())
self.setGeometry(QGuiApplication.screenAt(QCursor.pos()).geometry())
self.showFullScreen()
self.redraw()
QShortcut(QKeySequence('ctrl+s'),
self).activated.connect(self.saveScreenshot)
QShortcut(QKeySequence('esc'), self).activated.connect(self.close)
@staticmethod
def take_screenshot(flags):
loop = QEventLoop()
screen_shot = Screenshot(flags)
screen_shot.show()
screen_shot.widget_closed.connect(loop.quit)
loop.exec_()
img = screen_shot.target_img
return img
@staticmethod
def take_screenshot_pos(flags):
loop = QEventLoop()
screen_shot = Screenshot(flags)
screen_shot.show()
screen_shot.widget_closed.connect(loop.quit)
loop.exec_()
pos = screen_shot.target_img_pos
return pos
def getscreenshot(self):
screen = QGuiApplication.screenAt(QCursor.pos())
self.screenPixel = screen.grabWindow(0)
def mousePressEvent(self, event):
"""
:type event: QMouseEvent
:param event:
:return:
"""
if event.button() != Qt.LeftButton:
return
if self.action is None:
self.action = ACTION_SELECT
self.startX, self.startY = event.x(), event.y()
if self.action == ACTION_SELECT:
if self.mousePosition == MousePosition.OUTSIDE_AREA:
self.mousePressed = True
self.selected_area = QRect()
self.selected_area.setTopLeft(QPoint(event.x(), event.y()))
self.selected_area.setBottomRight(QPoint(event.x(), event.y()))
self.redraw()
elif self.mousePosition == MousePosition.INSIDE_AREA:
self.mousePressed = True
else:
pass
elif self.action == ACTION_MOVE_SELECTED:
if self.mousePosition == MousePosition.OUTSIDE_AREA:
self.action = ACTION_SELECT
self.selected_area = QRect()
self.selected_area.setTopLeft(QPoint(event.x(), event.y()))
self.selected_area.setBottomRight(QPoint(event.x(), event.y()))
self.redraw()
self.mousePressed = True
elif self.action in DRAW_ACTION:
self.mousePressed = True
if self.action == ACTION_FREEPEN:
self.pointPath = QPainterPath()
self.pointPath.moveTo(QPoint(event.x(), event.y()))
elif self.action == ACTION_TEXT:
if self.textPosition is None:
self.textPosition = QPoint(event.x(), event.y())
self.textRect = None
self.redraw()
def mouseMoveEvent(self, event: QMouseEvent):
"""
:type event: QMouseEvent
:param event:
:return:
"""
self.mousePoint = QPoint(event.globalPos().x(), event.globalPos().y())
if self.action is None:
self.action = ACTION_SELECT
if not self.mousePressed:
point = QPoint(event.x(), event.y())
self.detect_mouse_position(point)
self.setCursorStyle()
self.redraw()
else:
self.endX, self.endY = event.x(), event.y()
# if self.mousePosition != OUTSIDE_AREA:
# self.action = ACTION_MOVE_SELECTED
if self.action == ACTION_SELECT:
self.selected_area.setBottomRight(QPoint(event.x(), event.y()))
self.redraw()
elif self.action == ACTION_MOVE_SELECTED:
self.selected_area = QRect(self.selectedAreaRaw)
if self.mousePosition == MousePosition.INSIDE_AREA:
move_to_x = event.x(
) - self.startX + self.selected_area.left()
move_to_y = event.y(
) - self.startY + self.selected_area.top()
if 0 <= move_to_x <= self.screenPixel.width(
) - 1 - self.selected_area.width():
self.selected_area.moveLeft(move_to_x)
if 0 <= move_to_y <= self.screenPixel.height(
) - 1 - self.selected_area.height():
self.selected_area.moveTop(move_to_y)
self.selected_area = self.selected_area.normalized()
self.selectedAreaRaw = QRect(self.selected_area)
self.startX, self.startY = event.x(), event.y()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_LEFT_SIDE:
move_to_x = event.x(
) - self.startX + self.selected_area.left()
if move_to_x <= self.selected_area.right():
self.selected_area.setLeft(move_to_x)
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_RIGHT_SIDE:
move_to_x = event.x(
) - self.startX + self.selected_area.right()
self.selected_area.setRight(move_to_x)
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_UP_SIDE:
move_to_y = event.y(
) - self.startY + self.selected_area.top()
self.selected_area.setTop(move_to_y)
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_DOWN_SIDE:
move_to_y = event.y(
) - self.startY + self.selected_area.bottom()
self.selected_area.setBottom(move_to_y)
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_TOP_LEFT_CORNER:
move_to_x = event.x(
) - self.startX + self.selected_area.left()
move_to_y = event.y(
) - self.startY + self.selected_area.top()
self.selected_area.setTopLeft(QPoint(move_to_x, move_to_y))
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_BOTTOM_RIGHT_CORNER:
move_to_x = event.x(
) - self.startX + self.selected_area.right()
move_to_y = event.y(
) - self.startY + self.selected_area.bottom()
self.selected_area.setBottomRight(
QPoint(move_to_x, move_to_y))
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_TOP_RIGHT_CORNER:
move_to_x = event.x(
) - self.startX + self.selected_area.right()
move_to_y = event.y(
) - self.startY + self.selected_area.top()
self.selected_area.setTopRight(QPoint(
move_to_x, move_to_y))
self.selected_area = self.selected_area.normalized()
self.redraw()
elif self.mousePosition == MousePosition.ON_THE_BOTTOM_LEFT_CORNER:
move_to_x = event.x(
) - self.startX + self.selected_area.left()
move_to_y = event.y(
) - self.startY + self.selected_area.bottom()
self.selected_area.setBottomLeft(
QPoint(move_to_x, move_to_y))
self.redraw()
else:
pass
elif self.action == ACTION_RECT:
self.drawRect(self.startX, self.startY, event.x(), event.y(),
False)
self.redraw()
pass
elif self.action == ACTION_ELLIPSE:
self.drawEllipse(self.startX, self.startY, event.x(),
event.y(), False)
self.redraw()
elif self.action == ACTION_ARROW:
self.drawArrow(self.startX, self.startY, event.x(), event.y(),
False)
self.redraw()
elif self.action == ACTION_LINE:
self.drawLine(self.startX, self.startY, event.x(), event.y(),
False)
self.redraw()
elif self.action == ACTION_FREEPEN:
y1, y2 = event.x(), event.y()
rect = self.selected_area.normalized()
if y1 <= rect.left():
y1 = rect.left()
elif y1 >= rect.right():
y1 = rect.right()
if y2 <= rect.top():
y2 = rect.top()
elif y2 >= rect.bottom():
y2 = rect.bottom()
self.pointPath.lineTo(y1, y2)
self.drawFreeLine(self.pointPath, False)
self.redraw()
def mouseReleaseEvent(self, event):
"""
:type event: QMouseEvent
:param event:
:return:
"""
if event.button() != Qt.LeftButton:
return
if self.mousePressed:
self.mousePressed = False
self.endX, self.endY = event.x(), event.y()
if self.action == ACTION_SELECT:
self.selected_area.setBottomRight(QPoint(event.x(), event.y()))
self.selectedAreaRaw = QRect(self.selected_area)
self.action = ACTION_MOVE_SELECTED
self.redraw()
elif self.action == ACTION_MOVE_SELECTED:
self.selectedAreaRaw = QRect(self.selected_area)
self.redraw()
# self.action = None
elif self.action == ACTION_RECT:
self.drawRect(self.startX, self.startY, event.x(), event.y(),
True)
self.redraw()
elif self.action == ACTION_ELLIPSE:
self.drawEllipse(self.startX, self.startY, event.x(),
event.y(), True)
self.redraw()
elif self.action == ACTION_ARROW:
self.drawArrow(self.startX, self.startY, event.x(), event.y(),
True)
self.redraw()
elif self.action == ACTION_LINE:
self.drawLine(self.startX, self.startY, event.x(), event.y(),
True)
self.redraw()
elif self.action == ACTION_FREEPEN:
self.drawFreeLine(self.pointPath, True)
self.redraw()
def detect_mouse_position(self, point):
"""
:type point: QPoint
:param point: the mouse position you want to check
:return:
"""
if self.selected_area == QRect():
self.mousePosition = MousePosition.OUTSIDE_AREA
return
if self.selected_area.left() - ERRORRANGE <= point.x(
) <= self.selected_area.left() and (
self.selected_area.top() - ERRORRANGE <= point.y() <=
self.selected_area.top()):
self.mousePosition = MousePosition.ON_THE_TOP_LEFT_CORNER
elif self.selected_area.right() <= point.x(
) <= self.selected_area.right() + ERRORRANGE and (
self.selected_area.top() - ERRORRANGE <= point.y() <=
self.selected_area.top()):
self.mousePosition = MousePosition.ON_THE_TOP_RIGHT_CORNER
elif self.selected_area.left() - ERRORRANGE <= point.x(
) <= self.selected_area.left() and (
self.selected_area.bottom() <= point.y() <=
self.selected_area.bottom() + ERRORRANGE):
self.mousePosition = MousePosition.ON_THE_BOTTOM_LEFT_CORNER
elif self.selected_area.right() <= point.x(
) <= self.selected_area.right() + ERRORRANGE and (
self.selected_area.bottom() <= point.y() <=
self.selected_area.bottom() + ERRORRANGE):
self.mousePosition = MousePosition.ON_THE_BOTTOM_RIGHT_CORNER
elif -ERRORRANGE <= point.x() - self.selected_area.left() <= 0 and (
self.selected_area.topLeft().y() < point.y() <
self.selected_area.bottomLeft().y()):
self.mousePosition = MousePosition.ON_THE_LEFT_SIDE
elif 0 <= point.x() - self.selected_area.right() <= ERRORRANGE and (
self.selected_area.topRight().y() < point.y() <
self.selected_area.bottomRight().y()):
self.mousePosition = MousePosition.ON_THE_RIGHT_SIDE
elif -ERRORRANGE <= point.y() - self.selected_area.top() <= 0 and (
self.selected_area.topLeft().x() < point.x() <
self.selected_area.topRight().x()):
self.mousePosition = MousePosition.ON_THE_UP_SIDE
elif 0 <= point.y() - self.selected_area.bottom() <= ERRORRANGE and (
self.selected_area.bottomLeft().x() < point.x() <
self.selected_area.bottomRight().x()):
self.mousePosition = MousePosition.ON_THE_DOWN_SIDE
elif not self.selected_area.contains(point):
self.mousePosition = MousePosition.OUTSIDE_AREA
else:
self.mousePosition = MousePosition.INSIDE_AREA
def setCursorStyle(self):
if self.action in DRAW_ACTION:
self.setCursor(Qt.CrossCursor)
return
if self.mousePosition == MousePosition.ON_THE_LEFT_SIDE or \
self.mousePosition == MousePosition.ON_THE_RIGHT_SIDE:
self.setCursor(Qt.SizeHorCursor)
elif self.mousePosition == MousePosition.ON_THE_UP_SIDE or \
self.mousePosition == MousePosition.ON_THE_DOWN_SIDE:
self.setCursor(Qt.SizeVerCursor)
elif self.mousePosition == MousePosition.ON_THE_TOP_LEFT_CORNER or \
self.mousePosition == MousePosition.ON_THE_BOTTOM_RIGHT_CORNER:
self.setCursor(Qt.SizeFDiagCursor)
elif self.mousePosition == MousePosition.ON_THE_TOP_RIGHT_CORNER or \
self.mousePosition == MousePosition.ON_THE_BOTTOM_LEFT_CORNER:
self.setCursor(Qt.SizeBDiagCursor)
elif self.mousePosition == MousePosition.OUTSIDE_AREA:
self.setCursor(Qt.ArrowCursor)
elif self.mousePosition == MousePosition.INSIDE_AREA:
self.setCursor(Qt.OpenHandCursor)
else:
self.setCursor(Qt.ArrowCursor)
pass
def drawMagnifier(self):
# First, calculate the magnifier position due to the mouse position
watch_area_width = 16
watch_area_height = 16
cursor_pos = self.mousePoint
watch_area = QRect(
QPoint(cursor_pos.x() - watch_area_width / 2,
cursor_pos.y() - watch_area_height / 2),
QPoint(cursor_pos.x() + watch_area_width / 2,
cursor_pos.y() + watch_area_height / 2))
if watch_area.left() < 0:
watch_area.moveLeft(0)
watch_area.moveRight(watch_area_width)
if self.mousePoint.x(
) + watch_area_width / 2 >= self.screenPixel.width():
watch_area.moveRight(self.screenPixel.width() - 1)
watch_area.moveLeft(watch_area.right() - watch_area_width)
if self.mousePoint.y() - watch_area_height / 2 < 0:
watch_area.moveTop(0)
watch_area.moveBottom(watch_area_height)
if self.mousePoint.y(
) + watch_area_height / 2 >= self.screenPixel.height():
watch_area.moveBottom(self.screenPixel.height() - 1)
watch_area.moveTop(watch_area.bottom() - watch_area_height)
# tricks to solve the hidpi impact on QCursor.pos()
watch_area.setTopLeft(
QPoint(watch_area.topLeft().x() * self.scale,
watch_area.topLeft().y() * self.scale))
watch_area.setBottomRight(
QPoint(watch_area.bottomRight().x() * self.scale,
watch_area.bottomRight().y() * self.scale))
watch_area_pixmap = self.screenPixel.copy(watch_area)
# second, calculate the magnifier area
magnifier_area_width = watch_area_width * 10
magnifier_area_height = watch_area_height * 10
font_area_height = 40
cursor_size = 24
magnifier_area = QRectF(
QPoint(QCursor.pos().x() + cursor_size,
QCursor.pos().y() + cursor_size),
QPoint(QCursor.pos().x() + cursor_size + magnifier_area_width,
QCursor.pos().y() + cursor_size + magnifier_area_height))
if magnifier_area.right() >= self.screenPixel.width():
magnifier_area.moveLeft(QCursor.pos().x() - magnifier_area_width -
cursor_size / 2)
if magnifier_area.bottom(
) + font_area_height >= self.screenPixel.height():
magnifier_area.moveTop(QCursor.pos().y() - magnifier_area_height -
cursor_size / 2 - font_area_height)
# third, draw the watch area to magnifier area
watch_area_scaled = watch_area_pixmap.scaled(
QSize(magnifier_area_width * self.scale,
magnifier_area_height * self.scale))
magnifier_pixmap = self.graphics_scene.addPixmap(watch_area_scaled)
magnifier_pixmap.setOffset(magnifier_area.topLeft())
# then draw lines and text
self.graphics_scene.addRect(QRectF(magnifier_area),
QPen(QColor(255, 255, 255), 2))
self.graphics_scene.addLine(
QLineF(
QPointF(magnifier_area.center().x(), magnifier_area.top()),
QPointF(magnifier_area.center().x(), magnifier_area.bottom())),
QPen(QColor(0, 255, 255), 2))
self.graphics_scene.addLine(
QLineF(
QPointF(magnifier_area.left(),
magnifier_area.center().y()),
QPointF(magnifier_area.right(),
magnifier_area.center().y())),
QPen(QColor(0, 255, 255), 2))
# get the rgb of mouse point
point_rgb = QColor(self.screenPixel.toImage().pixel(self.mousePoint))
# draw information
self.graphics_scene.addRect(
QRectF(
magnifier_area.bottomLeft(),
magnifier_area.bottomRight() +
QPoint(0, font_area_height + 30)), QPen(Qt.black),
QBrush(Qt.black))
rgb_info = self.graphics_scene.addSimpleText(
' Rgb: ({0}, {1}, {2})'.format(point_rgb.red(), point_rgb.green(),
point_rgb.blue()))
rgb_info.setPos(magnifier_area.bottomLeft() + QPoint(0, 5))
rgb_info.setPen(QPen(QColor(255, 255, 255), 2))
rect = self.selected_area.normalized()
size_info = self.graphics_scene.addSimpleText(
' Size: {0} x {1}'.format(rect.width() * self.scale,
rect.height() * self.scale))
size_info.setPos(magnifier_area.bottomLeft() + QPoint(0, 15) +
QPoint(0, font_area_height / 2))
size_info.setPen(QPen(QColor(255, 255, 255), 2))
def get_scale(self):
return self.devicePixelRatio()
def saveScreenshot(self,
clipboard=False,
fileName='screenshot.png',
picType='png'):
fullWindow = QRect(0, 0, self.width() - 1, self.height() - 1)
selected = QRect(self.selected_area)
if selected.left() < 0:
selected.setLeft(0)
if selected.right() >= self.width():
selected.setRight(self.width() - 1)
if selected.top() < 0:
selected.setTop(0)
if selected.bottom() >= self.height():
selected.setBottom(self.height() - 1)
source = (fullWindow & selected)
source.setTopLeft(
QPoint(source.topLeft().x() * self.scale,
source.topLeft().y() * self.scale))
source.setBottomRight(
QPoint(source.bottomRight().x() * self.scale,
source.bottomRight().y() * self.scale))
image = self.screenPixel.copy(source)
if clipboard:
QGuiApplication.clipboard().setImage(image.toImage(),
QClipboard.Clipboard)
else:
image.save(fileName, picType, 10)
self.target_img = image
self.target_img_pos = source
self.screen_shot_grabed.emit(image.toImage())
self.screen_shot_pos_grabed.emit(source)
def redraw(self):
self.graphics_scene.clear()
# draw screenshot
self.graphics_scene.addPixmap(self.screenPixel)
# prepare for drawing selected area
rect = QRectF(self.selected_area)
rect = rect.normalized()
top_left_point = rect.topLeft()
top_right_point = rect.topRight()
bottom_left_point = rect.bottomLeft()
bottom_right_point = rect.bottomRight()
top_middle_point = (top_left_point + top_right_point) / 2
left_middle_point = (top_left_point + bottom_left_point) / 2
bottom_middle_point = (bottom_left_point + bottom_right_point) / 2
right_middle_point = (top_right_point + bottom_right_point) / 2
# draw the picture mask
mask = QColor(0, 0, 0, 155)
if self.selected_area == QRect():
self.graphics_scene.addRect(0, 0, self.screenPixel.width(),
self.screenPixel.height(),
QPen(Qt.NoPen), mask)
else:
self.graphics_scene.addRect(0, 0, self.screenPixel.width(),
top_right_point.y(), QPen(Qt.NoPen),
mask)
self.graphics_scene.addRect(0, top_left_point.y(),
top_left_point.x(), rect.height(),
QPen(Qt.NoPen), mask)
self.graphics_scene.addRect(
top_right_point.x(), top_right_point.y(),
self.screenPixel.width() - top_right_point.x(), rect.height(),
QPen(Qt.NoPen), mask)
self.graphics_scene.addRect(
0, bottom_left_point.y(), self.screenPixel.width(),
self.screenPixel.height() - bottom_left_point.y(),
QPen(Qt.NoPen), mask)
# draw the toolBar
if self.action != ACTION_SELECT:
spacing = 5
# show the toolbar first, then move it to the correct position
# because the width of it may be wrong if this is the first time it shows
self.tooBar.show()
dest = QPointF(rect.bottomRight() -
QPointF(self.tooBar.width(), 0) -
QPointF(spacing, -spacing))
if dest.x() < spacing:
dest.setX(spacing)
pen_set_bar_height = self.penSetBar.height(
) if self.penSetBar is not None else 0
if dest.y() + self.tooBar.height(
) + pen_set_bar_height >= self.height():
if rect.top() - self.tooBar.height(
) - pen_set_bar_height < spacing:
dest.setY(rect.top() + spacing)
else:
dest.setY(rect.top() - self.tooBar.height() -
pen_set_bar_height - spacing)
self.tooBar.move(dest.toPoint())
if self.penSetBar is not None:
self.penSetBar.show()
self.penSetBar.move(dest.toPoint() +
QPoint(0,
self.tooBar.height() + spacing))
if self.action == ACTION_TEXT:
self.penSetBar.showFontWidget()
else:
self.penSetBar.showPenWidget()
else:
self.tooBar.hide()
if self.penSetBar is not None:
self.penSetBar.hide()
# draw the list
for step in self.drawListResult:
self.drawOneStep(step)
if self.drawListProcess is not None:
self.drawOneStep(self.drawListProcess)
if self.action != ACTION_TEXT:
self.drawListProcess = None
if self.selected_area != QRect():
self.items_to_remove = []
# draw the selected rectangle
pen = QPen(QColor(0, 255, 255), 2)
self.items_to_remove.append(self.graphics_scene.addRect(rect, pen))
# draw the drag point
radius = QPoint(3, 3)
brush = QBrush(QColor(0, 255, 255))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(top_left_point - radius, top_left_point + radius),
pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(top_middle_point - radius,
top_middle_point + radius), pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(top_right_point - radius, top_right_point + radius),
pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(left_middle_point - radius,
left_middle_point + radius), pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(right_middle_point - radius,
right_middle_point + radius), pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(bottom_left_point - radius,
bottom_left_point + radius), pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(bottom_middle_point - radius,
bottom_middle_point + radius), pen, brush))
self.items_to_remove.append(
self.graphics_scene.addEllipse(
QRectF(bottom_right_point - radius,
bottom_right_point + radius), pen, brush))
# draw the textedit
if self.textPosition is not None:
textSpacing = 50
position = QPoint()
if self.textPosition.x() + self.textInput.width(
) >= self.screenPixel.width():
position.setX(self.textPosition.x() - self.textInput.width())
else:
position.setX(self.textPosition.x())
if self.textRect is not None:
if self.textPosition.y() + self.textInput.height(
) + self.textRect.height() >= self.screenPixel.height():
position.setY(self.textPosition.y() -
self.textInput.height() -
self.textRect.height())
else:
position.setY(self.textPosition.y() +
self.textRect.height())
else:
if self.textPosition.y() + self.textInput.height(
) >= self.screenPixel.height():
position.setY(self.textPosition.y() -
self.textInput.height())
else:
position.setY(self.textPosition.y())
self.textInput.move(position)
self.textInput.show()
# self.textInput.getFocus()
# draw the magnifier
if self.action == ACTION_SELECT:
self.drawMagnifier()
if self.mousePressed:
self.drawSizeInfo()
if self.action == ACTION_MOVE_SELECTED:
self.drawSizeInfo()
# deal with every step in drawList
def drawOneStep(self, step):
"""
:type step: tuple
"""
if step[0] == ACTION_RECT:
self.graphics_scene.addRect(
QRectF(QPointF(step[1], step[2]), QPointF(step[3], step[4])),
step[5])
elif step[0] == ACTION_ELLIPSE:
self.graphics_scene.addEllipse(
QRectF(QPointF(step[1], step[2]), QPointF(step[3], step[4])),
step[5])
elif step[0] == ACTION_ARROW:
arrow = QPolygonF()
linex = float(step[1] - step[3])
liney = float(step[2] - step[4])
line = sqrt(pow(linex, 2) + pow(liney, 2))
# in case to divided by 0
if line == 0:
return
sinAngel = liney / line
cosAngel = linex / line
# sideLength is the length of bottom side of the body of an arrow
# arrowSize is the size of the head of an arrow, left and right
# sides' size is arrowSize, and the bottom side's size is arrowSize / 2
sideLength = step[5].width()
arrowSize = 8
bottomSize = arrowSize / 2
tmpPoint = QPointF(step[3] + arrowSize * sideLength * cosAngel,
step[4] + arrowSize * sideLength * sinAngel)
point1 = QPointF(step[1] + sideLength * sinAngel,
step[2] - sideLength * cosAngel)
point2 = QPointF(step[1] - sideLength * sinAngel,
step[2] + sideLength * cosAngel)
point3 = QPointF(tmpPoint.x() - sideLength * sinAngel,
tmpPoint.y() + sideLength * cosAngel)
point4 = QPointF(tmpPoint.x() - bottomSize * sideLength * sinAngel,
tmpPoint.y() + bottomSize * sideLength * cosAngel)
point5 = QPointF(step[3], step[4])
point6 = QPointF(tmpPoint.x() + bottomSize * sideLength * sinAngel,
tmpPoint.y() - bottomSize * sideLength * cosAngel)
point7 = QPointF(tmpPoint.x() + sideLength * sinAngel,
tmpPoint.y() - sideLength * cosAngel)
arrow.append(point1)
arrow.append(point2)
arrow.append(point3)
arrow.append(point4)
arrow.append(point5)
arrow.append(point6)
arrow.append(point7)
arrow.append(point1)
self.graphics_scene.addPolygon(arrow, step[5], step[6])
elif step[0] == ACTION_LINE:
self.graphics_scene.addLine(
QLineF(QPointF(step[1], step[2]), QPointF(step[3], step[4])),
step[5])
elif step[0] == ACTION_FREEPEN:
self.graphics_scene.addPath(step[1], step[2])
elif step[0] == ACTION_TEXT:
textAdd = self.graphics_scene.addSimpleText(step[1], step[2])
textAdd.setPos(step[3])
textAdd.setBrush(QBrush(step[4]))
self.textRect = textAdd.boundingRect()
# draw the size information on the top left corner
def drawSizeInfo(self):
sizeInfoAreaWidth = 200
sizeInfoAreaHeight = 30
spacing = 5
rect = self.selected_area.normalized()
sizeInfoArea = QRect(rect.left(),
rect.top() - spacing - sizeInfoAreaHeight,
sizeInfoAreaWidth, sizeInfoAreaHeight)
if sizeInfoArea.top() < 0:
sizeInfoArea.moveTopLeft(rect.topLeft() + QPoint(spacing, spacing))
if sizeInfoArea.right() >= self.screenPixel.width():
sizeInfoArea.moveTopLeft(rect.topLeft() -
QPoint(spacing, spacing) -
QPoint(sizeInfoAreaWidth, 0))
if sizeInfoArea.left() < spacing:
sizeInfoArea.moveLeft(spacing)
if sizeInfoArea.top() < spacing:
sizeInfoArea.moveTop(spacing)
self.items_to_remove.append(
self.graphics_scene.addRect(QRectF(sizeInfoArea), QPen(Qt.white),
QBrush(Qt.black)))
sizeInfo = self.graphics_scene.addSimpleText(' {0} x {1}'.format(
rect.width() * self.scale,
rect.height() * self.scale))
sizeInfo.setPos(sizeInfoArea.topLeft() + QPoint(0, 2))
sizeInfo.setPen(QPen(QColor(255, 255, 255), 2))
self.items_to_remove.append(sizeInfo)
def drawRect(self, x1, x2, y1, y2, result):
rect = self.selected_area.normalized()
tmpRect = QRect(QPoint(x1, x2), QPoint(y1, y2)).normalized()
resultRect = rect & tmpRect
tmp = [
ACTION_RECT,
resultRect.topLeft().x(),
resultRect.topLeft().y(),
resultRect.bottomRight().x(),
resultRect.bottomRight().y(),
QPen(QColor(self.penColorNow), int(self.penSizeNow))
]
if result:
self.drawListResult.append(tmp)
else:
self.drawListProcess = tmp
def drawEllipse(self, x1, x2, y1, y2, result):
rect = self.selected_area.normalized()
tmpRect = QRect(QPoint(x1, x2), QPoint(y1, y2)).normalized()
resultRect = rect & tmpRect
tmp = [
ACTION_ELLIPSE,
resultRect.topLeft().x(),
resultRect.topLeft().y(),
resultRect.bottomRight().x(),
resultRect.bottomRight().y(),
QPen(QColor(self.penColorNow), int(self.penSizeNow))
]
if result:
self.drawListResult.append(tmp)
else:
self.drawListProcess = tmp
def drawArrow(self, x1, x2, y1, y2, result):
rect = self.selected_area.normalized()
if y1 <= rect.left():
y1 = rect.left()
elif y1 >= rect.right():
y1 = rect.right()
if y2 <= rect.top():
y2 = rect.top()
elif y2 >= rect.bottom():
y2 = rect.bottom()
tmp = [
ACTION_ARROW, x1, x2, y1, y2,
QPen(QColor(self.penColorNow), int(self.penSizeNow)),
QBrush(QColor(self.penColorNow))
]
if result:
self.drawListResult.append(tmp)
else:
self.drawListProcess = tmp
def drawLine(self, x1, x2, y1, y2, result):
rect = self.selected_area.normalized()
if y1 <= rect.left():
y1 = rect.left()
elif y1 >= rect.right():
y1 = rect.right()
if y2 <= rect.top():
y2 = rect.top()
elif y2 >= rect.bottom():
y2 = rect.bottom()
tmp = [
ACTION_LINE, x1, x2, y1, y2,
QPen(QColor(self.penColorNow), int(self.penSizeNow))
]
if result:
self.drawListResult.append(tmp)
else:
self.drawListProcess = tmp
def drawFreeLine(self, pointPath, result):
tmp = [
ACTION_FREEPEN,
QPainterPath(pointPath),
QPen(QColor(self.penColorNow), int(self.penSizeNow))
]
if result:
self.drawListResult.append(tmp)
else:
self.drawListProcess = tmp
def textChange(self):
if self.textPosition is None:
return
self.text = self.textInput.getText()
self.drawListProcess = [
ACTION_TEXT,
str(self.text),
QFont(self.fontNow),
QPoint(self.textPosition),
QColor(self.penColorNow)
]
self.redraw()
def undoOperation(self):
if len(self.drawListResult) == 0:
self.action = ACTION_SELECT
self.selected_area = QRect()
self.selectedAreaRaw = QRect()
self.tooBar.hide()
if self.penSetBar is not None:
self.penSetBar.hide()
else:
self.drawListResult.pop()
self.redraw()
def saveOperation(self):
filename = QFileDialog.getSaveFileName(self, 'Save file',
'./screenshot.png',
'*.png;;*.jpg')
if len(filename[0]) == 0:
return
else:
self.saveScreenshot(False, filename[0], filename[1][2:])
self.close()
def close(self):
self.widget_closed.emit()
super().close()
self.tooBar.close()
if self.penSetBar is not None:
self.penSetBar.close()
def saveToClipboard(self):
QApplication.clipboard().setText('Test in save function')
self.saveScreenshot(True)
self.close()
# slots
def changeAction(self, nextAction):
QApplication.clipboard().setText('Test in changeAction function')
if nextAction == ACTION_UNDO:
self.undoOperation()
elif nextAction == ACTION_SAVE:
self.saveOperation()
elif nextAction == ACTION_CANCEL:
self.close()
elif nextAction == ACTION_SURE:
self.saveToClipboard()
else:
self.action = nextAction
self.setFocus()
def changePenSize(self, nextPenSize):
self.penSizeNow = nextPenSize
def changePenColor(self, nextPenColor):
self.penColorNow = nextPenColor
def cancelInput(self):
self.drawListProcess = None
self.textPosition = None
self.textRect = None
self.textInput.hide()
self.textInput.clearText()
self.redraw()
def okInput(self):
self.text = self.textInput.getText()
self.drawListResult.append([
ACTION_TEXT,
str(self.text),
QFont(self.fontNow),
QPoint(self.textPosition),
QColor(self.penColorNow)
])
self.textPosition = None
self.textRect = None
self.textInput.hide()
self.textInput.clearText()
self.redraw()
def changeFont(self, font):
self.fontNow = font
class MandelbrotWidget(QWidget):
def __init__(self, parent=None):
super(MandelbrotWidget, self).__init__(parent)
self.thread = RenderThread()
self.pixmap = QPixmap()
self.pixmapOffset = QPoint()
self.lastDragPos = QPoint()
self.centerX = DefaultCenterX
self.centerY = DefaultCenterY
self.pixmapScale = DefaultScale
self.curScale = DefaultScale
self.thread.renderedImage.connect(self.updatePixmap)
self.setWindowTitle("Mandelbrot")
self.setCursor(Qt.CrossCursor)
self.resize(550, 400)
def paintEvent(self, event):
painter = QPainter(self)
painter.fillRect(self.rect(), Qt.black)
if self.pixmap.isNull():
painter.setPen(Qt.white)
painter.drawText(self.rect(), Qt.AlignCenter,
"Rendering initial image, please wait...")
return
if self.curScale == self.pixmapScale:
painter.drawPixmap(self.pixmapOffset, self.pixmap)
else:
scaleFactor = self.pixmapScale / self.curScale
newWidth = int(self.pixmap.width() * scaleFactor)
newHeight = int(self.pixmap.height() * scaleFactor)
newX = self.pixmapOffset.x() + (self.pixmap.width() - newWidth) / 2
newY = self.pixmapOffset.y() + (self.pixmap.height() - newHeight) / 2
painter.save()
painter.translate(newX, newY)
painter.scale(scaleFactor, scaleFactor)
exposed, _ = painter.matrix().inverted()
exposed = exposed.mapRect(self.rect()).adjusted(-1, -1, 1, 1)
painter.drawPixmap(exposed, self.pixmap, exposed)
painter.restore()
text = "Use mouse wheel or the '+' and '-' keys to zoom. Press and " \
"hold left mouse button to scroll."
metrics = painter.fontMetrics()
textWidth = metrics.width(text)
painter.setPen(Qt.NoPen)
painter.setBrush(QColor(0, 0, 0, 127))
painter.drawRect((self.width() - textWidth) / 2 - 5, 0, textWidth + 10,
metrics.lineSpacing() + 5)
painter.setPen(Qt.white)
painter.drawText((self.width() - textWidth) / 2,
metrics.leading() + metrics.ascent(), text)
def resizeEvent(self, event):
self.thread.render(self.centerX, self.centerY, self.curScale, self.size())
def keyPressEvent(self, event):
if event.key() == Qt.Key_Plus:
self.zoom(ZoomInFactor)
elif event.key() == Qt.Key_Minus:
self.zoom(ZoomOutFactor)
elif event.key() == Qt.Key_Left:
self.scroll(-ScrollStep, 0)
elif event.key() == Qt.Key_Right:
self.scroll(+ScrollStep, 0)
elif event.key() == Qt.Key_Down:
self.scroll(0, -ScrollStep)
elif event.key() == Qt.Key_Up:
self.scroll(0, +ScrollStep)
else:
super(MandelbrotWidget, self).keyPressEvent(event)
def wheelEvent(self, event):
numDegrees = event.angleDelta().y() / 8
numSteps = numDegrees / 15.0
self.zoom(pow(ZoomInFactor, numSteps))
def mousePressEvent(self, event):
if event.buttons() == Qt.LeftButton:
self.lastDragPos = QPoint(event.pos())
def mouseMoveEvent(self, event):
if event.buttons() & Qt.LeftButton:
self.pixmapOffset += event.pos() - self.lastDragPos
self.lastDragPos = QPoint(event.pos())
self.update()
def mouseReleaseEvent(self, event):
if event.button() == Qt.LeftButton:
self.pixmapOffset += event.pos() - self.lastDragPos
self.lastDragPos = QPoint()
deltaX = (self.width() - self.pixmap.width()) / 2 - self.pixmapOffset.x()
deltaY = (self.height() - self.pixmap.height()) / 2 - self.pixmapOffset.y()
self.scroll(deltaX, deltaY)
def updatePixmap(self, image, scaleFactor):
if not self.lastDragPos.isNull():
return
self.pixmap = QPixmap.fromImage(image)
self.pixmapOffset = QPoint()
self.lastDragPosition = QPoint()
self.pixmapScale = scaleFactor
self.update()
def zoom(self, zoomFactor):
self.curScale *= zoomFactor
self.update()
self.thread.render(self.centerX, self.centerY, self.curScale,
self.size())
def scroll(self, deltaX, deltaY):
self.centerX += deltaX * self.curScale
self.centerY += deltaY * self.curScale
self.update()
self.thread.render(self.centerX, self.centerY, self.curScale,
self.size())
class ImageWithMouseControl(QWidget):
wheelSignal = Signal(QWheelEvent, QSize)
moveSignal = Signal(QMouseEvent, QPoint, QPoint)
def __init__(self, imgPath, parent=None):
super().__init__(parent)
self.parent = parent
self.img = QPixmap(imgPath)
self.scaled_img = self.img
self.point = QPoint(0, 0)
def setImagePath(self, imgPath, scaleSize=''):
self.img = QPixmap(imgPath)
if scaleSize:
self.scaled_img = self.img.scaled(scaleSize)
else:
self.scaled_img = self.img
if self.scaled_img.height() and self.scaled_img.width():
self.heightScale = 100 * self.scaled_img.height() // self.scaled_img.width()
else:
self.heightScale = 56
self.point = QPoint(0, 0)
self.update()
def paintEvent(self, e):
painter = QPainter()
painter.begin(self)
self.draw_img(painter)
painter.end()
def recivePaintEvent(self, point):
painter = QPainter()
painter.begin(self)
self.draw_img(painter, point)
painter.end()
def draw_img(self, painter, point=''):
if not point:
painter.drawPixmap(self.point, self.scaled_img)
else:
self.point = point
painter.drawPixmap(self.point, self.scaled_img)
def mouseMoveEvent(self, e): # 重写移动事件
if self.left_click:
self._endPos = e.pos() - self._startPos
self.point = self.point + self._endPos
self.moveSignal.emit(e, self._startPos, self.point) # 向主函数发送坐标信息
self._startPos = e.pos()
self.repaint()
def reciveMoveEvent(self, e, startPos):
self._endPos = e.pos() - startPos
self.point = self.point + self._endPos
self._startPos = e.pos()
self.repaint()
def mousePressEvent(self, e):
if e.button() == Qt.LeftButton:
self.left_click = True
self._startPos = e.pos()
def mouseReleaseEvent(self, e):
if e.button() == Qt.LeftButton:
self.left_click = False
def wheelEvent(self, e):
if e.angleDelta().y() < 0:
# 放大图片
self.scaled_img = self.img.scaled(self.scaled_img.width() - 100, self.scaled_img.height() - self.heightScale)
new_w = e.x() - (self.scaled_img.width() * (e.x() - self.point.x())) / (self.scaled_img.width() + 100)
new_h = e.y() - (self.scaled_img.height() * (e.y() - self.point.y())) / (self.scaled_img.height() + self.heightScale)
self.point = QPoint(new_w, new_h)
self.repaint()
elif e.angleDelta().y() > 0:
# 缩小图片
self.scaled_img = self.img.scaled(self.scaled_img.width() + 100, self.scaled_img.height() + self.heightScale)
try:
new_w = e.x() - (self.scaled_img.width() * (e.x() - self.point.x())) / (self.scaled_img.width() - 100)
new_h = e.y() - (self.scaled_img.height() * (e.y() - self.point.y())) / (self.scaled_img.height() - self.heightScale)
self.point = QPoint(new_w, new_h)
self.repaint()
except:
pass
self.wheelSignal.emit(e, self.scaled_img.size())
def reciveWheelEvent(self, e):
if e.angleDelta().y() < 0:
# 放大图片
self.scaled_img = self.img.scaled(self.scaled_img.width() - 100, self.scaled_img.height() - self.heightScale)
new_w = e.x() - (self.scaled_img.width() * (e.x() - self.point.x())) / (self.scaled_img.width() + 100)
new_h = e.y() - (self.scaled_img.height() * (e.y() - self.point.y())) / (self.scaled_img.height() + self.heightScale)
self.point = QPoint(new_w, new_h)
self.repaint()
elif e.angleDelta().y() > 0:
# 缩小图片
self.scaled_img = self.img.scaled(self.scaled_img.width() + 100, self.scaled_img.height() + self.heightScale)
try:
new_w = e.x() - (self.scaled_img.width() * (e.x() - self.point.x())) / (self.scaled_img.width() - 100)
new_h = e.y() - (self.scaled_img.height() * (e.y() - self.point.y())) / (self.scaled_img.height() - self.heightScale)
self.point = QPoint(new_w, new_h)
self.repaint()
except:
pass
def resizeEvent(self, e):
if self.parent is not None:
self.scaled_img = self.img.scaled(self.size())
self.point = QPoint(0, 0)
self.update()
def drawSeparationLine(self, painter, pen, start: QPoint, width):
""" Helper function for `self.paint()`, drawing the separation line """
end = QPoint(start.x() + width,
start.y() + self._separationLineQThickness)
separationRect = QRect(start, end)
painter.fillRect(separationRect, QColor("lightGray"))
