From e7570d5bab1f07a071a3c1c32153c39caf7f16a1 Mon Sep 17 00:00:00 2001
From: HinGwenWoong <peterhuang0323@outlook.com>
Date: Sun, 20 Feb 2022 16:34:14 +0800
Subject: [PATCH] Load label successful

---
 PPOCRLabel/PPOCRLabel.py | 54 +++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 17 deletions(-)

diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py
index 9aaa2632..832dc4a6 100644
--- a/PPOCRLabel/PPOCRLabel.py
+++ b/PPOCRLabel/PPOCRLabel.py
@@ -79,14 +79,19 @@ class MainWindow(QMainWindow):
         self.settings.load()
         settings = self.settings
         self.lang = lang
-        self.kie_mode = kie_mode
-        self.key_previous_text = ""
+
         # Load string bundle for i18n
         if lang not in ['ch', 'en']:
             lang = 'en'
         self.stringBundle = StringBundle.getBundle(localeStr='zh-CN' if lang == 'ch' else 'en')  # 'en'
         getStr = lambda strId: self.stringBundle.getString(strId)
 
+        # KIE setting
+        self.kie_mode = kie_mode
+        self.key_previous_text = ""
+        self.existed_key_cls_set = set()
+        self.key_dialog_tip = getStr('keyDialogTip')
+
         self.defaultSaveDir = default_save_dir
         self.ocr = PaddleOCR(use_pdserving=False,
                              use_angle_cls=True,
@@ -425,26 +430,11 @@ class MainWindow(QMainWindow):
         }
 
         #  ================== New Actions ==================
-        # key list dialog
-        if kie_mode:
-            self.keyDialog = KeyDialog(
-                text=getStr('keyDialogTip'),
-                parent=self,
-                labels=None,
-                sort_labels=True,
-                show_text_field=True,
-                completion="startswith",
-                fit_to_content={'column': True, 'row': False},
-                flags=None
-            )
-        else:
-            self.keyDialog = None
 
         edit = action(getStr('editLabel'), self.editLabel,
                       'Ctrl+E', 'edit', getStr('editLabelDetail'),
                       enabled=False)
 
-        #  ================== New Actions ==================
         AutoRec = action(getStr('autoRecognition'), self.autoRecognition,
                          '', 'Auto', getStr('autoRecognition'), enabled=False)
 
@@ -652,6 +642,8 @@ class MainWindow(QMainWindow):
         elif self.filePath:
             self.queueEvent(partial(self.loadFile, self.filePath or ""))
 
+        self.keyDialog = None
+
         # Callbacks:
         self.zoomWidget.valueChanged.connect(self.paintCanvas)
 
@@ -1595,6 +1587,34 @@ class MainWindow(QMainWindow):
             self.Cachelabel = self.loadLabelFile(self.Cachelabelpath)
             if self.Cachelabel:
                 self.PPlabel = dict(self.Cachelabel, **self.PPlabel)
+
+            for image, info in self.PPlabel.items():
+                for box in info:
+                    if "key_cls" not in box:
+                        continue
+                    self.existed_key_cls_set.add(box["key_cls"])
+            if len(self.existed_key_cls_set) > 0:
+                for key_text in self.existed_key_cls_set:
+                    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)
+
+        # key list dialog
+        if self.kie_mode:
+            self.keyDialog = KeyDialog(
+                text=self.key_dialog_tip,
+                parent=self,
+                labels=self.existed_key_cls_set,
+                sort_labels=True,
+                show_text_field=True,
+                completion="startswith",
+                fit_to_content={'column': True, 'row': False},
+                flags=None
+            )
+
+
         self.lastOpenDir = dirpath
         self.dirname = dirpath
 
-- 
GitLab