user friendly image loading (apache#7974)

fix modify log message fix pylint
crazy-cat · Oct 26, 2017 · d80e1cb · d80e1cb
1 parent 93dcf05
commit d80e1cb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/python/mxnet/image/image.py b/python/mxnet/image/image.py
@@ -1180,7 +1180,23 @@ def check_valid_image(self, data):
     def imdecode(self, s):
         """Decodes a string or byte string to an NDArray.
         See mx.img.imdecode for more details."""
-        return imdecode(s)
+        def locate():
+            """Locate the image file/index if decode fails."""
+            if self.seq is not None:
+                idx = self.seq[self.cur - 1]
+            else:
+                idx = self.cur - 1
+            if self.imglist is not None:
+                _, fname = self.imglist[idx]
+                msg = "filename: {}".format(fname)
+            else:
+                msg = "index: {}".format(idx)
+            return "Broken image " + msg
+        try:
+            img = imdecode(s)
+        except Exception as e:
+            raise RuntimeError("{}, {}".format(locate(), e))
+        return img
 
     def read_image(self, fname):
         """Reads an input image `fname` and returns the decoded raw bytes.

diff --git a/src/io/image_io.cc b/src/io/image_io.cc
@@ -142,7 +142,7 @@ void ImdecodeImpl(int flag, bool to_rgb, void* data, size_t size,
   if (out->is_none()) {
     cv::Mat res = cv::imdecode(buf, flag);
     if (res.empty()) {
-      LOG(INFO) << "Invalid image file. Only supports png and jpg.";
+      LOG(INFO) << "Decoding failed. Invalid image file.";
       *out = NDArray();
       return;
     }
@@ -151,18 +151,20 @@ void ImdecodeImpl(int flag, bool to_rgb, void* data, size_t size,
     dst = cv::Mat(out->shape()[0], out->shape()[1], flag == 0 ? CV_8U : CV_8UC3,
                   out->data().dptr_);
     res.copyTo(dst);
+    CHECK(!dst.empty()) << "Failed copying buffer to output.";
   } else {
     dst = cv::Mat(out->shape()[0], out->shape()[1], flag == 0 ? CV_8U : CV_8UC3,
                 out->data().dptr_);
 #if (CV_MAJOR_VERSION > 2 || (CV_MAJOR_VERSION == 2 && CV_MINOR_VERSION >=4))
     cv::imdecode(buf, flag, &dst);
+    CHECK(!dst.empty()) << "Decoding failed. Invalid image file.";
 #else
     cv::Mat tmp = cv::imdecode(buf, flag);
-    CHECK(!tmp.empty());
+    CHECK(!tmp.empty()) << "Decoding failed. Invalid image file.";
     tmp.copyTo(dst);
+    CHECK(!dst.empty()) << "Failed copying buffer to output.";
 #endif
   }
-  CHECK(!dst.empty());
   CHECK_EQ(static_cast<void*>(dst.ptr()), out->data().dptr_);
   if (to_rgb && flag != 0) {
     cv::cvtColor(dst, dst, CV_BGR2RGB);