[GitHub] szha closed pull request #12912: Add bytearray support back to imdecode (#12855, #12868)

[GitBox <gi...@apache.org>]

szha closed pull request #12912: Add bytearray support back to imdecode (#12855, #12868)
URL: https://github.com/apache/incubator-mxnet/pull/12912
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/python/mxnet/image/image.py b/python/mxnet/image/image.py
index eee2ccf14a8..b846700fb50 100644
--- a/python/mxnet/image/image.py
+++ b/python/mxnet/image/image.py
@@ -93,7 +93,7 @@ def imdecode(buf, *args, **kwargs):
 
     Parameters
     ----------
-    buf : str/bytes or numpy.ndarray
+    buf : str/bytes/bytearray or numpy.ndarray
         Binary image data as string or numpy ndarray.
     flag : int, optional, default=1
         1 for three channel color output. 0 for grayscale output.
@@ -135,10 +135,15 @@ def imdecode(buf, *args, **kwargs):
     <NDArray 224x224x3 @cpu(0)>
     """
     if not isinstance(buf, nd.NDArray):
-        if sys.version_info[0] == 3 and not isinstance(buf, (bytes, np.ndarray)):
-            raise ValueError('buf must be of type bytes or numpy.ndarray,'
+        if sys.version_info[0] == 3 and not isinstance(buf, (bytes, bytearray, np.ndarray)):
+            raise ValueError('buf must be of type bytes, bytearray or numpy.ndarray,'
                              'if you would like to input type str, please convert to bytes')
         buf = nd.array(np.frombuffer(buf, dtype=np.uint8), dtype=np.uint8)
+
+    if len(buf) == 0:
+        # empty buf causes OpenCV crash.
+        raise ValueError("input buf cannot be empty.")
+
     return _internal._cvimdecode(buf, *args, **kwargs)
 
 
diff --git a/tests/python/unittest/test_image.py b/tests/python/unittest/test_image.py
index 0df08af317a..c8022b67bee 100644
--- a/tests/python/unittest/test_image.py
+++ b/tests/python/unittest/test_image.py
@@ -92,6 +92,22 @@ def test_imdecode(self):
             cv_image = cv2.imread(img)
             assert_almost_equal(image.asnumpy(), cv_image)
 
+    def test_imdecode_bytearray(self):
+        try:
+            import cv2
+        except ImportError:
+            return
+        for img in TestImage.IMAGES:
+            with open(img, 'rb') as fp:
+                str_image = bytearray(fp.read())
+                image = mx.image.imdecode(str_image, to_rgb=0)
+            cv_image = cv2.imread(img)
+            assert_almost_equal(image.asnumpy(), cv_image)
+
+    @raises(ValueError)
+    def test_imdecode_empty_buffer(self):
+        mx.image.imdecode(b'', to_rgb=0)
+
     def test_scale_down(self):
         assert mx.image.scale_down((640, 480), (720, 120)) == (640, 106)
         assert mx.image.scale_down((360, 1000), (480, 500)) == (360, 375)


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services