Files
Gourieff-ComfyUI-ReActor/scripts/reactor_sfw.py
T
Gourieff | 古仁 e620530232 FIX: Threshold adjustment
0.965 and less - is safety content
2025-01-20 14:57:59 +07:00

14 lines
429 B
Python

from transformers import pipeline
from PIL import Image
import logging
SCORE = 0.965 # 0.965 and less - is safety content
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