diff --git a/tests/_test.py b/tests/_test.py index cc3b7fc..7909634 100644 --- a/tests/_test.py +++ b/tests/_test.py @@ -6,6 +6,7 @@ from typing import Optional, Sequence from keras import backend as keras_backend +from PIL import Image import opennsfw2 as n2 @@ -61,9 +62,14 @@ def test_predict_images_simple_preprocessing(self) -> None: self._assert(expected_probabilities, predicted_probabilities) def test_predict_image(self) -> None: + # From path. self.assertAlmostEqual( 0.983, n2.predict_image(IMAGE_PATHS[1]), places=3 ) + # From PIL Image. + self.assertAlmostEqual( + 0.983, n2.predict_image(Image.open(IMAGE_PATHS[1])), places=3 + ) def test_predict_video_frames(self) -> None: elapsed_seconds, nsfw_probabilities = n2.predict_video_frames(