Merge pull request #6240 from PeterH0323/enhance-kie

[PPOCRLabel BUG] Fix re-reco will clear all key message 🐛 

PPOCRLabel/PPOCRLabel.py

@@ -1149,6 +1149,9 @@ class MainWindow(QMainWindow):
         for box in self.result_dic:
             trans_dic = {"label": box[1][0], "points": box[0], "difficult": False}
             if self.kie_mode:
+                if len(box) == 3:
+                    trans_dic.update({"key_cls": box[2]})
+                else:
                     trans_dic.update({"key_cls": "None"})
             if trans_dic["label"] == "" and mode == 'Auto':
                 continue
@@ -2047,6 +2050,7 @@ class MainWindow(QMainWindow):
             rec_flag = 0
             for shape in self.canvas.shapes:
                 box = [[int(p.x()), int(p.y())] for p in shape.points]
+                kie_cls = shape.key_cls
 
                 if len(box) > 4:
                     box = self.gen_quad_from_poly(np.array(box))
@@ -2062,15 +2066,25 @@ class MainWindow(QMainWindow):
                     if shape.line_color == DEFAULT_LOCK_COLOR:
                         shape.label = result[0][0]
                         result.insert(0, box)
+                        if self.kie_mode:
+                            result.append(kie_cls)
                         self.result_dic_locked.append(result)
                     else:
                         result.insert(0, box)
+                        if self.kie_mode:
+                            result.append(kie_cls)
                         self.result_dic.append(result)
                 else:
                     print('Can not recognise the box')
                     if shape.line_color == DEFAULT_LOCK_COLOR:
                         shape.label = result[0][0]
+                        if self.kie_mode:
+                            self.result_dic_locked.append([box, (self.noLabelText, 0), kie_cls])
+                        else:
                             self.result_dic_locked.append([box, (self.noLabelText, 0)])
+                    else:
+                        if self.kie_mode:
+                            self.result_dic.append([box, (self.noLabelText, 0), kie_cls])
                         else:
                             self.result_dic.append([box, (self.noLabelText, 0)])
                 try: