├── 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 | ![Performence of app](https://github.com/zihuaweng/image-matting-opencv-qt/blob/master/results.gif) 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 | --------------------------------------------------------------------------------