Add feature of adding vertex

labelme/canvas.py

@@ -52,6 +52,7 @@ class Canvas(QtWidgets.QWidget):
         self.hideBackround = False
         self.hShape = None
         self.hVertex = None
+        self.addVertex = None
         self.movingShape = False
         self._painter = QtGui.QPainter()
         self._cursor = CURSOR_DEFAULT
@@ -116,6 +117,9 @@ class Canvas(QtWidgets.QWidget):
     def selectedVertex(self):
         return self.hVertex is not None
 
+    def addedVertex(self):
+        return self.addVertex is not None
+
     def mouseMoveEvent(self, ev):
         """Update line with last point and current coordinates."""
         if QT5:
@@ -183,6 +187,7 @@ class Canvas(QtWidgets.QWidget):
             # Look for a nearby vertex to highlight. If that fails,
             # check if we happen to be inside a shape.
             index = shape.nearestVertex(pos, self.epsilon)
+            index2 = shape.nearestEdge(pos, self.epsilon)
             if index is not None:
                 if self.selectedVertex():
                     self.hShape.highlightClear()
@@ -193,10 +198,20 @@ class Canvas(QtWidgets.QWidget):
                 self.setStatusTip(self.toolTip())
                 self.update()
                 break
+            elif index2 is not None:
+                if self.selectedVertex():
+                    self.hShape.highlightClear()
+                self.addVertex, self.hShape = index2, shape
+                self.hVertex = None
+                self.overrideCursor(CURSOR_DRAW)
+                self.setToolTip("Click & drag to add point")
+                self.setStatusTip(self.toolTip())
+                self.update()
+                break
             elif shape.containsPoint(pos):
                 if self.selectedVertex():
                     self.hShape.highlightClear()
-                self.hVertex, self.hShape = None, shape
+                self.hVertex, self.addVertex, self.hShape = None, None, shape
                 self.setToolTip(
                     "Click & drag to move shape '%s'" % shape.label)
                 self.setStatusTip(self.toolTip())
@@ -305,6 +320,13 @@ class Canvas(QtWidgets.QWidget):
             index, shape = self.hVertex, self.hShape
             shape.highlightVertex(index, shape.MOVE_VERTEX)
             return
+        if self.addedVertex() and self.hShape is not None:
+            index2, shape = self.addVertex, self.hShape
+            shape.insertPoint(index2, point)
+            shape.highlightVertex(index2, shape.MOVE_VERTEX)
+            self.addVertex = None
+            self.hVertex = index2
+            return
         for shape in reversed(self.shapes):
             if self.isVisible(shape) and shape.containsPoint(point):
                 shape.selected = True

labelme/lib.py

 from math import sqrt
 import os.path as osp
+from numpy.linalg import norm
+import numpy as np
 
 from qtpy import QtCore
 from qtpy import QtGui
@@ -68,6 +70,18 @@ def distance(p):
     return sqrt(p.x() * p.x() + p.y() * p.y())
 
 
+def distancetoline(point, line):
+    p1, p2 = line
+    p1 = np.array([p1.x(), p1.y()])
+    p2 = np.array([p2.x(), p2.y()])
+    p3 = np.array([point.x(), point.y()])
+    if np.dot((p3 - p1), (p2 - p1)) < 0:
+        return norm(p3 - p1)
+    if np.dot((p3 - p2), (p1 - p2)) < 0:
+        return norm(p3 - p2)
+    return norm(np.cross(p2 - p1, p1 - p3)) / norm(p2 - p1)
+
+
 def fmtShortcut(text):
     mod, key = text.split('+', 1)
     return '<b>%s</b>+<b>%s</b>' % (mod, key)

labelme/shape.py

 from qtpy import QtGui
 
 from labelme.lib import distance
+from labelme.lib import distancetoline
 from labelme import logger
 
 
@@ -71,6 +72,9 @@ class Shape(object):
             return self.points.pop()
         return None
 
+    def insertPoint(self, i, point):
+        self.points.insert(i, point)
+
     def isClosed(self):
         return self._closed
 
@@ -137,6 +141,17 @@ class Shape(object):
                 min_i = i
         return min_i
 
+    def nearestEdge(self, point, epsilon):
+        min_distance = float('inf')
+        post_i = None
+        for i in range(len(self.points)):
+            line = [self.points[i - 1], self.points[i]]
+            dist = distancetoline(point, line)
+            if dist <= epsilon and dist < min_distance:
+                min_distance = dist
+                post_i = i
+        return post_i
+
     def containsPoint(self, point):
         return self.makePath().contains(point)