diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py index 6376173bad63bcd0656c07d37a86ec48a69f506d..03653a998f824ac66d30a67260e704b67b03e7c3 100644 --- a/PPOCRLabel/PPOCRLabel.py +++ b/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) diff --git a/PPOCRLabel/libs/keyDialog.py b/PPOCRLabel/libs/keyDialog.py index 028227aa98d5a5496ce529b843535ce4d8d81232..a910428bc13aaca54e255c2b2397694327faad97 100644 --- a/PPOCRLabel/libs/keyDialog.py +++ b/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 diff --git a/PPOCRLabel/libs/labelColor.py b/PPOCRLabel/libs/labelColor.py new file mode 100644 index 0000000000000000000000000000000000000000..c6f933981f3ca13981910a88fca76f884d727a14 --- /dev/null +++ b/PPOCRLabel/libs/labelColor.py @@ -0,0 +1,88 @@ +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