├── results.gif
├── testing_images
└── testing.jpeg
├── testing_images_out
├── testing_1.png
├── testing_2.png
└── testing_3.png
├── ustr.py
├── zoomWidget.py
├── LICENSE
├── toolBar.py
├── README.md
├── lib.py
├── grab_cut.py
├── shape.py
├── app.py
└── canvas.py
/results.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zihuaweng/Interactive-image-segmentation-opencv-qt/HEAD/results.gif
--------------------------------------------------------------------------------
/testing_images/testing.jpeg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zihuaweng/Interactive-image-segmentation-opencv-qt/HEAD/testing_images/testing.jpeg
--------------------------------------------------------------------------------
/testing_images_out/testing_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zihuaweng/Interactive-image-segmentation-opencv-qt/HEAD/testing_images_out/testing_1.png
--------------------------------------------------------------------------------
/testing_images_out/testing_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zihuaweng/Interactive-image-segmentation-opencv-qt/HEAD/testing_images_out/testing_2.png
--------------------------------------------------------------------------------
/testing_images_out/testing_3.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zihuaweng/Interactive-image-segmentation-opencv-qt/HEAD/testing_images_out/testing_3.png
--------------------------------------------------------------------------------
/ustr.py:
--------------------------------------------------------------------------------
1 | import sys
2 |
3 |
4 | def ustr(x):
5 | '''py2/py3 unicode helper'''
6 |
7 | if sys.version_info < (3, 0, 0):
8 | from PyQt4.QtCore import QString
9 | if type(x) == str:
10 | return x.decode('utf-8')
11 | if type(x) == QString:
12 | return unicode(x)
13 | return x
14 | else:
15 | return x # py3
16 |
--------------------------------------------------------------------------------
/zoomWidget.py:
--------------------------------------------------------------------------------
1 | from PyQt5.QtGui import *
2 | from PyQt5.QtCore import *
3 | from PyQt5.QtWidgets import *
4 |
5 |
6 | class ZoomWidget(QSpinBox):
7 |
8 | def __init__(self, value=100):
9 | super(ZoomWidget, self).__init__()
10 | self.setButtonSymbols(QAbstractSpinBox.NoButtons)
11 | self.setRange(1, 500)
12 | self.setSuffix(' %')
13 | self.setValue(value)
14 | self.setToolTip(u'Zoom Level')
15 | self.setStatusTip(self.toolTip())
16 | self.setAlignment(Qt.AlignCenter)
17 |
18 | def minimumSizeHint(self):
19 | height = super(ZoomWidget, self).minimumSizeHint().height()
20 | fm = QFontMetrics(self.font())
21 | width = fm.width(str(self.maximum()))
22 | return QSize(width, height)
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2018 Harrietong
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/toolBar.py:
--------------------------------------------------------------------------------
1 | from PyQt5.QtGui import *
2 | from PyQt5.QtCore import *
3 | from PyQt5.QtWidgets import *
4 |
5 |
6 | class ToolBar(QToolBar):
7 |
8 | def __init__(self, title):
9 | super(ToolBar, self).__init__(title)
10 | layout = self.layout()
11 | m = (0, 0, 0, 0)
12 | layout.setSpacing(0)
13 | layout.setContentsMargins(*m)
14 | self.setContentsMargins(*m)
15 | self.setWindowFlags(self.windowFlags() | Qt.FramelessWindowHint)
16 |
17 | def addAction(self, action):
18 | if isinstance(action, QWidgetAction):
19 | return super(ToolBar, self).addAction(action)
20 | btn = ToolButton()
21 | btn.setDefaultAction(action)
22 | btn.setToolButtonStyle(self.toolButtonStyle())
23 | self.addWidget(btn)
24 |
25 |
26 | class ToolButton(QToolButton):
27 | """ToolBar companion class which ensures all buttons have the same size."""
28 | minSize = (60, 60)
29 |
30 | def minimumSizeHint(self):
31 | ms = super(ToolButton, self).minimumSizeHint()
32 | w1, h1 = ms.width(), ms.height()
33 | w2, h2 = self.minSize
34 | ToolButton.minSize = max(w1, w2), max(h1, h2)
35 | return QSize(*ToolButton.minSize)
36 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Interactive-image-segmentation-opencv-qt V-1.0
2 | ## 使用opencv进行交互式抠图
3 | This is a simple interactive image segmentation app written by python opencv and pyqt.
4 |
5 | This app applys Grabcut algorithm in opencv for matting images. Grabcut is an improved version of Graphcut algorithm. Check these papers([paper1](http://www.cs.cornell.edu/~rdz/Papers/BVZ-pami01-final.pdf), [paper2](http://www.csd.uwo.ca/~yuri/Papers/iccv01.pdf)) for detail information~~
6 |
7 | The gui part is mostly from this great work [labelImg](https://github.com/tzutalin/labelImg). It is a very good example for pygt beginner!
8 |
9 |
10 | ## Requirement:
11 | - Ubuntu 16.04
12 | - python3
13 | - pyqt5
14 | - cv2
15 |
16 | ## Usage:
17 | python app.py
18 |
19 | ## Result:
20 |
21 | 
22 |
23 | ### input:
24 |
25 |
26 | ### output:
27 |
28 |
29 |
30 | ## TODO:
31 | - Auto resize images in mainwindow according to mouse movement
32 | - Image item switch
33 | - Rect edit and auto matting
34 | - Show only one Rect every selection
35 |
36 |
37 |
38 |
--------------------------------------------------------------------------------
/lib.py:
--------------------------------------------------------------------------------
1 | from math import sqrt
2 | from ustr import ustr
3 | import hashlib
4 | from PyQt5.QtGui import *
5 | from PyQt5.QtCore import *
6 | from PyQt5.QtWidgets import *
7 |
8 |
9 | def newIcon(icon):
10 | return QIcon(':/' + icon)
11 |
12 |
13 | def newButton(text, icon=None, slot=None):
14 | b = QPushButton(text)
15 | if icon is not None:
16 | b.setIcon(newIcon(icon))
17 | if slot is not None:
18 | b.clicked.connect(slot)
19 | return b
20 |
21 |
22 | def newAction(parent, text, slot=None, shortcut=None, icon=None,
23 | tip=None, checkable=False, enabled=True):
24 | """Create a new action and assign callbacks, shortcuts, etc."""
25 | a = QAction(text, parent)
26 | if icon is not None:
27 | a.setIcon(newIcon(icon))
28 | if shortcut is not None:
29 | if isinstance(shortcut, (list, tuple)):
30 | a.setShortcuts(shortcut)
31 | else:
32 | a.setShortcut(shortcut)
33 | if tip is not None:
34 | a.setToolTip(tip)
35 | a.setStatusTip(tip)
36 | if slot is not None:
37 | a.triggered.connect(slot)
38 | if checkable:
39 | a.setCheckable(True)
40 | a.setEnabled(enabled)
41 | return a
42 |
43 |
44 | def addActions(widget, actions):
45 | for action in actions:
46 | if action is None:
47 | widget.addSeparator()
48 | elif isinstance(action, QMenu):
49 | widget.addMenu(action)
50 | else:
51 | widget.addAction(action)
52 |
53 |
54 | def labelValidator():
55 | return QRegExpValidator(QRegExp(r'^[^ \t].+'), None)
56 |
57 |
58 | class struct(object):
59 |
60 | def __init__(self, **kwargs):
61 | self.__dict__.update(kwargs)
62 |
63 |
64 | def distance(p):
65 | return sqrt(p.x() * p.x() + p.y() * p.y())
66 |
67 |
68 | def fmtShortcut(text):
69 | mod, key = text.split('+', 1)
70 | return '%s+%s' % (mod, key)
71 |
72 |
73 | def generateColorByText(text):
74 | s = str(ustr(text))
75 | hashCode = int(hashlib.sha256(s.encode('utf-8')).hexdigest(), 16)
76 | r = int((hashCode / 255) % 255)
77 | g = int((hashCode / 65025) % 255)
78 | b = int((hashCode / 16581375) % 255)
79 | return QColor(r, g, b, 100)
80 |
--------------------------------------------------------------------------------
/grab_cut.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import cv2
3 |
4 |
5 | class Grab_cut(object):
6 | suffix = '.jpg'
7 |
8 | def __init__(self, filename=None):
9 | self.filename = filename
10 | self.height = None
11 | self.width = None
12 |
13 | def image_matting(self, image_file, shape, iteration=10):
14 | points = shape['points']
15 | xmin, ymin, xmax, ymax = Grab_cut.convertPoints2BndBox(points)
16 | self.width = xmax - xmin
17 | self.height = ymax - ymin
18 |
19 | src_img = cv2.imread(image_file)
20 |
21 | mask = np.zeros(src_img.shape[:2], np.uint8)
22 | bgdModel = np.zeros((1, 65), np.float64)
23 | fgdModel = np.zeros((1, 65), np.float64)
24 | rect = (xmin, ymin, self.width, self.height)
25 |
26 | # Grabcut
27 | cv2.grabCut(src_img, mask, rect, bgdModel, fgdModel,
28 | iteration, cv2.GC_INIT_WITH_RECT)
29 |
30 | r_channel, g_channel, b_channel = cv2.split(src_img)
31 | a_channel = np.where((mask == 2) | (mask == 0), 0, 255).astype('uint8')
32 |
33 | # crop image space
34 | for row in range(ymin, ymax):
35 | if sum(r_channel[row, xmin:xmax + 1]) > 0:
36 | out_ymin = row
37 | break
38 | for row in range(ymin, ymax)[::-1]:
39 | if sum(r_channel[row, xmin:xmax + 1]) > 0:
40 | out_ymax = row + 1
41 | break
42 | for col in range(xmin, xmax):
43 | if sum(a_channel[ymin:ymax + 1, col]) > 0:
44 | out_xmin = col
45 | break
46 | for col in range(xmin, xmax)[::-1]:
47 | if sum(a_channel[ymin:ymax + 1, col]) > 0:
48 | out_xmax = col + 1
49 | break
50 |
51 | # output image
52 | img_RGBA = cv2.merge((r_channel[out_ymin:out_ymax, out_xmin:out_xmax],
53 | g_channel[out_ymin:out_ymax, out_xmin:out_xmax],
54 | b_channel[out_ymin:out_ymax, out_xmin:out_xmax],
55 | a_channel[out_ymin:out_ymax, out_xmin:out_xmax]))
56 |
57 | return img_RGBA
58 |
59 | @staticmethod
60 | def convertPoints2BndBox(points):
61 | xmin = float('inf')
62 | ymin = float('inf')
63 | xmax = float('-inf')
64 | ymax = float('-inf')
65 | for p in points:
66 | x = p[0]
67 | y = p[1]
68 | xmin = min(x, xmin)
69 | ymin = min(y, ymin)
70 | xmax = max(x, xmax)
71 | ymax = max(y, ymax)
72 |
73 | # Martin Kersner, 2015/11/12
74 | # 0-valued coordinates of BB caused an error while
75 | # training faster-rcnn object detector.
76 | if xmin < 1:
77 | xmin = 1
78 |
79 | if ymin < 1:
80 | ymin = 1
81 |
82 | return (int(xmin), int(ymin), int(xmax), int(ymax))
83 |
84 | @staticmethod
85 | def resultSave(save_path, image_np):
86 | cv2.imwrite(save_path, image_np)
87 |
--------------------------------------------------------------------------------
/shape.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/python
2 | # -*- coding: utf-8 -*-
3 |
4 |
5 | from PyQt5.QtGui import *
6 | from PyQt5.QtCore import *
7 | from lib import distance
8 |
9 | DEFAULT_LINE_COLOR = QColor(0, 255, 0, 128)
10 | DEFAULT_FILL_COLOR = QColor(255, 0, 0, 128)
11 | DEFAULT_SELECT_LINE_COLOR = QColor(255, 255, 255)
12 | DEFAULT_SELECT_FILL_COLOR = QColor(0, 128, 255, 155)
13 | DEFAULT_VERTEX_FILL_COLOR = QColor(0, 255, 0, 255)
14 | DEFAULT_HVERTEX_FILL_COLOR = QColor(255, 0, 0)
15 |
16 |
17 | class Shape(object):
18 | P_SQUARE, P_ROUND = range(2)
19 |
20 | MOVE_VERTEX, NEAR_VERTEX = range(2)
21 |
22 | # The following class variables influence the drawing
23 | # of _all_ shape objects.
24 | line_color = DEFAULT_LINE_COLOR
25 | fill_color = DEFAULT_FILL_COLOR
26 | select_line_color = DEFAULT_SELECT_LINE_COLOR
27 | select_fill_color = DEFAULT_SELECT_FILL_COLOR
28 | vertex_fill_color = DEFAULT_VERTEX_FILL_COLOR
29 | hvertex_fill_color = DEFAULT_HVERTEX_FILL_COLOR
30 | point_type = P_ROUND
31 | point_size = 8
32 | scale = 1.0
33 |
34 | def __init__(self, label=None, line_color=None,difficult = False):
35 | self.label = label
36 | self.points = []
37 | self.fill = False
38 | self.selected = False
39 | self.difficult = difficult
40 |
41 | self._highlightIndex = None
42 | self._highlightMode = self.NEAR_VERTEX
43 | self._highlightSettings = {
44 | self.NEAR_VERTEX: (4, self.P_ROUND),
45 | self.MOVE_VERTEX: (1.5, self.P_SQUARE),
46 | }
47 |
48 | self._closed = False
49 |
50 | if line_color is not None:
51 | # Override the class line_color attribute
52 | # with an object attribute. Currently this
53 | # is used for drawing the pending line a different color.
54 | self.line_color = line_color
55 |
56 | def close(self):
57 | self._closed = True
58 |
59 | def reachMaxPoints(self):
60 | if len(self.points) >= 4:
61 | return True
62 | return False
63 |
64 | def addPoint(self, point):
65 | if not self.reachMaxPoints():
66 | self.points.append(point)
67 |
68 | def popPoint(self):
69 | if self.points:
70 | return self.points.pop()
71 | return None
72 |
73 | def isClosed(self):
74 | return self._closed
75 |
76 | def setOpen(self):
77 | self._closed = False
78 |
79 | def paint(self, painter):
80 | if self.points:
81 | color = self.select_line_color if self.selected else self.line_color
82 | pen = QPen(color)
83 | # Try using integer sizes for smoother drawing(?)
84 | pen.setWidth(max(1, int(round(2.0 / self.scale))))
85 | painter.setPen(pen)
86 |
87 | line_path = QPainterPath()
88 | vrtx_path = QPainterPath()
89 |
90 | line_path.moveTo(self.points[0])
91 | # Uncommenting the following line will draw 2 paths
92 | # for the 1st vertex, and make it non-filled, which
93 | # may be desirable.
94 | #self.drawVertex(vrtx_path, 0)
95 |
96 | for i, p in enumerate(self.points):
97 | line_path.lineTo(p)
98 | self.drawVertex(vrtx_path, i)
99 | if self.isClosed():
100 | line_path.lineTo(self.points[0])
101 |
102 | painter.drawPath(line_path)
103 | painter.drawPath(vrtx_path)
104 | painter.fillPath(vrtx_path, self.vertex_fill_color)
105 | if self.fill:
106 | color = self.select_fill_color if self.selected else self.fill_color
107 | painter.fillPath(line_path, color)
108 |
109 | def drawVertex(self, path, i):
110 | d = self.point_size / self.scale
111 | shape = self.point_type
112 | point = self.points[i]
113 | if i == self._highlightIndex:
114 | size, shape = self._highlightSettings[self._highlightMode]
115 | d *= size
116 | if self._highlightIndex is not None:
117 | self.vertex_fill_color = self.hvertex_fill_color
118 | else:
119 | self.vertex_fill_color = Shape.vertex_fill_color
120 | if shape == self.P_SQUARE:
121 | path.addRect(point.x() - d / 2, point.y() - d / 2, d, d)
122 | elif shape == self.P_ROUND:
123 | path.addEllipse(point, d / 2.0, d / 2.0)
124 | else:
125 | assert False, "unsupported vertex shape"
126 |
127 | def nearestVertex(self, point, epsilon):
128 | for i, p in enumerate(self.points):
129 | if distance(p - point) <= epsilon:
130 | return i
131 | return None
132 |
133 | def containsPoint(self, point):
134 | return self.makePath().contains(point)
135 |
136 | def makePath(self):
137 | path = QPainterPath(self.points[0])
138 | for p in self.points[1:]:
139 | path.lineTo(p)
140 | return path
141 |
142 | def boundingRect(self):
143 | return self.makePath().boundingRect()
144 |
145 | def moveBy(self, offset):
146 | self.points = [p + offset for p in self.points]
147 |
148 | def moveVertexBy(self, i, offset):
149 | self.points[i] = self.points[i] + offset
150 |
151 | def highlightVertex(self, i, action):
152 | self._highlightIndex = i
153 | self._highlightMode = action
154 |
155 | def highlightClear(self):
156 | self._highlightIndex = None
157 |
158 | def copy(self):
159 | shape = Shape("%s" % self.label)
160 | shape.points = [p for p in self.points]
161 | shape.fill = self.fill
162 | shape.selected = self.selected
163 | shape._closed = self._closed
164 | if self.line_color != Shape.line_color:
165 | shape.line_color = self.line_color
166 | if self.fill_color != Shape.fill_color:
167 | shape.fill_color = self.fill_color
168 | shape.difficult = self.difficult
169 | return shape
170 |
171 | def __len__(self):
172 | return len(self.points)
173 |
174 | def __getitem__(self, key):
175 | return self.points[key]
176 |
177 | def __setitem__(self, key, value):
178 | self.points[key] = value
179 |
--------------------------------------------------------------------------------
/app.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 |
4 | from PyQt5.QtGui import *
5 | from PyQt5.QtCore import *
6 | from PyQt5.QtWidgets import *
7 |
8 | from functools import partial
9 |
10 | from toolBar import ToolBar
11 | from canvas import Canvas
12 | from lib import newIcon
13 | from zoomWidget import ZoomWidget
14 | from grab_cut import Grab_cut
15 | import cv2
16 |
17 | __appname__ = 'ImageMatting'
18 | defaultFilename = '.'
19 |
20 |
21 | class WindowMixin(object):
22 |
23 | def menu(self, title, actions=None):
24 | menu = self.menuBar().addMenu(title)
25 | if actions:
26 | addActions(menu, actions)
27 | return menu
28 |
29 | def toolbar(self, title, actions=None):
30 | toolbar = ToolBar(title)
31 | toolbar.setObjectName('{}ToolBar'.format(title))
32 | toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
33 | if actions:
34 | addActions(toolbar, actions)
35 | self.addToolBar(Qt.LeftToolBarArea, toolbar)
36 | return toolbar
37 |
38 |
39 | class ResizedQWidget(QWidget):
40 | def sizeHint(self):
41 | return QSize(100, 150)
42 |
43 |
44 | def newAction(parent, text, slot=None, shortcut=None,
45 | tip=None, icon=None, checkable=False,
46 | enable=True):
47 | a = QAction(text, parent)
48 | if icon is not None:
49 | a.setIcon(QIcon(icon))
50 | if shortcut is not None:
51 | a.setShortcut(shortcut)
52 | if tip is not None:
53 | a.setToolTip(tip)
54 | a.setStatusTip(tip)
55 | if slot is not None:
56 | a.triggered.connect(slot)
57 | if checkable:
58 | a.setCheckable(True)
59 | a.setEnabled(enable)
60 | return a
61 |
62 |
63 | def addActions(widget, actions):
64 | for action in actions:
65 | if action is None:
66 | widget.addSeparator()
67 | elif isinstance(action, QMenu):
68 | widget.addMenu(action)
69 | else:
70 | widget.addAction(action)
71 |
72 |
73 | class struct(object):
74 |
75 | def __init__(self, **kwargs):
76 | self.__dict__.update(kwargs)
77 |
78 |
79 | class MainWindow(QMainWindow, WindowMixin):
80 | FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
81 |
82 | def __init__(self, defaultFilename=None):
83 | super().__init__()
84 |
85 | self.dirty = True
86 | self.mImgList = []
87 | self.dirname = None
88 | self._beginner = True
89 |
90 | self.image_out_np = None
91 | self.default_save_dir = None
92 | # Application state
93 | self.filePath = None
94 | self.mattingFile = None
95 |
96 | listLayout = QVBoxLayout()
97 | listLayout.setContentsMargins(0, 0, 0, 0)
98 | matResultShow = ResizedQWidget()
99 | matResultShow.resize(150, 150)
100 |
101 | self.pic = QLabel(matResultShow)
102 | self.pic.resize(150, 150)
103 | self.pic.setGeometry(50, 20, 150, 150)
104 |
105 | # self.pic.resize(matResultShow.width(), matResultShow.height())
106 | # self.pic.setScaledContents(True)
107 |
108 | matResultShow.setLayout(listLayout)
109 | self.resultdock = QDockWidget('Result Image', self)
110 | # self.resultdock.adjustSize()
111 | self.resultdock.setObjectName('result')
112 | self.resultdock.setWidget(matResultShow)
113 | self.resultdock.resize(150, 150)
114 |
115 | self.fileListWidget = QListWidget()
116 | self.fileListWidget.itemDoubleClicked.connect(
117 | self.fileitemDoubleClicked)
118 | fileListLayout = QVBoxLayout()
119 | fileListLayout.setContentsMargins(0, 0, 0, 0)
120 | fileListLayout.addWidget(self.fileListWidget)
121 | fileListContainer = QWidget()
122 | fileListContainer.setLayout(fileListLayout)
123 | self.filedock = QDockWidget('File List', self)
124 | self.filedock.setObjectName('Files')
125 | self.filedock.setWidget(fileListContainer)
126 |
127 | self.zoomWidget = ZoomWidget()
128 |
129 | self.canvas = Canvas(parent=self)
130 | scroll = QScrollArea()
131 | scroll.setWidget(self.canvas)
132 | scroll.setWidgetResizable(True)
133 | self.scrollBars = {
134 | Qt.Vertical: scroll.verticalScrollBar(),
135 | Qt.Horizontal: scroll.horizontalScrollBar()
136 | }
137 | self.scrollArea = scroll
138 | self.canvas.scrollRequest.connect(self.scrollRequest)
139 |
140 | self.setCentralWidget(scroll)
141 | self.addDockWidget(Qt.RightDockWidgetArea, self.resultdock)
142 | self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
143 | self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
144 |
145 | self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
146 | self.resultdock.setFeatures(
147 | self.resultdock.features() ^ self.dockFeatures)
148 |
149 | # Actions
150 | action = partial(newAction, self)
151 |
152 | open_file = action('&Open', self.openFile, 'Ctrl+O', 'Open image')
153 | open_dir = action('&Open Dir', self.openDir,
154 | 'Ctrl+D', 'Open image dir')
155 | change_save_dir = action('&Change Save Dir', self.changeSavedirDialog)
156 | # open_next_img = action('&Next Image', self.openNextImg,
157 | # 'Ctrl+N', 'Open next image')
158 | # open_pre_img = action('&Previous Image', self.openPreImg,
159 | # 'Ctrl+M', 'Open previous image')
160 | save = action('&Save', self.saveFile, 'Crl+S', 'Save output image')
161 | create = action('Create\nRectBox', self.createShape,
162 | 'w', 'Draw a new Box')
163 | matting = action('&Create\nMatting', self.grabcutMatting,
164 | 'e', 'GrabcutMatting')
165 |
166 | self.scalers = {
167 | self.FIT_WINDOW: self.scaleFitWindow,
168 | self.FIT_WIDTH: self.scaleFitWidth,
169 | # Set to one to scale to 100% when loading files.
170 | self.MANUAL_ZOOM: lambda: 1,
171 | }
172 |
173 | # store actions for further handling
174 | self.actions = struct(save=save, open_file=open_file,
175 | open_dir=open_dir, change_save_dir=change_save_dir,
176 | # open_next_img=open_next_img, open_pre_img=open_pre_img,
177 | create=create, matting=matting)
178 |
179 | # Auto saving: enable auto saving if pressing next
180 | # self.autoSaving = QAction('Auto Saving', self)
181 | # self.autoSaving.setCheckable(True)
182 | # self.autoSaving.setChecked()
183 |
184 | # set toolbar
185 | self.tools = self.toolbar('Tools')
186 | self.actions.all = (save, open_file, open_dir,
187 | change_save_dir, create,
188 | # open_pre_img, open_next_img,
189 | matting)
190 | addActions(self.tools, self.actions.all)
191 |
192 | # set status
193 | self.statusBar().showMessage('{} started.'.format(__appname__))
194 |
195 | def okToContinue(self):
196 | if self.dirty:
197 | reply = QMessageBox.question(self, "Attention",
198 | "you have unsaved changes, proceed anyway?",
199 | QMessageBox.Yes | QMessageBox.No | QMessageBox.Cancel)
200 | if reply == QMessageBox.Cancel:
201 | return False
202 | elif reply == QMessageBox.Yes:
203 | return self.fileSave
204 | return True
205 |
206 | def resetState(self):
207 | self.canvas.resetState()
208 |
209 | def errorMessage(self, title, message):
210 | return QMessageBox.critical(self, title,
211 | '
%s
%s' % (title, message)) 212 | 213 | def beginner(self): 214 | return self._beginner 215 | 216 | def advanced(self): 217 | return not self.beginner() 218 | 219 | def openFile(self, _value=False): 220 | path = os.path.dirname(self.filePath) if self.filePath else '.' 221 | formats = ['*.%s' % fmt.data().decode("ascii").lower() 222 | for fmt in QImageReader.supportedImageFormats()] 223 | filters = "Image (%s)" % ' '.join(formats) 224 | filename = QFileDialog.getOpenFileName( 225 | self, '%s - Choose Image or Label file' % __appname__, path, filters) 226 | if filename: 227 | if isinstance(filename, (tuple, list)): 228 | filename = filename[0] 229 | self.loadFile(filename) 230 | 231 | def openDir(self, dirpath=None): 232 | defaultOpenDirPath = dirpth if dirpath else '.' 233 | targetDirPath = QFileDialog.getExistingDirectory(self, 234 | '%s - Open Directory' % __appname__, defaultOpenDirPath, 235 | QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks) 236 | self.importDirImages(targetDirPath) 237 | 238 | def importDirImages(self, dirpath): 239 | self.fileListWidget.clear() 240 | self.mImgList = self.scanAllImages(dirpath) 241 | # self.openNextImg() 242 | for imgPath in self.mImgList: 243 | item = QListWidgetItem(imgPath) 244 | self.fileListWidget.addItem(item) 245 | 246 | def scanAllImages(self, folderPath): 247 | extensions = ['.%s' % fmt.data().decode("ascii").lower() 248 | for fmt in QImageReader.supportedImageFormats()] 249 | imageList = [] 250 | 251 | for root, dirs, files in os.walk(folderPath): 252 | for file in files: 253 | if file.lower().endswith(tuple(extensions)): 254 | relativePath = os.path.join(root, file) 255 | path = os.path.abspath(relativePath) 256 | imageList.append(path) 257 | imageList.sort(key=lambda x: x.lower()) 258 | return imageList 259 | 260 | def fileitemDoubleClicked(self, item=None): 261 | currIndex = self.mImgList.index(item.text()) 262 | if currIndex < len(self.mImgList): 263 | filename = self.mImgList[currIndex] 264 | if filename: 265 | self.loadFile(filename) 266 | 267 | def loadFile(self, filePath=None): 268 | self.resetState() 269 | self.canvas.setEnabled(False) 270 | 271 | # highlight the file item 272 | if filePath and self.fileListWidget.count() > 0: 273 | index = self.mImgList.index(filePath) 274 | fileWidgetItem = self.fileListWidget.item(index) 275 | fileWidgetItem.setSelected(True) 276 | 277 | if filePath and os.path.exists(filePath): 278 | # load image 279 | self.imageData = read(filePath, None) 280 | 281 | image = QImage.fromData(self.imageData) 282 | if image.isNull(): 283 | self.errorMessage(u'Error opening file', 284 | u'Make sure %s is a valid image file.' % filePath) 285 | self.status('Error reading %s' % filePath) 286 | return False 287 | self.status('Loaded %s' % os.path.basename(filePath)) 288 | self.image = image 289 | self.filePath = filePath 290 | self.canvas.loadPixmap(QPixmap.fromImage(image)) 291 | self.canvas.setEnabled(True) 292 | self.adjustScale(initial=True) 293 | self.paintCanvas() 294 | # self.toggleActions(True) 295 | 296 | def status(self, message, delay=5000): 297 | self.statusBar().showMessage(message, delay) 298 | 299 | def adjustScale(self, initial=False): 300 | value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]() 301 | self.zoomWidget.setValue(int(100 * value)) 302 | 303 | def scaleFitWindow(self): 304 | """Figure out the size of the pixmap in order to fit the main widget.""" 305 | e = 2.0 # So that no scrollbars are generated. 306 | w1 = self.centralWidget().width() - e 307 | h1 = self.centralWidget().height() - e 308 | a1 = w1 / h1 309 | # Calculate a new scale value based on the pixmap's aspect ratio. 310 | w2 = self.canvas.pixmap.width() - 0.0 311 | h2 = self.canvas.pixmap.height() - 0.0 312 | a2 = w2 / h2 313 | return w1 / w2 if a2 >= a1 else h1 / h2 314 | 315 | def scaleFitWidth(self): 316 | # The epsilon does not seem to work too well here. 317 | w = self.centralWidget().width() - 2.0 318 | return w / self.canvas.pixmap.width() 319 | 320 | def paintCanvas(self): 321 | assert not self.image.isNull(), "cannot paint null image" 322 | self.canvas.scale = 0.01 * self.zoomWidget.value() 323 | self.canvas.adjustSize() 324 | self.canvas.update() 325 | 326 | def createShape(self): 327 | assert self.beginner() 328 | self.canvas.setEditing(False) 329 | self.actions.create.setEnabled(False) 330 | 331 | def toggleDrawMode(self, edit=True): 332 | self.canvas.setEditing(edit) 333 | self.actions.createMode.setEnabled(edit) 334 | self.actions.editMode.setEnabled(not edit) 335 | 336 | def grabcutMatting(self): 337 | 338 | if self.mattingFile is None: 339 | self.mattingFile = Grab_cut() 340 | 341 | def format_shape(s): 342 | return dict(line_color=s.line_color.getRgb(), 343 | fill_color=s.fill_color.getRgb(), 344 | points=[(p.x(), p.y()) for p in s.points]) 345 | 346 | shape = format_shape(self.canvas.shapes[-1]) 347 | self.image_out_np = self.mattingFile.image_matting(self.filePath, 348 | shape, iteration=10) 349 | self.showResultImg(self.image_out_np) 350 | self.actions.save.setEnabled(True) 351 | 352 | def showResultImg(self, image_np): 353 | # resize to pic 354 | factor = min(self.pic.width() / 355 | image_np.shape[1], self.pic.height() / image_np.shape[0]) 356 | image_np = cv2.resize(image_np, None, fx=factor, 357 | fy=factor, interpolation=cv2.INTER_AREA) 358 | # image_np = cv2.resize((self.pic.height(), self.pic.width())) 359 | image = QImage(image_np, image_np.shape[1], 360 | image_np.shape[0], QImage.Format_ARGB32) 361 | matImg = QPixmap(image) 362 | self.pic.setPixmap(matImg) 363 | 364 | def saveFile(self): 365 | self._saveFile(self.saveFileDialog()) 366 | 367 | def _saveFile(self, saved_path): 368 | print(saved_path) 369 | if saved_path: 370 | Grab_cut.resultSave(saved_path, self.image_out_np) 371 | self.setClean() 372 | self.statusBar().showMessage('Saved to %s' % saved_path) 373 | self.statusBar().show() 374 | 375 | def saveFileDialog(self): 376 | caption = '%s - Choose File' % __appname__ 377 | filters = 'File (*%s)' % 'png' 378 | if self.default_save_dir is not None and len(self.default_save_dir): 379 | openDialogPath = self.default_save_dir 380 | else: 381 | openDialogPath = self.currentPath() 382 | 383 | print(openDialogPath) 384 | dlg = QFileDialog(self, caption, openDialogPath, filters) 385 | dlg.setDefaultSuffix('png') 386 | dlg.setAcceptMode(QFileDialog.AcceptSave) 387 | filenameWithoutExtension = os.path.splitext(self.filePath)[0] 388 | dlg.selectFile(filenameWithoutExtension) 389 | dlg.setOption(QFileDialog.DontUseNativeDialog, False) 390 | if dlg.exec_(): 391 | return dlg.selectedFiles()[0] 392 | return '' 393 | 394 | def currentPath(self): 395 | return os.path.dirname(self.filePath) if self.filePath else '.' 396 | 397 | def changeSavedirDialog(self, _value=False): 398 | if self.default_save_dir is not None: 399 | path = self.default_save_dir 400 | else: 401 | path = '.' 402 | 403 | dirpath = QFileDialog.getExistingDirectory(self, 404 | '%s - Save annotations to the directory' % __appname__, path, QFileDialog.ShowDirsOnly 405 | | QFileDialog.DontResolveSymlinks) 406 | 407 | if dirpath is not None and len(dirpath) > 1: 408 | self.default_save_dir = dirpath 409 | 410 | self.statusBar().showMessage('%s . Annotation will be saved to %s' % 411 | ('Change saved folder', self.default_save_dir)) 412 | self.statusBar().show() 413 | 414 | def setClean(self): 415 | self.dirty = False 416 | self.actions.save.setEnabled(False) 417 | 418 | self.actions.create.setEnabled(True) 419 | 420 | def openNextImg(): 421 | pass 422 | 423 | def openPreImg(): 424 | pass 425 | 426 | def scrollRequest(self, delta, orientation): 427 | units = - delta / (8 * 15) 428 | bar = self.scrollBars[orientation] 429 | bar.setValue(bar.value() + bar.singleStep() * units) 430 | 431 | 432 | def read(filename, default=None): 433 | try: 434 | with open(filename, 'rb') as f: 435 | return f.read() 436 | except Exception: 437 | return default 438 | 439 | 440 | def get_main_app(argv=[]): 441 | app = QApplication(argv) 442 | app.setApplicationName(__appname__) 443 | app.setWindowIcon(newIcon("app")) 444 | ex = MainWindow() 445 | ex.show() 446 | return app, ex 447 | 448 | 449 | def main(argv=[]): 450 | app, ex = get_main_app(argv) 451 | return app.exec_() 452 | 453 | 454 | if __name__ == '__main__': 455 | sys.exit(main(sys.argv)) 456 | -------------------------------------------------------------------------------- /canvas.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtGui import * 2 | from PyQt5.QtCore import * 3 | from PyQt5.QtWidgets import * 4 | 5 | from shape import Shape 6 | from lib import distance 7 | 8 | CURSOR_DEFAULT = Qt.ArrowCursor 9 | CURSOR_POINT = Qt.PointingHandCursor 10 | CURSOR_DRAW = Qt.CrossCursor 11 | CURSOR_MOVE = Qt.ClosedHandCursor 12 | CURSOR_GRAB = Qt.OpenHandCursor 13 | 14 | 15 | class Canvas(QWidget): 16 | zoomRequest = pyqtSignal(int) 17 | scrollRequest = pyqtSignal(int, int) 18 | newShape = pyqtSignal() 19 | selectionChanged = pyqtSignal(bool) 20 | shapeMoved = pyqtSignal() 21 | drawingPolygon = pyqtSignal(bool) 22 | 23 | CREATE, EDIT = list(range(2)) 24 | 25 | epsilon = 11.0 26 | 27 | def __init__(self, *args, **kwargs): 28 | super(Canvas, self).__init__(*args, **kwargs) 29 | # Initialise local state. 30 | self.mode = self.EDIT 31 | self.shapes = [] 32 | self.current = None 33 | self.selectedShape = None # save the selected shape here 34 | self.selectedShapeCopy = None 35 | self.drawingLineColor = QColor(0, 0, 255) 36 | self.drawingRectColor = QColor(0, 0, 255) 37 | self.line = Shape(line_color=self.drawingLineColor) 38 | self.prevPoint = QPointF() 39 | self.offsets = QPointF(), QPointF() 40 | self.scale = 1.0 41 | self.pixmap = QPixmap() 42 | self.visible = {} 43 | self._hideBackround = False 44 | self.hideBackround = False 45 | self.hShape = None 46 | self.hVertex = None 47 | self._painter = QPainter() 48 | self._cursor = CURSOR_DEFAULT 49 | # Menus: 50 | self.menus = (QMenu(), QMenu()) 51 | # Set widget options. 52 | self.setMouseTracking(True) 53 | self.setFocusPolicy(Qt.WheelFocus) 54 | self.verified = False 55 | 56 | def setDrawingColor(self, qColor): 57 | self.drawingLineColor = qColor 58 | self.drawingRectColor = qColor 59 | 60 | def enterEvent(self, ev): 61 | self.overrideCursor(self._cursor) 62 | 63 | def leaveEvent(self, ev): 64 | self.restoreCursor() 65 | 66 | def focusOutEvent(self, ev): 67 | self.restoreCursor() 68 | 69 | def isVisible(self, shape): 70 | return self.visible.get(shape, True) 71 | 72 | def drawing(self): 73 | return self.mode == self.CREATE 74 | 75 | def editing(self): 76 | return self.mode == self.EDIT 77 | 78 | def setEditing(self, value=True): 79 | self.mode = self.EDIT if value else self.CREATE 80 | if not value: # Create 81 | self.unHighlight() 82 | self.deSelectShape() 83 | self.prevPoint = QPointF() 84 | self.repaint() 85 | 86 | def unHighlight(self): 87 | if self.hShape: 88 | self.hShape.highlightClear() 89 | self.hVertex = self.hShape = None 90 | 91 | def selectedVertex(self): 92 | return self.hVertex is not None 93 | 94 | def mouseMoveEvent(self, ev): 95 | """Update line with last point and current coordinates.""" 96 | pos = self.transformPos(ev.pos()) 97 | 98 | # Polygon drawing. 99 | if self.drawing(): 100 | self.overrideCursor(CURSOR_DRAW) 101 | if self.current: 102 | color = self.drawingLineColor 103 | if self.outOfPixmap(pos): 104 | # Don't allow the user to draw outside the pixmap. 105 | # Project the point to the pixmap's edges. 106 | pos = self.intersectionPoint(self.current[-1], pos) 107 | elif len(self.current) > 1 and self.closeEnough(pos, self.current[0]): 108 | # Attract line to starting point and colorise to alert the 109 | # user: 110 | pos = self.current[0] 111 | color = self.current.line_color 112 | self.overrideCursor(CURSOR_POINT) 113 | self.current.highlightVertex(0, Shape.NEAR_VERTEX) 114 | self.line[1] = pos 115 | self.line.line_color = color 116 | self.prevPoint = QPointF() 117 | self.current.highlightClear() 118 | else: 119 | self.prevPoint = pos 120 | self.repaint() 121 | return 122 | 123 | # Polygon copy moving. 124 | if Qt.RightButton & ev.buttons(): 125 | if self.selectedShapeCopy and self.prevPoint: 126 | self.overrideCursor(CURSOR_MOVE) 127 | self.boundedMoveShape(self.selectedShapeCopy, pos) 128 | self.repaint() 129 | elif self.selectedShape: 130 | self.selectedShapeCopy = self.selectedShape.copy() 131 | self.repaint() 132 | return 133 | 134 | # Polygon/Vertex moving. 135 | if Qt.LeftButton & ev.buttons(): 136 | if self.selectedVertex(): 137 | self.boundedMoveVertex(pos) 138 | self.shapeMoved.emit() 139 | self.repaint() 140 | elif self.selectedShape and self.prevPoint: 141 | self.overrideCursor(CURSOR_MOVE) 142 | self.boundedMoveShape(self.selectedShape, pos) 143 | self.shapeMoved.emit() 144 | self.repaint() 145 | return 146 | 147 | # Just hovering over the canvas, 2 posibilities: 148 | # - Highlight shapes 149 | # - Highlight vertex 150 | # Update shape/vertex fill and tooltip value accordingly. 151 | self.setToolTip("Image") 152 | for shape in reversed([s for s in self.shapes if self.isVisible(s)]): 153 | # Look for a nearby vertex to highlight. If that fails, 154 | # check if we happen to be inside a shape. 155 | index = shape.nearestVertex(pos, self.epsilon) 156 | if index is not None: 157 | if self.selectedVertex(): 158 | self.hShape.highlightClear() 159 | self.hVertex, self.hShape = index, shape 160 | shape.highlightVertex(index, shape.MOVE_VERTEX) 161 | self.overrideCursor(CURSOR_POINT) 162 | self.setToolTip("Click & drag to move point") 163 | self.setStatusTip(self.toolTip()) 164 | self.update() 165 | break 166 | elif shape.containsPoint(pos): 167 | if self.selectedVertex(): 168 | self.hShape.highlightClear() 169 | self.hVertex, self.hShape = None, shape 170 | self.setToolTip( 171 | "Click & drag to move shape '%s'" % shape.label) 172 | self.setStatusTip(self.toolTip()) 173 | self.overrideCursor(CURSOR_GRAB) 174 | self.update() 175 | break 176 | else: # Nothing found, clear highlights, reset state. 177 | if self.hShape: 178 | self.hShape.highlightClear() 179 | self.update() 180 | self.hVertex, self.hShape = None, None 181 | self.overrideCursor(CURSOR_DEFAULT) 182 | 183 | def mousePressEvent(self, ev): 184 | pos = self.transformPos(ev.pos()) 185 | 186 | if ev.button() == Qt.LeftButton: 187 | if self.drawing(): 188 | self.handleDrawing(pos) 189 | else: 190 | self.selectShapePoint(pos) 191 | self.prevPoint = pos 192 | self.repaint() 193 | elif ev.button() == Qt.RightButton and self.editing(): 194 | self.selectShapePoint(pos) 195 | self.prevPoint = pos 196 | self.repaint() 197 | 198 | def mouseReleaseEvent(self, ev): 199 | if ev.button() == Qt.RightButton: 200 | menu = self.menus[bool(self.selectedShapeCopy)] 201 | self.restoreCursor() 202 | if not menu.exec_(self.mapToGlobal(ev.pos()))\ 203 | and self.selectedShapeCopy: 204 | # Cancel the move by deleting the shadow copy. 205 | self.selectedShapeCopy = None 206 | self.repaint() 207 | elif ev.button() == Qt.LeftButton and self.selectedShape: 208 | if self.selectedVertex(): 209 | self.overrideCursor(CURSOR_POINT) 210 | else: 211 | self.overrideCursor(CURSOR_GRAB) 212 | elif ev.button() == Qt.LeftButton: 213 | pos = self.transformPos(ev.pos()) 214 | if self.drawing(): 215 | self.handleDrawing(pos) 216 | 217 | def endMove(self, copy=False): 218 | assert self.selectedShape and self.selectedShapeCopy 219 | shape = self.selectedShapeCopy 220 | #del shape.fill_color 221 | #del shape.line_color 222 | if copy: 223 | self.shapes.append(shape) 224 | self.selectedShape.selected = False 225 | self.selectedShape = shape 226 | self.repaint() 227 | else: 228 | self.selectedShape.points = [p for p in shape.points] 229 | self.selectedShapeCopy = None 230 | 231 | def hideBackroundShapes(self, value): 232 | self.hideBackround = value 233 | if self.selectedShape: 234 | # Only hide other shapes if there is a current selection. 235 | # Otherwise the user will not be able to select a shape. 236 | self.setHiding(True) 237 | self.repaint() 238 | 239 | def handleDrawing(self, pos): 240 | if self.current and self.current.reachMaxPoints() is False: 241 | initPos = self.current[0] 242 | minX = initPos.x() 243 | minY = initPos.y() 244 | targetPos = self.line[1] 245 | maxX = targetPos.x() 246 | maxY = targetPos.y() 247 | self.current.addPoint(QPointF(maxX, minY)) 248 | self.current.addPoint(targetPos) 249 | self.current.addPoint(QPointF(minX, maxY)) 250 | self.finalise() 251 | elif not self.outOfPixmap(pos): 252 | self.current = Shape() 253 | self.current.addPoint(pos) 254 | self.line.points = [pos, pos] 255 | self.setHiding() 256 | self.drawingPolygon.emit(True) 257 | self.update() 258 | 259 | def setHiding(self, enable=True): 260 | self._hideBackround = self.hideBackround if enable else False 261 | 262 | def canCloseShape(self): 263 | return self.drawing() and self.current and len(self.current) > 2 264 | 265 | def mouseDoubleClickEvent(self, ev): 266 | # We need at least 4 points here, since the mousePress handler 267 | # adds an extra one before this handler is called. 268 | if self.canCloseShape() and len(self.current) > 3: 269 | self.current.popPoint() 270 | self.finalise() 271 | 272 | def selectShape(self, shape): 273 | self.deSelectShape() 274 | shape.selected = True 275 | self.selectedShape = shape 276 | self.setHiding() 277 | self.selectionChanged.emit(True) 278 | self.update() 279 | 280 | def selectShapePoint(self, point): 281 | """Select the first shape created which contains this point.""" 282 | self.deSelectShape() 283 | if self.selectedVertex(): # A vertex is marked for selection. 284 | index, shape = self.hVertex, self.hShape 285 | shape.highlightVertex(index, shape.MOVE_VERTEX) 286 | self.selectShape(shape) 287 | return 288 | for shape in reversed(self.shapes): 289 | if self.isVisible(shape) and shape.containsPoint(point): 290 | self.selectShape(shape) 291 | self.calculateOffsets(shape, point) 292 | return 293 | 294 | def calculateOffsets(self, shape, point): 295 | rect = shape.boundingRect() 296 | x1 = rect.x() - point.x() 297 | y1 = rect.y() - point.y() 298 | x2 = (rect.x() + rect.width()) - point.x() 299 | y2 = (rect.y() + rect.height()) - point.y() 300 | self.offsets = QPointF(x1, y1), QPointF(x2, y2) 301 | 302 | def boundedMoveVertex(self, pos): 303 | index, shape = self.hVertex, self.hShape 304 | point = shape[index] 305 | if self.outOfPixmap(pos): 306 | pos = self.intersectionPoint(point, pos) 307 | 308 | shiftPos = pos - point 309 | shape.moveVertexBy(index, shiftPos) 310 | 311 | lindex = (index + 1) % 4 312 | rindex = (index + 3) % 4 313 | lshift = None 314 | rshift = None 315 | if index % 2 == 0: 316 | rshift = QPointF(shiftPos.x(), 0) 317 | lshift = QPointF(0, shiftPos.y()) 318 | else: 319 | lshift = QPointF(shiftPos.x(), 0) 320 | rshift = QPointF(0, shiftPos.y()) 321 | shape.moveVertexBy(rindex, rshift) 322 | shape.moveVertexBy(lindex, lshift) 323 | 324 | def boundedMoveShape(self, shape, pos): 325 | if self.outOfPixmap(pos): 326 | return False # No need to move 327 | o1 = pos + self.offsets[0] 328 | if self.outOfPixmap(o1): 329 | pos -= QPointF(min(0, o1.x()), min(0, o1.y())) 330 | o2 = pos + self.offsets[1] 331 | if self.outOfPixmap(o2): 332 | pos += QPointF(min(0, self.pixmap.width() - o2.x()), 333 | min(0, self.pixmap.height() - o2.y())) 334 | # The next line tracks the new position of the cursor 335 | # relative to the shape, but also results in making it 336 | # a bit "shaky" when nearing the border and allows it to 337 | # go outside of the shape's area for some reason. XXX 338 | #self.calculateOffsets(self.selectedShape, pos) 339 | dp = pos - self.prevPoint 340 | if dp: 341 | shape.moveBy(dp) 342 | self.prevPoint = pos 343 | return True 344 | return False 345 | 346 | def deSelectShape(self): 347 | if self.selectedShape: 348 | self.selectedShape.selected = False 349 | self.selectedShape = None 350 | self.setHiding(False) 351 | self.selectionChanged.emit(False) 352 | self.update() 353 | 354 | def deleteSelected(self): 355 | if self.selectedShape: 356 | shape = self.selectedShape 357 | self.shapes.remove(self.selectedShape) 358 | self.selectedShape = None 359 | self.update() 360 | return shape 361 | 362 | def copySelectedShape(self): 363 | if self.selectedShape: 364 | shape = self.selectedShape.copy() 365 | self.deSelectShape() 366 | self.shapes.append(shape) 367 | shape.selected = True 368 | self.selectedShape = shape 369 | self.boundedShiftShape(shape) 370 | return shape 371 | 372 | def boundedShiftShape(self, shape): 373 | # Try to move in one direction, and if it fails in another. 374 | # Give up if both fail. 375 | point = shape[0] 376 | offset = QPointF(2.0, 2.0) 377 | self.calculateOffsets(shape, point) 378 | self.prevPoint = point 379 | if not self.boundedMoveShape(shape, point - offset): 380 | self.boundedMoveShape(shape, point + offset) 381 | 382 | def paintEvent(self, event): 383 | if not self.pixmap: 384 | return super(Canvas, self).paintEvent(event) 385 | 386 | p = self._painter 387 | p.begin(self) 388 | p.setRenderHint(QPainter.Antialiasing) 389 | p.setRenderHint(QPainter.HighQualityAntialiasing) 390 | p.setRenderHint(QPainter.SmoothPixmapTransform) 391 | 392 | p.scale(self.scale, self.scale) 393 | p.translate(self.offsetToCenter()) 394 | 395 | p.drawPixmap(0, 0, self.pixmap) 396 | Shape.scale = self.scale 397 | for shape in self.shapes: 398 | if (shape.selected or not self._hideBackround) and self.isVisible(shape): 399 | shape.fill = shape.selected or shape == self.hShape 400 | shape.paint(p) 401 | if self.current: 402 | self.current.paint(p) 403 | self.line.paint(p) 404 | if self.selectedShapeCopy: 405 | self.selectedShapeCopy.paint(p) 406 | 407 | # Paint rect 408 | if self.current is not None and len(self.line) == 2: 409 | leftTop = self.line[0] 410 | rightBottom = self.line[1] 411 | rectWidth = rightBottom.x() - leftTop.x() 412 | rectHeight = rightBottom.y() - leftTop.y() 413 | p.setPen(self.drawingRectColor) 414 | brush = QBrush(Qt.BDiagPattern) 415 | p.setBrush(brush) 416 | p.drawRect(leftTop.x(), leftTop.y(), rectWidth, rectHeight) 417 | 418 | if self.drawing() and not self.prevPoint.isNull() and not self.outOfPixmap(self.prevPoint): 419 | p.setPen(QColor(0, 0, 0)) 420 | p.drawLine(self.prevPoint.x(), 0, 421 | self.prevPoint.x(), self.pixmap.height()) 422 | p.drawLine(0, self.prevPoint.y(), 423 | self.pixmap.width(), self.prevPoint.y()) 424 | 425 | self.setAutoFillBackground(True) 426 | if self.verified: 427 | pal = self.palette() 428 | pal.setColor(self.backgroundRole(), QColor(184, 239, 38, 128)) 429 | self.setPalette(pal) 430 | else: 431 | pal = self.palette() 432 | pal.setColor(self.backgroundRole(), QColor(232, 232, 232, 255)) 433 | self.setPalette(pal) 434 | 435 | p.end() 436 | 437 | def transformPos(self, point): 438 | """Convert from widget-logical coordinates to painter-logical coordinates.""" 439 | return point / self.scale - self.offsetToCenter() 440 | 441 | def offsetToCenter(self): 442 | s = self.scale 443 | area = super(Canvas, self).size() 444 | w, h = self.pixmap.width() * s, self.pixmap.height() * s 445 | aw, ah = area.width(), area.height() 446 | x = (aw - w) / (2 * s) if aw > w else 0 447 | y = (ah - h) / (2 * s) if ah > h else 0 448 | return QPointF(x, y) 449 | 450 | def outOfPixmap(self, p): 451 | w, h = self.pixmap.width(), self.pixmap.height() 452 | return not (0 <= p.x() <= w and 0 <= p.y() <= h) 453 | 454 | def finalise(self): 455 | assert self.current 456 | if self.current.points[0] == self.current.points[-1]: 457 | self.current = None 458 | self.drawingPolygon.emit(False) 459 | self.update() 460 | return 461 | 462 | self.current.close() 463 | self.shapes.append(self.current) 464 | self.current = None 465 | self.setHiding(False) 466 | self.newShape.emit() 467 | self.update() 468 | 469 | def closeEnough(self, p1, p2): 470 | #d = distance(p1 - p2) 471 | #m = (p1-p2).manhattanLength() 472 | # print "d %.2f, m %d, %.2f" % (d, m, d - m) 473 | return distance(p1 - p2) < self.epsilon 474 | 475 | def intersectionPoint(self, p1, p2): 476 | # Cycle through each image edge in clockwise fashion, 477 | # and find the one intersecting the current line segment. 478 | # http://paulbourke.net/geometry/lineline2d/ 479 | size = self.pixmap.size() 480 | points = [(0, 0), 481 | (size.width(), 0), 482 | (size.width(), size.height()), 483 | (0, size.height())] 484 | x1, y1 = p1.x(), p1.y() 485 | x2, y2 = p2.x(), p2.y() 486 | d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points)) 487 | x3, y3 = points[i] 488 | x4, y4 = points[(i + 1) % 4] 489 | if (x, y) == (x1, y1): 490 | # Handle cases where previous point is on one of the edges. 491 | if x3 == x4: 492 | return QPointF(x3, min(max(0, y2), max(y3, y4))) 493 | else: # y3 == y4 494 | return QPointF(min(max(0, x2), max(x3, x4)), y3) 495 | return QPointF(x, y) 496 | 497 | def intersectingEdges(self, x1y1, x2y2, points): 498 | """For each edge formed by `points', yield the intersection 499 | with the line segment `(x1,y1) - (x2,y2)`, if it exists. 500 | Also return the distance of `(x2,y2)' to the middle of the 501 | edge along with its index, so that the one closest can be chosen.""" 502 | x1, y1 = x1y1 503 | x2, y2 = x2y2 504 | for i in range(4): 505 | x3, y3 = points[i] 506 | x4, y4 = points[(i + 1) % 4] 507 | denom = (y4 - y3) * (x2 - x1) - (x4 - x3) * (y2 - y1) 508 | nua = (x4 - x3) * (y1 - y3) - (y4 - y3) * (x1 - x3) 509 | nub = (x2 - x1) * (y1 - y3) - (y2 - y1) * (x1 - x3) 510 | if denom == 0: 511 | # This covers two cases: 512 | # nua == nub == 0: Coincident 513 | # otherwise: Parallel 514 | continue 515 | ua, ub = nua / denom, nub / denom 516 | if 0 <= ua <= 1 and 0 <= ub <= 1: 517 | x = x1 + ua * (x2 - x1) 518 | y = y1 + ua * (y2 - y1) 519 | m = QPointF((x3 + x4) / 2, (y3 + y4) / 2) 520 | d = distance(m - QPointF(x2, y2)) 521 | yield d, i, (x, y) 522 | 523 | # These two, along with a call to adjustSize are required for the 524 | # scroll area. 525 | def sizeHint(self): 526 | return self.minimumSizeHint() 527 | 528 | def minimumSizeHint(self): 529 | if self.pixmap: 530 | return self.scale * self.pixmap.size() 531 | return super(Canvas, self).minimumSizeHint() 532 | 533 | def wheelEvent(self, ev): 534 | qt_version = 4 if hasattr(ev, "delta") else 5 535 | if qt_version == 4: 536 | if ev.orientation() == Qt.Vertical: 537 | v_delta = ev.delta() 538 | h_delta = 0 539 | else: 540 | h_delta = ev.delta() 541 | v_delta = 0 542 | else: 543 | delta = ev.angleDelta() 544 | h_delta = delta.x() 545 | v_delta = delta.y() 546 | 547 | mods = ev.modifiers() 548 | if Qt.ControlModifier == int(mods) and v_delta: 549 | self.zoomRequest.emit(v_delta) 550 | else: 551 | v_delta and self.scrollRequest.emit(v_delta, Qt.Vertical) 552 | h_delta and self.scrollRequest.emit(h_delta, Qt.Horizontal) 553 | ev.accept() 554 | 555 | def keyPressEvent(self, ev): 556 | key = ev.key() 557 | if key == Qt.Key_Escape and self.current: 558 | print('ESC press') 559 | self.current = None 560 | self.drawingPolygon.emit(False) 561 | self.update() 562 | elif key == Qt.Key_Return and self.canCloseShape(): 563 | self.finalise() 564 | elif key == Qt.Key_Left and self.selectedShape: 565 | self.moveOnePixel('Left') 566 | elif key == Qt.Key_Right and self.selectedShape: 567 | self.moveOnePixel('Right') 568 | elif key == Qt.Key_Up and self.selectedShape: 569 | self.moveOnePixel('Up') 570 | elif key == Qt.Key_Down and self.selectedShape: 571 | self.moveOnePixel('Down') 572 | 573 | def moveOnePixel(self, direction): 574 | # print(self.selectedShape.points) 575 | if direction == 'Left' and not self.moveOutOfBound(QPointF(-1.0, 0)): 576 | # print("move Left one pixel") 577 | self.selectedShape.points[0] += QPointF(-1.0, 0) 578 | self.selectedShape.points[1] += QPointF(-1.0, 0) 579 | self.selectedShape.points[2] += QPointF(-1.0, 0) 580 | self.selectedShape.points[3] += QPointF(-1.0, 0) 581 | elif direction == 'Right' and not self.moveOutOfBound(QPointF(1.0, 0)): 582 | # print("move Right one pixel") 583 | self.selectedShape.points[0] += QPointF(1.0, 0) 584 | self.selectedShape.points[1] += QPointF(1.0, 0) 585 | self.selectedShape.points[2] += QPointF(1.0, 0) 586 | self.selectedShape.points[3] += QPointF(1.0, 0) 587 | elif direction == 'Up' and not self.moveOutOfBound(QPointF(0, -1.0)): 588 | # print("move Up one pixel") 589 | self.selectedShape.points[0] += QPointF(0, -1.0) 590 | self.selectedShape.points[1] += QPointF(0, -1.0) 591 | self.selectedShape.points[2] += QPointF(0, -1.0) 592 | self.selectedShape.points[3] += QPointF(0, -1.0) 593 | elif direction == 'Down' and not self.moveOutOfBound(QPointF(0, 1.0)): 594 | # print("move Down one pixel") 595 | self.selectedShape.points[0] += QPointF(0, 1.0) 596 | self.selectedShape.points[1] += QPointF(0, 1.0) 597 | self.selectedShape.points[2] += QPointF(0, 1.0) 598 | self.selectedShape.points[3] += QPointF(0, 1.0) 599 | self.shapeMoved.emit() 600 | self.repaint() 601 | 602 | def moveOutOfBound(self, step): 603 | points = [p1+p2 for p1, p2 in zip(self.selectedShape.points, [step]*4)] 604 | return True in map(self.outOfPixmap, points) 605 | 606 | def undoLastLine(self): 607 | assert self.shapes 608 | self.current = self.shapes.pop() 609 | self.current.setOpen() 610 | self.line.points = [self.current[-1], self.current[0]] 611 | self.drawingPolygon.emit(True) 612 | 613 | def resetAllLines(self): 614 | assert self.shapes 615 | self.current = self.shapes.pop() 616 | self.current.setOpen() 617 | self.line.points = [self.current[-1], self.current[0]] 618 | self.drawingPolygon.emit(True) 619 | self.current = None 620 | self.drawingPolygon.emit(False) 621 | self.update() 622 | 623 | def loadPixmap(self, pixmap): 624 | self.pixmap = pixmap 625 | self.shapes = [] 626 | self.repaint() 627 | 628 | def loadShapes(self, shapes): 629 | self.shapes = list(shapes) 630 | self.current = None 631 | self.repaint() 632 | 633 | def setShapeVisible(self, shape, value): 634 | self.visible[shape] = value 635 | self.repaint() 636 | 637 | def currentCursor(self): 638 | cursor = QApplication.overrideCursor() 639 | if cursor is not None: 640 | cursor = cursor.shape() 641 | return cursor 642 | 643 | def overrideCursor(self, cursor): 644 | self._cursor = cursor 645 | if self.currentCursor() is None: 646 | QApplication.setOverrideCursor(cursor) 647 | else: 648 | QApplication.changeOverrideCursor(cursor) 649 | 650 | def restoreCursor(self): 651 | QApplication.restoreOverrideCursor() 652 | 653 | def resetState(self): 654 | self.restoreCursor() 655 | self.pixmap = None 656 | self.update() 657 | --------------------------------------------------------------------------------