浏览代码

Support PyQt4 and PyQt5

Kentaro Wada 7 年之前
父节点
当前提交
dc2dce6f8b
共有 11 个文件被更改,包括 235 次插入154 次删除
  1. 99 97
      labelme/app.py
  2. 41 15
      labelme/canvas.py
  3. 9 2
      labelme/colorDialog.py
  4. 20 6
      labelme/labelDialog.py
  5. 2 2
      labelme/labelFile.py
  6. 7 2
      labelme/lib.py
  7. 10 5
      labelme/shape.py
  8. 8 2
      labelme/toolBar.py
  9. 9 6
      labelme/utils.py
  10. 9 3
      labelme/zoomWidget.py
  11. 21 14
      setup.py

+ 99 - 97
labelme/app.py

@@ -28,8 +28,13 @@ import subprocess
 from functools import partial
 from collections import defaultdict
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
 
 from labelme import resources
 from labelme.lib import struct, newAction, newIcon, addActions, fmtShortcut
@@ -71,7 +76,7 @@ class WindowMixin(object):
 
     def toolbar(self, title, actions=None):
         toolbar = ToolBar(title)
-        toolbar.setObjectName(u'%sToolBar' % title)
+        toolbar.setObjectName('%sToolBar' % title)
         #toolbar.setOrientation(Qt.Vertical)
         toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
         if actions:
@@ -81,7 +86,7 @@ class WindowMixin(object):
 
 
 class MainWindow(QMainWindow, WindowMixin):
-    FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = range(3)
+    FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = 0, 1, 2
 
     def __init__(self, filename=None, output=None):
         super(MainWindow, self).__init__()
@@ -99,8 +104,7 @@ class MainWindow(QMainWindow, WindowMixin):
         self.labelDialog = LabelDialog(parent=self)
 
         self.labelList = QListWidget()
-        self.itemsToShapes = {}
-        self.shapesToItems = {}
+        self.itemsToShapes = []
 
         self.labelList.itemActivated.connect(self.labelSelectionChanged)
         self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
@@ -119,8 +123,8 @@ class MainWindow(QMainWindow, WindowMixin):
         listLayout.addWidget(self.labelList)
 
 
-        self.dock = QDockWidget(u'Polygon Labels', self)
-        self.dock.setObjectName(u'Labels')
+        self.dock = QDockWidget('Polygon Labels', self)
+        self.dock.setObjectName('Labels')
         self.dock.setWidget(self.labelListContainer)
 
         self.zoomWidget = ZoomWidget()
@@ -152,67 +156,67 @@ class MainWindow(QMainWindow, WindowMixin):
         # Actions
         action = partial(newAction, self)
         quit = action('&Quit', self.close,
-                'Ctrl+Q', 'quit', u'Quit application')
+                'Ctrl+Q', 'quit', 'Quit application')
         open = action('&Open', self.openFile,
-                'Ctrl+O', 'open', u'Open image or label file')
+                'Ctrl+O', 'open', 'Open image or label file')
         save = action('&Save', self.saveFile,
-                'Ctrl+S', 'save', u'Save labels to file', enabled=False)
+                'Ctrl+S', 'save', 'Save labels to file', enabled=False)
         saveAs = action('&Save As', self.saveFileAs,
-                'Ctrl+Shift+S', 'save-as', u'Save labels to a different file',
+                'Ctrl+Shift+S', 'save-as', 'Save labels to a different file',
                 enabled=False)
         close = action('&Close', self.closeFile,
-                'Ctrl+W', 'close', u'Close current file')
+                'Ctrl+W', 'close', 'Close current file')
         color1 = action('Polygon &Line Color', self.chooseColor1,
-                'Ctrl+L', 'color_line', u'Choose polygon line color')
+                'Ctrl+L', 'color_line', 'Choose polygon line color')
         color2 = action('Polygon &Fill Color', self.chooseColor2,
-                'Ctrl+Shift+L', 'color', u'Choose polygon fill color')
+                'Ctrl+Shift+L', 'color', 'Choose polygon fill color')
 
         createMode = action('Create\nPolygo&ns', self.setCreateMode,
-                'Ctrl+N', 'new', u'Start drawing polygons', enabled=False)
+                'Ctrl+N', 'new', 'Start drawing polygons', enabled=False)
         editMode = action('&Edit\nPolygons', self.setEditMode,
-                'Ctrl+J', 'edit', u'Move and edit polygons', enabled=False)
+                'Ctrl+J', 'edit', 'Move and edit polygons', enabled=False)
 
         create = action('Create\nPolygo&n', self.createShape,
-                'Ctrl+N', 'new', u'Draw a new polygon', enabled=False)
+                'Ctrl+N', 'new', 'Draw a new polygon', enabled=False)
         delete = action('Delete\nPolygon', self.deleteSelectedShape,
-                'Delete', 'delete', u'Delete', enabled=False)
+                'Delete', 'delete', 'Delete', enabled=False)
         copy = action('&Duplicate\nPolygon', self.copySelectedShape,
-                'Ctrl+D', 'copy', u'Create a duplicate of the selected polygon',
+                'Ctrl+D', 'copy', 'Create a duplicate of the selected polygon',
                 enabled=False)
 
         advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
-                'Ctrl+Shift+A', 'expert', u'Switch to advanced mode',
+                'Ctrl+Shift+A', 'expert', 'Switch to advanced mode',
                 checkable=True)
 
         hideAll = action('&Hide\nPolygons', partial(self.togglePolygons, False),
-                'Ctrl+H', 'hide', u'Hide all polygons',
+                'Ctrl+H', 'hide', 'Hide all polygons',
                 enabled=False)
         showAll = action('&Show\nPolygons', partial(self.togglePolygons, True),
-                'Ctrl+A', 'hide', u'Show all polygons',
+                'Ctrl+A', 'hide', 'Show all polygons',
                 enabled=False)
 
         help = action('&Tutorial', self.tutorial, 'Ctrl+T', 'help',
-                u'Show screencast of introductory tutorial')
+                'Show screencast of introductory tutorial')
 
         zoom = QWidgetAction(self)
         zoom.setDefaultWidget(self.zoomWidget)
         self.zoomWidget.setWhatsThis(
-            u"Zoom in or out of the image. Also accessible with"\
+            "Zoom in or out of the image. Also accessible with"\
              " %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
                  fmtShortcut("Ctrl+Wheel")))
         self.zoomWidget.setEnabled(False)
 
         zoomIn = action('Zoom &In', partial(self.addZoom, 10),
-                'Ctrl++', 'zoom-in', u'Increase zoom level', enabled=False)
+                'Ctrl++', 'zoom-in', 'Increase zoom level', enabled=False)
         zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
-                'Ctrl+-', 'zoom-out', u'Decrease zoom level', enabled=False)
+                'Ctrl+-', 'zoom-out', 'Decrease zoom level', enabled=False)
         zoomOrg = action('&Original size', partial(self.setZoom, 100),
-                'Ctrl+=', 'zoom', u'Zoom to original size', enabled=False)
+                'Ctrl+=', 'zoom', 'Zoom to original size', enabled=False)
         fitWindow = action('&Fit Window', self.setFitWindow,
-                'Ctrl+F', 'fit-window', u'Zoom follows window size',
+                'Ctrl+F', 'fit-window', 'Zoom follows window size',
                 checkable=True, enabled=False)
         fitWidth = action('Fit &Width', self.setFitWidth,
-                'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
+                'Ctrl+Shift+F', 'fit-width', 'Zoom follows window width',
                 checkable=True, enabled=False)
         # Group zoom controls into a list for easier toggling.
         zoomActions = (self.zoomWidget, zoomIn, zoomOut, zoomOrg, fitWindow, fitWidth)
@@ -225,15 +229,15 @@ class MainWindow(QMainWindow, WindowMixin):
         }
 
         edit = action('&Edit Label', self.editLabel,
-                'Ctrl+E', 'edit', u'Modify the label of the selected polygon',
+                'Ctrl+E', 'edit', 'Modify the label of the selected polygon',
                 enabled=False)
         self.editButton.setDefaultAction(edit)
 
         shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
-                icon='color_line', tip=u'Change the line color for this specific shape',
+                icon='color_line', tip='Change the line color for this specific shape',
                 enabled=False)
         shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
-                icon='color', tip=u'Change the fill color for this specific shape',
+                icon='color', tip='Change the fill color for this specific shape',
                 enabled=False)
 
         labels = self.dock.toggleViewAction()
@@ -316,30 +320,21 @@ class MainWindow(QMainWindow, WindowMixin):
 
         # XXX: Could be completely declarative.
         # Restore application settings.
-        types = {
-            'filename': QString,
-            'recentFiles': QStringList,
-            'window/size': QSize,
-            'window/position': QPoint,
-            'window/geometry': QByteArray,
-            # Docks and toolbars:
-            'window/state': QByteArray,
-        }
-        self.settings = settings = Settings(types)
-        self.recentFiles = list(settings['recentFiles'])
-        size = settings.get('window/size', QSize(600, 500))
-        position = settings.get('window/position', QPoint(0, 0))
+        self.settings = {}
+        self.recentFiles = self.settings.get('recentFiles', [])
+        size = self.settings.get('window/size', QSize(600, 500))
+        position = self.settings.get('window/position', QPoint(0, 0))
         self.resize(size)
         self.move(position)
         # or simply:
         #self.restoreGeometry(settings['window/geometry']
-        self.restoreState(settings['window/state'])
-        self.lineColor = QColor(settings.get('line/color', Shape.line_color))
-        self.fillColor = QColor(settings.get('fill/color', Shape.fill_color))
+        self.restoreState(self.settings.get('window/state', QByteArray()))
+        self.lineColor = QColor(self.settings.get('line/color', Shape.line_color))
+        self.fillColor = QColor(self.settings.get('fill/color', Shape.fill_color))
         Shape.line_color = self.lineColor
         Shape.fill_color = self.fillColor
 
-        if settings.get('advanced', QVariant()).toBool():
+        if self.settings.get('advanced', QVariant()):
             self.actions.advancedMode.setChecked(True)
             self.toggleAdvancedMode()
 
@@ -419,8 +414,7 @@ class MainWindow(QMainWindow, WindowMixin):
         self.statusBar().showMessage(message, delay)
 
     def resetState(self):
-        self.itemsToShapes.clear()
-        self.shapesToItems.clear()
+        self.itemsToShapes = []
         self.labelList.clear()
         self.filename = None
         self.imageData = None
@@ -480,7 +474,7 @@ class MainWindow(QMainWindow, WindowMixin):
     def updateFileMenu(self):
         current = self.filename
         def exists(filename):
-            return os.path.exists(unicode(filename))
+            return os.path.exists(str(filename))
         menu = self.menus.recentFiles
         menu.clear()
         files = [f for f in self.recentFiles if f != current and exists(f)]
@@ -510,7 +504,10 @@ class MainWindow(QMainWindow, WindowMixin):
         else:
             shape = self.canvas.selectedShape
             if shape:
-                self.labelList.setItemSelected(self.shapesToItems[shape], True)
+                for item, shape_ in self.itemsToShapes:
+                    if shape_ == shape:
+                        break
+                item.setSelected(True)
             else:
                 self.labelList.clearSelection()
         self.actions.delete.setEnabled(selected)
@@ -523,17 +520,17 @@ class MainWindow(QMainWindow, WindowMixin):
         item = QListWidgetItem(shape.label)
         item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
         item.setCheckState(Qt.Checked)
-        self.itemsToShapes[item] = shape
-        self.shapesToItems[shape] = item
+        self.itemsToShapes.append((item, shape))
         self.labelList.addItem(item)
         for action in self.actions.onShapesPresent:
             action.setEnabled(True)
 
     def remLabel(self, shape):
-        item = self.shapesToItems[shape]
+        for index, (item, shape_) in enumerate(self.itemsToShapes):
+            if shape_ == shape:
+                break
+        self.itemsToShapes.pop(index)
         self.labelList.takeItem(self.labelList.row(item))
-        del self.shapesToItems[shape]
-        del self.itemsToShapes[item]
 
     def loadLabels(self, shapes):
         s = []
@@ -553,7 +550,7 @@ class MainWindow(QMainWindow, WindowMixin):
     def saveLabels(self, filename):
         lf = LabelFile()
         def format_shape(s):
-            return dict(label=unicode(s.label),
+            return dict(label=str(s.label),
                         line_color=s.line_color.getRgb()\
                                 if s.line_color != self.lineColor else None,
                         fill_color=s.fill_color.getRgb()\
@@ -562,14 +559,14 @@ class MainWindow(QMainWindow, WindowMixin):
 
         shapes = [format_shape(shape) for shape in self.canvas.shapes]
         try:
-            lf.save(filename, shapes, unicode(self.filename), self.imageData,
+            lf.save(filename, shapes, str(self.filename), self.imageData,
                 self.lineColor.getRgb(), self.fillColor.getRgb())
             self.labelFile = lf
             self.filename = filename
             return True
-        except LabelFileError, e:
-            self.errorMessage(u'Error saving label data',
-                    u'<b>%s</b>' % e)
+        except LabelFileError as e:
+            self.errorMessage('Error saving label data',
+                    '<b>%s</b>' % e)
             return False
 
     def copySelectedShape(self):
@@ -581,13 +578,18 @@ class MainWindow(QMainWindow, WindowMixin):
         item = self.currentItem()
         if item and self.canvas.editing():
             self._noSelectionSlot = True
-            self.canvas.selectShape(self.itemsToShapes[item])
+            for item_, shape in self.itemsToShapes:
+                if item_ == item:
+                    break
+            self.canvas.selectShape(shape)
 
     def labelItemChanged(self, item):
-        shape = self.itemsToShapes[item]
-        label = unicode(item.text())
+        for item_, shape in self.itemsToShapes:
+            if item_ == item:
+                break
+        label = str(item.text())
         if label != shape.label:
-            shape.label = unicode(item.text())
+            shape.label = str(item.text())
             self.setDirty()
         else: # User probably changed item visibility
             self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
@@ -642,7 +644,7 @@ class MainWindow(QMainWindow, WindowMixin):
         self.adjustScale()
 
     def togglePolygons(self, value):
-        for item, shape in self.itemsToShapes.iteritems():
+        for item, shape in self.itemsToShapes:
             item.setCheckState(Qt.Checked if value else Qt.Unchecked)
 
     def loadFile(self, filename=None):
@@ -650,16 +652,16 @@ class MainWindow(QMainWindow, WindowMixin):
         self.resetState()
         self.canvas.setEnabled(False)
         if filename is None:
-            filename = self.settings['filename']
-        filename = unicode(filename)
+            filename = self.settings.get('filename', '')
+        filename = str(filename)
         if QFile.exists(filename):
             if LabelFile.isLabelFile(filename):
                 try:
                     self.labelFile = LabelFile(filename)
-                except LabelFileError, e:
-                    self.errorMessage(u'Error opening file',
-                            (u"<p><b>%s</b></p>"
-                             u"<p>Make sure <i>%s</i> is a valid label file.")\
+                except LabelFileError as e:
+                    self.errorMessage('Error opening file',
+                            ("<p><b>%s</b></p>"
+                             "<p>Make sure <i>%s</i> is a valid label file.")\
                             % (e, filename))
                     self.status("Error reading %s" % filename)
                     return False
@@ -673,11 +675,11 @@ class MainWindow(QMainWindow, WindowMixin):
                 self.labelFile = None
             image = QImage.fromData(self.imageData)
             if image.isNull():
-                self.errorMessage(u'Error opening file',
-                        u"<p>Make sure <i>%s</i> is a valid image file." % filename)
+                self.errorMessage('Error opening file',
+                        "<p>Make sure <i>%s</i> is a valid image file." % filename)
                 self.status("Error reading %s" % filename)
                 return False
-            self.status("Loaded %s" % os.path.basename(unicode(filename)))
+            self.status("Loaded %s" % os.path.basename(str(filename)))
             self.image = image
             self.filename = filename
             self.canvas.loadPixmap(QPixmap.fromImage(image))
@@ -729,7 +731,7 @@ class MainWindow(QMainWindow, WindowMixin):
         if not self.mayContinue():
             event.ignore()
         s = self.settings
-        s['filename'] = self.filename if self.filename else QString()
+        s['filename'] = self.filename if self.filename else ''
         s['window/size'] = self.size()
         s['window/position'] = self.pos()
         s['window/state'] = self.saveState()
@@ -749,14 +751,14 @@ class MainWindow(QMainWindow, WindowMixin):
     def openFile(self, _value=False):
         if not self.mayContinue():
             return
-        path = os.path.dirname(unicode(self.filename))\
+        path = os.path.dirname(str(self.filename))\
                 if self.filename else '.'
-        formats = ['*.%s' % unicode(fmt).lower()\
+        formats = ['*.%s' % str(fmt).lower()\
                 for fmt in QImageReader.supportedImageFormats()]
         filters = "Image & Label files (%s)" % \
                 ' '.join(formats + ['*%s' % LabelFile.suffix])
-        filename = unicode(QFileDialog.getOpenFileName(self,
-            '%s - Choose Image or Label file' % __appname__, path, filters))
+        filename = str(QFileDialog.getOpenFileName(self,
+            '%s - Choose Image or Label file' % __appname__, path, filters)[0])
         if filename:
             self.loadFile(filename)
 
@@ -781,15 +783,15 @@ class MainWindow(QMainWindow, WindowMixin):
         dlg = QFileDialog(self, caption, self.currentPath(), filters)
         dlg.setDefaultSuffix(LabelFile.suffix[1:])
         dlg.setAcceptMode(QFileDialog.AcceptSave)
-        dlg.setConfirmOverwrite(True)
+        dlg.setOption(QFileDialog.DontConfirmOverwrite, False)
         dlg.setOption(QFileDialog.DontUseNativeDialog, False)
         basename = os.path.splitext(self.filename)[0]
         default_labelfile_name = os.path.join(self.currentPath(),
                                               basename + LabelFile.suffix)
         filename = dlg.getSaveFileName(
             self, 'Choose File', default_labelfile_name,
-            'Label files (*%s)' % LabelFile.suffix)
-        return unicode(filename)
+            'Label files (*%s)' % LabelFile.suffix)[0]
+        return str(filename)
 
     def _saveFile(self, filename):
         if filename and self.saveLabels(filename):
@@ -810,8 +812,8 @@ class MainWindow(QMainWindow, WindowMixin):
     # Message Dialogs. #
     def hasLabels(self):
         if not self.itemsToShapes:
-            self.errorMessage(u'No objects labeled',
-                    u'You must label at least one object to save the file.')
+            self.errorMessage('No objects labeled',
+                    'You must label at least one object to save the file.')
             return False
         return True
 
@@ -820,18 +822,18 @@ class MainWindow(QMainWindow, WindowMixin):
 
     def discardChangesDialog(self):
         yes, no = QMessageBox.Yes, QMessageBox.No
-        msg = u'You have unsaved changes, proceed anyway?'
-        return yes == QMessageBox.warning(self, u'Attention', msg, yes|no)
+        msg = 'You have unsaved changes, proceed anyway?'
+        return yes == QMessageBox.warning(self, 'Attention', msg, yes|no)
 
     def errorMessage(self, title, message):
         return QMessageBox.critical(self, title,
                 '<p><b>%s</b></p>%s' % (title, message))
 
     def currentPath(self):
-        return os.path.dirname(unicode(self.filename)) if self.filename else '.'
+        return os.path.dirname(str(self.filename)) if self.filename else '.'
 
     def chooseColor1(self):
-        color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
+        color = self.colorDialog.getColor(self.lineColor, 'Choose line color',
                 default=DEFAULT_LINE_COLOR)
         if color:
             self.lineColor = color
@@ -841,7 +843,7 @@ class MainWindow(QMainWindow, WindowMixin):
             self.setDirty()
 
     def chooseColor2(self):
-       color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
+       color = self.colorDialog.getColor(self.fillColor, 'Choose fill color',
                 default=DEFAULT_FILL_COLOR)
        if color:
             self.fillColor = color
@@ -851,8 +853,8 @@ class MainWindow(QMainWindow, WindowMixin):
 
     def deleteSelectedShape(self):
         yes, no = QMessageBox.Yes, QMessageBox.No
-        msg = u'You are about to permanently delete this polygon, proceed anyway?'
-        if yes == QMessageBox.warning(self, u'Attention', msg, yes|no):
+        msg = 'You are about to permanently delete this polygon, proceed anyway?'
+        if yes == QMessageBox.warning(self, 'Attention', msg, yes|no):
             self.remLabel(self.canvas.deleteSelected())
             self.setDirty()
             if self.noShapes():
@@ -860,7 +862,7 @@ class MainWindow(QMainWindow, WindowMixin):
                     action.setEnabled(False)
 
     def chshapeLineColor(self):
-        color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
+        color = self.colorDialog.getColor(self.lineColor, 'Choose line color',
                 default=DEFAULT_LINE_COLOR)
         if color:
             self.canvas.selectedShape.line_color = color
@@ -868,7 +870,7 @@ class MainWindow(QMainWindow, WindowMixin):
             self.setDirty()
 
     def chshapeFillColor(self):
-        color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
+        color = self.colorDialog.getColor(self.fillColor, 'Choose fill color',
                 default=DEFAULT_FILL_COLOR)
         if color:
             self.canvas.selectedShape.fill_color = color

+ 41 - 15
labelme/canvas.py

@@ -17,9 +17,15 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
-#from PyQt4.QtOpenGL import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+    PYQT5 = True
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
+    PYQT5 = False
 
 from labelme.shape import Shape
 from labelme.lib import distance
@@ -42,7 +48,7 @@ class Canvas(QWidget):
     shapeMoved = pyqtSignal()
     drawingPolygon = pyqtSignal(bool)
 
-    CREATE, EDIT = range(2)
+    CREATE, EDIT = 0, 1
 
     epsilon = 11.0
 
@@ -107,7 +113,10 @@ class Canvas(QWidget):
 
     def mouseMoveEvent(self, ev):
         """Update line with last point and current coordinates."""
-        pos = self.transformPos(ev.posF())
+        if PYQT5:
+            pos = self.transformPos(ev.pos())
+        else:
+            pos = self.transformPos(ev.posF())
 
         self.restoreCursor()
 
@@ -191,7 +200,10 @@ class Canvas(QWidget):
             self.hVertex, self.hShape = None, None
 
     def mousePressEvent(self, ev):
-        pos = self.transformPos(ev.posF())
+        if PYQT5:
+            pos = self.transformPos(ev.pos())
+        else:
+            pos = self.transformPos(ev.posF())
         if ev.button() == Qt.LeftButton:
             if self.drawing():
                 if self.current:
@@ -442,12 +454,14 @@ class Canvas(QWidget):
                 return QPointF(min(max(0, x2), max(x3, x4)), y3)
         return QPointF(x, y)
 
-    def intersectingEdges(self, (x1, y1), (x2, y2), points):
+    def intersectingEdges(self, point1, point2, points):
         """For each edge formed by `points', yield the intersection
         with the line segment `(x1,y1) - (x2,y2)`, if it exists.
         Also return the distance of `(x2,y2)' to the middle of the
         edge along with its index, so that the one closest can be chosen."""
-        for i in xrange(4):
+        (x1, y1) = point1
+        (x2, y2) = point2
+        for i in range(4):
             x3, y3 = points[i]
             x4, y4 = points[(i+1) % 4]
             denom = (y4-y3) * (x2 - x1) - (x4 - x3) * (y2 - y1)
@@ -477,16 +491,28 @@ class Canvas(QWidget):
         return super(Canvas, self).minimumSizeHint()
 
     def wheelEvent(self, ev):
-        if ev.orientation() == Qt.Vertical:
-            mods = ev.modifiers()
-            if Qt.ControlModifier == int(mods):
-                self.zoomRequest.emit(ev.delta())
+        if PYQT5:
+            if ev.inverted():
+                mods = ev.modifiers()
+                if Qt.ControlModifier == int(mods):
+                    self.zoomRequest.emit(ev.pixelDelta())
+                else:
+                    self.scrollRequest.emit(ev.pixelDelta(),
+                            Qt.Horizontal if (Qt.ShiftModifier == int(mods))\
+                                        else Qt.Vertical)
             else:
-                self.scrollRequest.emit(ev.delta(),
+                self.scrollRequest.emit(ev.pixelDelta(), Qt.Horizontal)
+        else:
+            if ev.orientation() == qt.Vertical:
+                mods = ev.modifiers()
+                if Qt.ControlModifier == int(mods):
+                    self.zoomRequest.emit(ev.delta())
+                else:
+                    self.scrollRequest.emit(ev.delta(),
                         Qt.Horizontal if (Qt.ShiftModifier == int(mods))\
                                       else Qt.Vertical)
-        else:
-            self.scrollRequest.emit(ev.delta(), Qt.Horizontal)
+            else:
+                self.scrollRequest.emit(ev.delta(), Qt.Horizontal)
         ev.accept()
 
     def keyPressEvent(self, ev):

+ 9 - 2
labelme/colorDialog.py

@@ -17,8 +17,15 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
+    from PyQt4.QtWidgets import *
+
 
 BB = QDialogButtonBox
 

+ 20 - 6
labelme/labelDialog.py

@@ -17,10 +17,17 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+    PYQT5 = True
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
+    PYQT5 = False
 
-from lib import newIcon, labelValidator
+from .lib import newIcon, labelValidator
 
 # TODO:
 # - Calculate optimal position so as not to go out of screen area.
@@ -46,11 +53,18 @@ class LabelDialog(QDialog):
         self.setLayout(layout)
 
     def validate(self):
-        if self.edit.text().trimmed():
-            self.accept()
+        if PYQT5:
+            if self.edit.text().strip():
+                self.accept()
+        else:
+            if self.edit.text().trimmed():
+                self.accept()
 
     def postProcess(self):
-        self.edit.setText(self.edit.text().trimmed())
+        if PYQT5:
+            self.edit.setText(self.edit.text().strip())
+        else:
+            self.edit.setText(self.edit.text().trimmed())
 
     def popUp(self, text='', move=True):
         self.edit.setText(text)

+ 2 - 2
labelme/labelFile.py

@@ -51,7 +51,7 @@ class LabelFile(object):
                 self.imageData = imageData
                 self.lineColor = lineColor
                 self.fillColor = fillColor
-        except Exception, e:
+        except Exception as e:
             raise LabelFileError(e)
 
     def save(self, filename, shapes, imagePath, imageData,
@@ -64,7 +64,7 @@ class LabelFile(object):
                     imagePath=imagePath,
                     imageData=b64encode(imageData)),
                     f, ensure_ascii=True, indent=2)
-        except Exception, e:
+        except Exception as e:
             raise LabelFileError(e)
 
     @staticmethod

+ 7 - 2
labelme/lib.py

@@ -19,8 +19,13 @@
 
 from math import sqrt
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
 
 
 def newIcon(icon):

+ 10 - 5
labelme/shape.py

@@ -19,10 +19,15 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
 
-from lib import distance
+
+from .lib import distance
 
 # TODO:
 # - [opt] Store paths instead of creating new ones at each paint.
@@ -35,9 +40,9 @@ DEFAULT_VERTEX_FILL_COLOR = QColor(0, 255, 0, 255)
 DEFAULT_HVERTEX_FILL_COLOR = QColor(255, 0, 0)
 
 class Shape(object):
-    P_SQUARE, P_ROUND = range(2)
+    P_SQUARE, P_ROUND = 0, 1
 
-    MOVE_VERTEX, NEAR_VERTEX = range(2)
+    MOVE_VERTEX, NEAR_VERTEX = 0, 1
 
     ## The following class variables influence the drawing
     ## of _all_ shape objects.

+ 8 - 2
labelme/toolBar.py

@@ -17,8 +17,14 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
+
 
 class ToolBar(QToolBar):
     def __init__(self, title):

+ 9 - 6
labelme/utils.py

@@ -1,5 +1,8 @@
 import base64
-import cStringIO as StringIO
+try:
+    import io
+except ImportError:
+    import io as io
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -15,10 +18,10 @@ def labelcolormap(N=256):
         return ((byteval & (1 << idx)) != 0)
 
     cmap = np.zeros((N, 3))
-    for i in xrange(0, N):
+    for i in range(0, N):
         id = i
         r, g, b = 0, 0, 0
-        for j in xrange(0, 8):
+        for j in range(0, 8):
             r = np.bitwise_or(r, (bitget(id, 0) << 7-j))
             g = np.bitwise_or(g, (bitget(id, 1) << 7-j))
             b = np.bitwise_or(b, (bitget(id, 2) << 7-j))
@@ -31,7 +34,7 @@ def labelcolormap(N=256):
 
 
 def img_b64_to_array(img_b64):
-    f = StringIO.StringIO()
+    f = io.BytesIO()
     f.write(base64.b64decode(img_b64))
     img_arr = np.array(PIL.Image.open(f))
     return img_arr
@@ -40,7 +43,7 @@ def img_b64_to_array(img_b64):
 def polygons_to_mask(img_shape, polygons):
     mask = np.zeros(img_shape[:2], dtype=np.uint8)
     mask = PIL.Image.fromarray(mask)
-    xy = map(tuple, polygons)
+    xy = list(map(tuple, polygons))
     PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
     mask = np.array(mask, dtype=bool)
     return mask
@@ -70,7 +73,7 @@ def draw_label(label, img, label_names, colormap=None):
         plt_titles.append(label_name)
     plt.legend(plt_handlers, plt_titles, loc='lower right', framealpha=.5)
 
-    f = StringIO.StringIO()
+    f = io.BytesIO()
     plt.savefig(f, bbox_inches='tight', pad_inches=0)
     plt.cla()
     plt.close()

+ 9 - 3
labelme/zoomWidget.py

@@ -17,8 +17,14 @@
 # along with Labelme.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from PyQt4.QtGui import *
-from PyQt4.QtCore import *
+try:
+    from PyQt5.QtGui import *
+    from PyQt5.QtCore import *
+    from PyQt5.QtWidgets import *
+except ImportError:
+    from PyQt4.QtGui import *
+    from PyQt4.QtCore import *
+
 
 class ZoomWidget(QSpinBox):
     def __init__(self, value=100):
@@ -27,7 +33,7 @@ class ZoomWidget(QSpinBox):
         self.setRange(1, 500)
         self.setSuffix(' %')
         self.setValue(value)
-        self.setToolTip(u'Zoom Level')
+        self.setToolTip('Zoom Level')
         self.setStatusTip(self.toolTip())
         self.setAlignment(Qt.AlignCenter)
 

+ 21 - 14
setup.py

@@ -8,6 +8,18 @@ import subprocess
 import sys
 
 
+try:
+    import PyQt5  # NOQA
+    PYQT_VERSION = 5
+except ImportError:
+    try:
+        import PyQt4  # NOQA
+        PYQT_VERSION = 4
+    except ImportError:
+        sys.stderr.write('Please install PyQt4 or PyQt5.\n')
+        sys.exit(1)
+
+
 version = '2.3.1'
 
 
@@ -20,31 +32,26 @@ if sys.argv[1] == 'release':
     sys.exit(0)
 
 
+here = osp.dirname(osp.abspath(__file__))
+
+
 class LabelmeBuildPyCommand(BuildPyCommand):
 
     def run(self):
-        if find_executable('pyrcc4') is None:
-            sys.stderr.write('Please install pyrcc4 command.\n')
+        pyrcc = 'pyrcc{:d}'.format(PYQT_VERSION)
+        if find_executable(pyrcc) is None:
+            sys.stderr.write('Please install {:s} command.\n'.format(pyrcc))
             sys.stderr.write('(See https://github.com/wkentaro/labelme.git)\n')
             sys.exit(1)
-        this_dir = osp.dirname(osp.abspath(__file__))
-        package_dir = osp.join(this_dir, 'labelme')
+        package_dir = osp.join(here, 'labelme')
         src = 'resources.qrc'
         dst = 'resources.py'
-        print('converting {0} -> {1}'
-              .format(osp.join(package_dir, src), osp.join(package_dir, dst)))
-        cmd = 'pyrcc4 -o {1} {0}'.format(src, dst)
+        cmd = '{pyrcc} -o {dst} {src}'.format(pyrcc=pyrcc, src=src, dst=dst)
+        print('+ {:s}'.format(cmd))
         subprocess.call(shlex.split(cmd), cwd=package_dir)
         BuildPyCommand.run(self)
 
 
-try:
-    import PyQt4
-except ImportError:
-    sys.stderr.write('Please install PyQt4.\n')
-    sys.exit(1)
-
-
 setup(
     name='labelme',
     version=version,