diff --git a/ppocr/data/imaug/label_ops.py b/ppocr/data/imaug/label_ops.py index c9bc2e7722e8027ce870e4969bfcdab720495c28..02a5187dad27b76d04e866de45333d79383c1347 100644 --- 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 anno.keys(): + labels.append(ann['label']) + elif 'key_cls' in anno.keys(): + labels.append(anno['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'],