imaginAIry/tests/test_safety.py
Bryce 6a80759016 feature: improved safety filter
- provides more informative logs
- provides a detailed safety score object
- adds non-bypassable filter for extreme content
2022-10-10 03:32:19 -05:00

16 lines
415 B
Python

from PIL import Image
from imaginairy.safety import create_safety_score
from tests import TESTS_FOLDER
def test_is_nsfw():
img = Image.open(f"{TESTS_FOLDER}/data/safety.jpg")
safety_score = create_safety_score(img)
assert safety_score.is_nsfw
img = Image.open(f"{TESTS_FOLDER}/data/girl_with_a_pearl_earring.jpg")
safety_score = create_safety_score(img)
assert not safety_score.is_nsfw