Key can add into key list after labeling a new one

PPOCRLabel/PPOCRLabel.py

@@ -42,6 +42,7 @@ sys.path.append("..")
 from paddleocr import PaddleOCR
 from libs.constants import *
 from libs.utils import *
+from libs.labelColor import label_colormap
 from libs.settings import Settings
 from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR, DEFAULT_LOCK_COLOR
 from libs.stringBundle import StringBundle
@@ -58,6 +59,8 @@ from libs.keyDialog import KeyDialog
 
 __appname__ = 'PPOCRLabel'
 
+LABEL_COLORMAP = label_colormap()
+
 
 class MainWindow(QMainWindow):
     FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
@@ -167,9 +170,8 @@ class MainWindow(QMainWindow):
 
         #  ================== Key List  ==================
         if self.kie_mode:
+            # self.keyList = QListWidget()
             self.keyList = UniqueLabelQListWidget()
-
-            # self.keyList.itemActivated.connect(self.boxSelectionChanged)
             self.keyList.itemSelectionChanged.connect(self.keyListSelectionChanged)
             self.keyList.itemDoubleClicked.connect(self.editBox)
             # Connect to itemChanged to detect checkbox changes.
@@ -521,8 +523,7 @@ class MainWindow(QMainWindow):
         addActions(labelMenu, (edit, delete))
 
         self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
-        self.labelList.customContextMenuRequested.connect(
-            self.popLabelListMenu)
+        self.labelList.customContextMenuRequested.connect(self.popLabelListMenu)
 
         # Draw squares/rectangles
         self.drawSquaresOption = QAction(getStr('drawSquares'), self)
@@ -1122,7 +1123,7 @@ class MainWindow(QMainWindow):
             for box in shapes:
                 trans_dic.append(
                     {"transcription": box['label'], "points": box['points'],
-                     "difficult": box['difficult'], "key": 0})
+                     "difficult": box['difficult'], "key": "None"})
             self.PPlabel[annotationFilePath] = trans_dic
             if mode == 'Auto':
                 self.Cachelabel[annotationFilePath] = trans_dic
@@ -1221,9 +1222,33 @@ class MainWindow(QMainWindow):
             self.canvas.resetAllLines()
 
         if self.kie_mode:
-            self.key_previous_text, flags = self.keyDialog.popUp(self.key_previous_text)
-            if not self.key_previous_text:
-                self.keyDialog.edit.setText(self.key_previous_text)
+            key_text, flags = self.keyDialog.popUp(self.key_previous_text)
+            if key_text is not None:
+                self.key_previous_text = key_text
+                if not self.keyList.findItemsByLabel(key_text):
+                    item = self.keyList.createItemFromLabel(key_text)
+                    self.keyList.addItem(item)
+                    rgb = self._get_rgb_by_label(key_text, self.kie_mode)
+                    self.keyList.setItemLabel(item, key_text, rgb)
+
+    def _update_shape_color(self, shape):
+        r, g, b = self._get_rgb_by_label(shape.label)
+        shape.line_color = QtGui.QColor(r, g, b)
+        shape.vertex_fill_color = QtGui.QColor(r, g, b)
+        shape.hvertex_fill_color = QtGui.QColor(255, 255, 255)
+        shape.fill_color = QtGui.QColor(r, g, b, 128)
+        shape.select_line_color = QtGui.QColor(255, 255, 255)
+        shape.select_fill_color = QtGui.QColor(r, g, b, 155)
+
+    def _get_rgb_by_label(self, label, kie_mode):
+        shift_auto_shape_color = 0  # use for random color
+        if kie_mode:
+            item = self.keyList.findItemsByLabel(label)[0]
+            label_id = self.keyList.indexFromItem(item).row() + 1
+            label_id += shift_auto_shape_color
+            return LABEL_COLORMAP[label_id % len(LABEL_COLORMAP)]
+        else:
+            return (0, 255, 0)
 
     def scrollRequest(self, delta, orientation):
         units = - delta / (8 * 15)
@@ -2188,7 +2213,7 @@ class MainWindow(QMainWindow):
             trans_dic = []
             for box in shapes:
                 trans_dic.append({"transcription": box['label'], "ratio": box['ratio'],
-                                  "difficult": box['difficult'], "key": 0 if "key" not in box else box["key"]})
+                                  "difficult": box['difficult'], "key": "None" if "key" not in box else box["key"]})
             self.canvas.lockedShapes = trans_dic
             self.actions.save.setEnabled(True)
 

PPOCRLabel/libs/keyDialog.py

@@ -212,4 +212,4 @@ class KeyDialog(QtWidgets.QDialog):
         if self.exec_():
             return self.edit.text(), self.getFlags()
         else:
-            return None, None, None
+            return None, None

PPOCRLabel/libs/labelColor.py

+import PIL.Image
+import numpy as np
+
+
+def rgb2hsv(rgb):
+    # type: (np.ndarray) -> np.ndarray
+    """Convert rgb to hsv.
+
+    Parameters
+    ----------
+    rgb: numpy.ndarray, (H, W, 3), np.uint8
+        Input rgb image.
+
+    Returns
+    -------
+    hsv: numpy.ndarray, (H, W, 3), np.uint8
+        Output hsv image.
+
+    """
+    hsv = PIL.Image.fromarray(rgb, mode="RGB")
+    hsv = hsv.convert("HSV")
+    hsv = np.array(hsv)
+    return hsv
+
+
+def hsv2rgb(hsv):
+    # type: (np.ndarray) -> np.ndarray
+    """Convert hsv to rgb.
+
+    Parameters
+    ----------
+    hsv: numpy.ndarray, (H, W, 3), np.uint8
+        Input hsv image.
+
+    Returns
+    -------
+    rgb: numpy.ndarray, (H, W, 3), np.uint8
+        Output rgb image.
+
+    """
+    rgb = PIL.Image.fromarray(hsv, mode="HSV")
+    rgb = rgb.convert("RGB")
+    rgb = np.array(rgb)
+    return rgb
+
+
+def label_colormap(n_label=256, value=None):
+    """Label colormap.
+
+    Parameters
+    ----------
+    n_label: int
+        Number of labels (default: 256).
+    value: float or int
+        Value scale or value of label color in HSV space.
+
+    Returns
+    -------
+    cmap: numpy.ndarray, (N, 3), numpy.uint8
+        Label id to colormap.
+
+    """
+
+    def bitget(byteval, idx):
+        return (byteval & (1 << idx)) != 0
+
+    cmap = np.zeros((n_label, 3), dtype=np.uint8)
+    for i in range(0, n_label):
+        id = i
+        r, g, b = 0, 0, 0
+        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))
+            id = id >> 3
+        cmap[i, 0] = r
+        cmap[i, 1] = g
+        cmap[i, 2] = b
+
+    if value is not None:
+        hsv = rgb2hsv(cmap.reshape(1, -1, 3))
+        if isinstance(value, float):
+            hsv[:, 1:, 2] = hsv[:, 1:, 2].astype(float) * value
+        else:
+            assert isinstance(value, int)
+            hsv[:, 1:, 2] = value
+        cmap = hsv2rgb(hsv).reshape(-1, 3)
+    return cmap