From 851df47e2d1e03e23ec3d3cea4f352ee688120a7 Mon Sep 17 00:00:00 2001 From: Artur Niederfahrenhorst Date: Tue, 19 Nov 2024 12:32:27 +0100 Subject: [PATCH] [RLlib] Make opencv-python-headless default over opencv-python (#48776) ## Why are these changes needed? Make `opencv-python-headless` defeault requirement for RLlib/ This is necessary, because `opencv-python` may not work in many headless environments. Signed-off-by: hjiang --- .../ml/rllib-test-requirements.txt | 2 +- rllib/utils/images.py | 19 ++++++++++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/python/requirements/ml/rllib-test-requirements.txt b/python/requirements/ml/rllib-test-requirements.txt index 027c57446e60..c67bf2cec445 100644 --- a/python/requirements/ml/rllib-test-requirements.txt +++ b/python/requirements/ml/rllib-test-requirements.txt @@ -5,7 +5,7 @@ # Atari ale_py==0.10.1 imageio==2.34.2 -opencv-python==4.8.1.78 +opencv-python-headless==4.8.1.78 # For testing MuJoCo envs with gymnasium. mujoco==3.2.4 diff --git a/rllib/utils/images.py b/rllib/utils/images.py index 0716ea5c45b4..7b0f1601d574 100644 --- a/rllib/utils/images.py +++ b/rllib/utils/images.py @@ -1,4 +1,5 @@ import logging +import importlib import numpy as np @@ -6,13 +7,29 @@ logger = logging.getLogger(__name__) + +@DeveloperAPI +def is_package_installed(package_name): + try: + importlib.metadata.version(package_name) + return True + except importlib.metadata.PackageNotFoundError: + return False + + try: import cv2 cv2.ocl.setUseOpenCL(False) logger.debug("CV2 found for image processing.") -except ImportError: +except ImportError as e: + if is_package_installed("opencv-python"): + raise ImportError( + f"OpenCV is installed, but we failed to import it. This may be because " + f"you need to install `opencv-python-headless` instead of " + f"`opencv-python`. Error message: {e}", + ) cv2 = None