|
import cv2 |
|
import torch |
|
import onnx |
|
import onnxruntime |
|
import numpy as np |
|
from tqdm import tqdm |
|
|
|
|
|
|
|
class NSFWChecker: |
|
def __init__(self, model_path=None, providers=["CPUExecutionProvider"]): |
|
model = onnx.load(model_path) |
|
self.input_name = model.graph.input[0].name |
|
session_options = onnxruntime.SessionOptions() |
|
self.session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers) |
|
|
|
def is_nsfw(self, img_paths, threshold = 0.85): |
|
skip_step = 1 |
|
total_len = len(img_paths) |
|
if total_len < 100: skip_step = 1 |
|
if total_len > 100 and total_len < 500: skip_step = 10 |
|
if total_len > 500 and total_len < 1000: skip_step = 20 |
|
if total_len > 1000 and total_len < 10000: skip_step = 50 |
|
if total_len > 10000: skip_step = 100 |
|
|
|
for idx in tqdm(range(0, total_len, skip_step), total=int(total_len // skip_step), desc="Checking for NSFW contents"): |
|
img = cv2.imread(img_paths[idx]) |
|
img = cv2.resize(img, (224,224)).astype('float32') |
|
img -= np.array([104, 117, 123], dtype=np.float32) |
|
img = np.expand_dims(img, axis=0) |
|
|
|
score = self.session.run(None, {self.input_name:img})[0][0][1] |
|
|
|
if score > threshold: |
|
print(f"Detected nsfw score:{score}") |
|
return True |
|
return False |
|
|
