From cf38b16f90109df1771731f19df17362599386a6 Mon Sep 17 00:00:00 2001
From: Tingquan Gao <gaotingquan@baidu.com>
Date: Thu, 25 Feb 2021 16:19:21 +0800
Subject: [PATCH] Enhance dataloader robustness (#615)

* Enhance dataloader robustness

* Fix the bug about unexpected exiting when reading image failed
---
 deploy/cpp_infer/src/main.cpp |  5 +++++
 ppcls/data/reader.py          | 15 ++++++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/deploy/cpp_infer/src/main.cpp b/deploy/cpp_infer/src/main.cpp
index ec08cde6..9c420b32 100644
--- a/deploy/cpp_infer/src/main.cpp
+++ b/deploy/cpp_infer/src/main.cpp
@@ -70,6 +70,11 @@ int main(int argc, char **argv) {
   for (int idx = 0; idx < img_files_list.size(); ++idx) {
     std::string img_path = img_files_list[idx];
     cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);
+    if (!srcimg.data) {
+      std::cerr << "[ERROR] image read failed! image path: " << img_path
+                << "\n";
+      exit(-1);
+    }
     cv::cvtColor(srcimg, srcimg, cv::COLOR_BGR2RGB);
 
     double run_time = classifier.Run(srcimg);
diff --git a/ppcls/data/reader.py b/ppcls/data/reader.py
index edb56a60..376f7262 100755
--- a/ppcls/data/reader.py
+++ b/ppcls/data/reader.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
 import numpy as np
 import imghdr
 import os
@@ -182,11 +183,15 @@ class CommonDataset(Dataset):
         return
 
     def __getitem__(self, idx):
-        line = self.full_lines[idx]
-        img_path, label = line.split(self.delimiter)
-        img_path = os.path.join(self.params['data_dir'], img_path)
-        with open(img_path, 'rb') as f:
-            img = f.read()
+        try:
+            line = self.full_lines[idx]
+            img_path, label = line.split(self.delimiter)
+            img_path = os.path.join(self.params['data_dir'], img_path)
+            with open(img_path, 'rb') as f:
+                img = f.read()
+        except Exception as e:
+            logger.error('data read faild: "{}"'.format(line))
+            return self.__getitem__(random.randint(0, len(self)))
         return (transform(img, self.ops), int(label))
 
     def __len__(self):
-- 
GitLab