diff --git a/python/paddle/dataset/image.py b/python/paddle/dataset/image.py index 1cd50bd1802095db07e5618f37b0d42d11e94760..b32736ee7c265e3a94207afc04673eec4fcf1c6e 100644 --- a/python/paddle/dataset/image.py +++ b/python/paddle/dataset/image.py @@ -36,11 +36,6 @@ import numpy as np try: import cv2 except ImportError: - import sys - sys.stderr.write( - '''Warning with paddle image module: opencv-python should be imported, - or paddle image module could NOT work; please install opencv-python first.''' - ) cv2 = None import os import tarfile @@ -53,6 +48,18 @@ __all__ = [ ] +def _check_cv2(): + if cv2 is None: + import sys + sys.stderr.write( + '''Warning with paddle image module: opencv-python should be imported, + or paddle image module could NOT work; please install opencv-python first.''' + ) + return False + else: + return True + + def batch_images_from_tar(data_file, dataset_name, img2label, @@ -134,7 +141,7 @@ def load_image_bytes(bytes, is_color=True): load and return a gray image. :type is_color: bool """ - assert cv2 is not None + assert _check_cv2() is True flag = 1 if is_color else 0 file_bytes = np.asarray(bytearray(bytes), dtype=np.uint8) @@ -159,7 +166,7 @@ def load_image(file, is_color=True): load and return a gray image. :type is_color: bool """ - assert cv2 is not None + assert _check_cv2() is True # cv2.IMAGE_COLOR for OpenCV3 # cv2.CV_LOAD_IMAGE_COLOR for older OpenCV Version @@ -188,7 +195,7 @@ def resize_short(im, size): :param size: the shorter edge size of image after resizing. :type size: int """ - assert cv2 is not None + assert _check_cv2() is True h, w = im.shape[:2] h_new, w_new = size, size