fix re infer bug

ppocr/data/imaug/label_ops.py

@@ -833,10 +833,11 @@ class VQATokenLabelEncode(object):
         segment_offset_id = []
         gt_label_list = []
 
-        if self.contains_re:
-            # for re
         entities = []
-            if not self.infer_mode:
+
+        # for re
+        train_re = self.contains_re and not self.infer_mode
+        if train_re:
             relations = []
             id2label = {}
             entity_id_to_index_map = {}
@@ -845,7 +846,7 @@ class VQATokenLabelEncode(object):
         data['ocr_info'] = copy.deepcopy(ocr_info)
 
         for info in ocr_info:
-            if self.contains_re and not self.infer_mode:
+            if train_re:
                 # for re
                 if len(info["text"]) == 0:
                     empty_entity.add(info["id"])
@@ -872,8 +873,7 @@ class VQATokenLabelEncode(object):
                 gt_label = self._parse_label(label, encode_res)
 
             # construct entities for re
-            if self.contains_re:
-                if not self.infer_mode:
+            if train_re:
                 if gt_label[0] != self.label2id_map["O"]:
                     entity_id_to_index_map[info["id"]] = len(entities)
                     label = label.upper()
@@ -886,8 +886,7 @@ class VQATokenLabelEncode(object):
             else:
                 entities.append({
                     "start": len(input_ids_list),
-                        "end":
-                        len(input_ids_list) + len(encode_res["input_ids"]),
+                    "end": len(input_ids_list) + len(encode_res["input_ids"]),
                     "label": 'O',
                 })
             input_ids_list.extend(encode_res["input_ids"])
@@ -908,12 +907,9 @@ class VQATokenLabelEncode(object):
             padding_side=self.tokenizer.padding_side,
             pad_token_type_id=self.tokenizer.pad_token_type_id,
             pad_token_id=self.tokenizer.pad_token_id)
-
-        if self.contains_re:
         data['entities'] = entities
-            if self.infer_mode:
-                data['ocr_info'] = ocr_info
-            else:
+
+        if train_re:
             data['relations'] = relations
             data['id2label'] = id2label
             data['empty_entity'] = empty_entity
@@ -921,6 +917,13 @@ class VQATokenLabelEncode(object):
         return data
 
     def _load_ocr_info(self, data):
+        def trans_poly_to_bbox(poly):
+            x1 = np.min([p[0] for p in poly])
+            x2 = np.max([p[0] for p in poly])
+            y1 = np.min([p[1] for p in poly])
+            y2 = np.max([p[1] for p in poly])
+            return [x1, y1, x2, y2]
+
         if self.infer_mode:
             ocr_result = self.ocr_engine.ocr(data['image'], cls=False)
             ocr_info = []