diff --git a/scripts/reactor_sfw.py b/scripts/reactor_sfw.py index b2dfd0c..7c730c1 100644 --- a/scripts/reactor_sfw.py +++ b/scripts/reactor_sfw.py @@ -3,6 +3,7 @@ from PIL import Image import logging import os from reactor_utils import download +from scripts.reactor_logger import logger def ensure_nsfw_model(nsfwdet_model_path): """Download NSFW detection model if it doesn't exist""" @@ -18,13 +19,16 @@ def ensure_nsfw_model(nsfwdet_model_path): model_path = os.path.join(nsfwdet_model_path, model_name) download(model_url, model_path, model_name) -SCORE = 0.972 +SCORE = 0.96 -logging.getLogger('transformers').setLevel(logging.ERROR) +logging.getLogger("transformers").setLevel(logging.ERROR) def nsfw_image(img_path: str, model_path: str): ensure_nsfw_model(model_path) with Image.open(img_path) as img: predict = pipeline("image-classification", model=model_path) result = predict(img) - return True if result[0]["score"] > SCORE else False + if result[0]["label"] == "nsfw" and result[0]["score"] > SCORE: + logger.status(f"NSFW content detected, skipping...") + return True + return False