HOTFIX: Black Image Issue

This commit is contained in:
Gourieff | 古仁
2025-03-02 00:49:01 +07:00
parent d2318ad140
commit 144a6fe224
+8 -3
View File
@@ -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