From 144a6fe2248cea2159d4457611efa9f9236ca657 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gourieff=20=7C=20=E5=8F=A4=E4=BB=81?= Date: Sun, 2 Mar 2025 00:49:01 +0700 Subject: [PATCH] HOTFIX: Black Image Issue --- scripts/reactor_sfw.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/scripts/reactor_sfw.py b/scripts/reactor_sfw.py index 524b074..50fbfc3 100644 --- a/scripts/reactor_sfw.py +++ b/scripts/reactor_sfw.py @@ -1,13 +1,18 @@ from transformers import pipeline from PIL import Image import logging +from scripts.reactor_logger import logger -SCORE = 0.965 # 0.965 and less - is safety content +SCORE = 0.96 -logging.getLogger('transformers').setLevel(logging.ERROR) +logging.getLogger("transformers").setLevel(logging.ERROR) def nsfw_image(img_path: str, model_path: str): + 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