152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
import os
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import ViTFeatureExtractor, ViTForImageClassification, pipeline
|
|
import torchvision.transforms as T
|
|
|
|
|
|
class AgeCheckerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"gate_enabled": ("BOOLEAN", {"default": True}),
|
|
"use_local_model": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BOOLEAN", "INT", "FLOAT", "STRING", "BOOLEAN")
|
|
RETURN_NAMES = ("is_underage", "predicted_age", "confidence", "status", "gate_output")
|
|
|
|
FUNCTION = "check_age"
|
|
CATEGORY = "Moderation/Age"
|
|
|
|
def __init__(self):
|
|
model_path = os.getenv("LOCAL_AGE_MODEL_PATH", "")
|
|
if model_path and os.path.exists(model_path):
|
|
self.model = ViTForImageClassification.from_pretrained(model_path).eval().to("cuda")
|
|
self.processor = ViTFeatureExtractor.from_pretrained(model_path)
|
|
else:
|
|
self.model = ViTForImageClassification.from_pretrained("nateraw/vit-age-classifier").eval().to("cuda")
|
|
self.processor = ViTFeatureExtractor.from_pretrained("nateraw/vit-age-classifier")
|
|
|
|
def check_age(self, image, gate_enabled, use_local_model):
|
|
if use_local_model:
|
|
model_path = os.getenv("LOCAL_AGE_MODEL_PATH", "")
|
|
if model_path and os.path.exists(model_path):
|
|
self.model = ViTForImageClassification.from_pretrained(model_path).eval().to("cuda")
|
|
self.processor = ViTFeatureExtractor.from_pretrained(model_path)
|
|
|
|
pil_image = Image.fromarray(image[0])
|
|
inputs = self.processor(images=pil_image, return_tensors="pt").to("cuda")
|
|
with torch.no_grad():
|
|
outputs = self.model(**inputs)
|
|
logits = outputs.logits
|
|
predicted_age = int(logits.argmax(-1).item())
|
|
confidence = float(logits.softmax(-1).max().item())
|
|
|
|
is_underage = predicted_age < 18
|
|
gate_output = not is_underage if gate_enabled else True
|
|
|
|
status = "Underage" if is_underage else "OK"
|
|
|
|
if gate_enabled and is_underage:
|
|
raise PermissionError(403, f"Blocked: Detected age {predicted_age} with confidence {confidence:.2f}.")
|
|
|
|
return (is_underage, predicted_age, confidence, status, gate_output)
|
|
|
|
|
|
class UnderageFilterNode:
|
|
def __init__(self):
|
|
self.classifier = pipeline("image-classification", model="nateraw/vit-age-classifier")
|
|
self.to_pil = T.ToPILImage()
|
|
self.underage_labels = {"0-2", "3-9", "10-19"}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"score": ("FLOAT", {
|
|
"default": 0.85,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"round": 0.001,
|
|
"display": "age_confidence_threshold"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("BOOLEAN",)
|
|
FUNCTION = "check_underage"
|
|
CATEGORY = "Moderation/Detection"
|
|
|
|
def check_underage(self, image, score):
|
|
img_tensor = image[0]
|
|
pil_img = self.to_pil(img_tensor.permute(2, 0, 1))
|
|
results = self.classifier(pil_img)
|
|
top = max(results, key=lambda r: r["score"])
|
|
is_underage = top["label"] in self.underage_labels and top["score"] >= score
|
|
return (is_underage,)
|
|
|
|
|
|
class MultiTypeGateNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": (["BOOLEAN", "INT", "FLOAT", "STRING"],),
|
|
"block_on": ("STRING", {
|
|
"default": "falsy", # options: 'falsy', 'truthy', 'equal'
|
|
}),
|
|
"match_value": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"message": ("STRING", {
|
|
"default": "Blocked by moderation gate."
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "evaluate"
|
|
CATEGORY = "Moderation/Blockers"
|
|
|
|
def evaluate(self, value, block_on, match_value, message):
|
|
should_block = False
|
|
|
|
if isinstance(value, str):
|
|
value = value.strip()
|
|
|
|
if block_on == "falsy":
|
|
should_block = not bool(value)
|
|
elif block_on == "truthy":
|
|
should_block = bool(value)
|
|
elif block_on == "equal":
|
|
try:
|
|
match_casted = type(value)(match_value)
|
|
except Exception:
|
|
match_casted = match_value
|
|
should_block = value == match_casted
|
|
|
|
if should_block:
|
|
raise PermissionError(403, message)
|
|
|
|
return ()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AgeCheckerNode": AgeCheckerNode,
|
|
"UnderageFilterNode": UnderageFilterNode,
|
|
"MultiTypeGateNode": MultiTypeGateNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AgeCheckerNode": "Age Checker",
|
|
"UnderageFilterNode": "Underage Filter",
|
|
"MultiTypeGateNode": "Flexible Gate",
|
|
}
|