Add test for PIL Image input

bhky · Aug 17, 2024 · d85edb2 · d85edb2
1 parent 3e08be5
commit d85edb2
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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(