Assign different labels with different colors

libs/canvas.py

@@ -41,8 +41,9 @@ class Canvas(QWidget):
         self.current = None
         self.selectedShape = None  # save the selected shape here
         self.selectedShapeCopy = None
-        self.lineColor = QColor(0, 0, 255)
-        self.line = Shape(line_color=self.lineColor)
+        self.drawingLineColor = QColor(0, 0, 255)
+        self.drawingRectColor = QColor(0, 0, 255) 
+        self.line = Shape(line_color=self.drawingLineColor)
         self.prevPoint = QPointF()
         self.offsets = QPointF(), QPointF()
         self.scale = 1.0
@@ -61,6 +62,10 @@ class Canvas(QWidget):
         self.setFocusPolicy(Qt.WheelFocus)
         self.verified = False
 
+    def setDrawingColor(self, qColor):
+        self.drawingLineColor = qColor
+        self.drawingRectColor = qColor
+
     def enterEvent(self, ev):
         self.overrideCursor(self._cursor)
 
@@ -103,7 +108,7 @@ class Canvas(QWidget):
         if self.drawing():
             self.overrideCursor(CURSOR_DRAW)
             if self.current:
-                color = self.lineColor
+                color = self.drawingLineColor
                 if self.outOfPixmap(pos):
                     # Don't allow the user to draw outside the pixmap.
                     # Project the point to the pixmap's edges.
@@ -414,8 +419,7 @@ class Canvas(QWidget):
             rightBottom = self.line[1]
             rectWidth = rightBottom.x() - leftTop.x()
             rectHeight = rightBottom.y() - leftTop.y()
-            color = QColor(0, 220, 0)
-            p.setPen(color)
+            p.setPen(self.drawingRectColor)
             brush = QBrush(Qt.BDiagPattern)
             p.setBrush(brush)
             p.drawRect(leftTop.x(), leftTop.y(), rectWidth, rectHeight)
@@ -606,9 +610,15 @@ class Canvas(QWidget):
         points = [p1+p2 for p1, p2 in zip(self.selectedShape.points, [step]*4)]
         return True in map(self.outOfPixmap, points)
 
-    def setLastLabel(self, text):
+    def setLastLabel(self, text, line_color  = None, fill_color = None):
         assert text
         self.shapes[-1].label = text
+        if line_color:
+            self.shapes[-1].line_color = line_color
+        
+        if fill_color:
+            self.shapes[-1].fill_color = fill_color
+
         return self.shapes[-1]
 
     def undoLastLine(self):

libs/labelFile.py

@@ -58,6 +58,35 @@ class LabelFile(object):
     def toggleVerify(self):
         self.verified = not self.verified
 
+    ''' ttf is disable
+    def load(self, filename):
+        import json
+        with open(filename, 'rb') as f:
+                data = json.load(f)
+                imagePath = data['imagePath']
+                imageData = b64decode(data['imageData'])
+                lineColor = data['lineColor']
+                fillColor = data['fillColor']
+                shapes = ((s['label'], s['points'], s['line_color'], s['fill_color'])\
+                        for s in data['shapes'])
+                # Only replace data after everything is loaded.
+                self.shapes = shapes
+                self.imagePath = imagePath
+                self.imageData = imageData
+                self.lineColor = lineColor
+                self.fillColor = fillColor
+
+    def save(self, filename, shapes, imagePath, imageData, lineColor=None, fillColor=None):
+        import json
+        with open(filename, 'wb') as f:
+                json.dump(dict(
+                    shapes=shapes,
+                    lineColor=lineColor, fillColor=fillColor,
+                    imagePath=imagePath,
+                    imageData=b64encode(imageData)),
+                    f, ensure_ascii=True, indent=2)
+    '''
+
     @staticmethod
     def isLabelFile(filename):
         fileSuffix = os.path.splitext(filename)[1].lower()

libs/lib.py

@@ -1,5 +1,5 @@
 from math import sqrt
-
+import hashlib
 try:
     from PyQt5.QtGui import *
     from PyQt5.QtCore import *
@@ -71,3 +71,23 @@ def distance(p):
 def fmtShortcut(text):
     mod, key = text.split('+', 1)
     return '<b>%s</b>+<b>%s</b>' % (mod, key)
+
+
+def generateColorByText(text):
+    color_table = []
+    color_table.append(QColor(0, 0, 50))
+    color_table.append(QColor(0, 0, 255))
+    color_table.append(QColor(0, 50, 0))
+    color_table.append(QColor(0, 255, 0))
+    color_table.append(QColor(50, 0, 0))
+    color_table.append(QColor(255, 0, 0))
+    color_table.append(QColor(0, 50, 50))
+    color_table.append(QColor(0, 255, 255))
+    color_table.append(QColor(50, 50, 0))
+    color_table.append(QColor(255, 255, 0))
+    color_table.append(QColor(50, 0, 50))
+    color_table.append(QColor(255, 0, 255))
+    color_table.append(QColor(50, 50, 50))
+    color_table.append(QColor(255, 255, 255))
+    colorInd = int(hashlib.sha1(text.encode('utf-8')).hexdigest(), 16) % len(color_table)
+    return color_table[colorInd]