35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
from transformers import pipeline
|
|
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"""
|
|
if not os.path.exists(nsfwdet_model_path):
|
|
os.makedirs(nsfwdet_model_path)
|
|
nd_urls = [
|
|
"https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/config.json",
|
|
"https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/model.safetensors",
|
|
"https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/preprocessor_config.json",
|
|
]
|
|
for model_url in nd_urls:
|
|
model_name = os.path.basename(model_url)
|
|
model_path = os.path.join(nsfwdet_model_path, model_name)
|
|
download(model_url, model_path, model_name)
|
|
|
|
SCORE = 0.96
|
|
|
|
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)
|
|
if result[0]["label"] == "nsfw" and result[0]["score"] > SCORE:
|
|
logger.status(f"NSFW content detected, skipping...")
|
|
return True
|
|
return False
|