diff --git a/tests/test_feature_extraction_common.py b/tests/test_feature_extraction_common.py index d6420ad3c58..f226244fd0a 100644 --- a/tests/test_feature_extraction_common.py +++ b/tests/test_feature_extraction_common.py @@ -68,10 +68,15 @@ def prepare_image_inputs(feature_extract_tester, equal_resolution=False, numpify ) else: image_inputs = [] + + # To avoid getting image width/height 0 + min_resolution = feature_extract_tester.min_resolution + if getattr(feature_extract_tester, "size_divisor", None): + # If `size_divisor` is defined, the image needs to have width/size >= `size_divisor` + min_resolution = max(feature_extract_tester.size_divisor, min_resolution) + for i in range(feature_extract_tester.batch_size): - width, height = np.random.choice( - np.arange(feature_extract_tester.min_resolution, feature_extract_tester.max_resolution), 2 - ) + width, height = np.random.choice(np.arange(min_resolution, feature_extract_tester.max_resolution), 2) image_inputs.append( np.random.randint(255, size=(feature_extract_tester.num_channels, width, height), dtype=np.uint8) )