From 2365c70e4a0db922f4bbbb0144159a8488ab429f Mon Sep 17 00:00:00 2001
From: Jethong <1147925384@qq.com>
Date: Wed, 21 Apr 2021 19:36:58 +0800
Subject: [PATCH] fix dataset

---
 ppocr/data/pgnet_dataset.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/ppocr/data/pgnet_dataset.py b/ppocr/data/pgnet_dataset.py
index b9819358..5adcd02c 100644
--- a/ppocr/data/pgnet_dataset.py
+++ b/ppocr/data/pgnet_dataset.py
@@ -72,6 +72,7 @@ class PGDataSet(Dataset):
     def __getitem__(self, idx):
         file_idx = self.data_idx_order_list[idx]
         data_line = self.data_lines[file_idx]
+        img_id = 0
         try:
             data_line = data_line.decode('utf-8')
             substr = data_line.strip("\n").split(self.delimiter)
@@ -79,9 +80,10 @@ class PGDataSet(Dataset):
             label = substr[1]
             img_path = os.path.join(self.data_dir, file_name)
             if self.mode.lower() == 'eval':
-                img_id = int(data_line.split(".")[0][7:])
-            else:
-                img_id = 0
+                try:
+                    img_id = int(data_line.split(".")[0][7:])
+                except:
+                    img_id = 0
             data = {'img_path': img_path, 'label': label, 'img_id': img_id}
             if not os.path.exists(img_path):
                 raise Exception("{} does not exist!".format(img_path))
-- 
GitLab