Update label_ops.py

2022-05-12 16:14:08 +08:00 · 2022-05-12 16:14:08 +08:00 · 6024fc9f39
parent 400989320b
commit 6024fc9f39
1 changed files with 6 additions and 1 deletions
--- a/ppocr/data/imaug/label_ops.py
+++ b/ppocr/data/imaug/label_ops.py
@ -438,7 +438,12 @@ class KieLabelEncode(object):
            texts.append(ann['transcription'])
            text_ind = [self.dict[c] for c in text if c in self.dict]
            text_inds.append(text_ind)
-            labels.append(ann['label'])
+            if 'label' in ann.keys():
+                labels.append(ann['label'])
+            elif 'key_cls' in ann.keys():
+                labels.append(ann['key_cls'])
+            else:
+                raise ValueError("Cannot found 'key_cls' in ann.keys(), please check your training annotation.")
            edges.append(ann.get('edge', 0))
        ann_infos = dict(
            image=data['image'],