Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
Add bytearray support back to imdecode (#12855, #12868)
Browse files Browse the repository at this point in the history
1. Avoid raise exception when input is bytearray.
2. Avoid OpenCV crash for empty input.
3. Added unittests.
  • Loading branch information
frankfliu committed Oct 23, 2018
1 parent 0874677 commit faf0fd5
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 3 deletions.
11 changes: 8 additions & 3 deletions python/mxnet/image/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)


Expand Down
16 changes: 16 additions & 0 deletions tests/python/unittest/test_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit faf0fd5

Please sign in to comment.