101 lines
2.0 KiB
Python
101 lines
2.0 KiB
Python
from PIL import Image
|
|
|
|
from app.config.image_policy import FORBIDDEN_IMAGE_LABELS
|
|
from app.config.settings import settings
|
|
from app.ml.image.clip_classifier import ClipClassifier
|
|
from app.ml.image.image_classifier import ImageClassifier
|
|
from app.ml.image.image_prediction_result import ImagePredictionResult
|
|
from app.moderation.image.validator import ImageValidator
|
|
|
|
class ImageModerationService:
|
|
|
|
def __init__(
|
|
self,
|
|
classifier: ImageClassifier,
|
|
validator: ImageValidator,
|
|
clip_classifier: ClipClassifier
|
|
):
|
|
|
|
self.classifier = classifier
|
|
self.clip_classifier = clip_classifier
|
|
self.validator = validator
|
|
|
|
def moderate(
|
|
self,
|
|
image: Image.Image
|
|
) -> ImagePredictionResult:
|
|
|
|
self.validator.validate(image)
|
|
|
|
# =========================
|
|
# 1. NSFW MODEL
|
|
# =========================
|
|
|
|
nsfw_prediction = self.classifier.predict(image)
|
|
|
|
if (
|
|
nsfw_prediction.label.lower() == "nsfw"
|
|
and nsfw_prediction.score >= settings.NSFW_THRESHOLD
|
|
):
|
|
nsfw_prediction.approved = False
|
|
|
|
nsfw_prediction.reason = "NSFW"
|
|
|
|
nsfw_prediction.detected_labels = [
|
|
nsfw_prediction.label
|
|
]
|
|
|
|
return nsfw_prediction
|
|
|
|
# =========================
|
|
# 2. CLIP MODEL
|
|
# =========================
|
|
|
|
clip_prediction = (
|
|
self.clip_classifier.predict(image)
|
|
)
|
|
|
|
detected_forbidden = [
|
|
|
|
label
|
|
|
|
for label in clip_prediction.detected_labels
|
|
|
|
if label.lower()
|
|
in FORBIDDEN_IMAGE_LABELS
|
|
|
|
]
|
|
|
|
if detected_forbidden:
|
|
return ImagePredictionResult(
|
|
|
|
label=clip_prediction.label,
|
|
|
|
score=clip_prediction.score,
|
|
|
|
approved=False,
|
|
|
|
reason="FORBIDDEN_CONTENT",
|
|
|
|
detected_labels=detected_forbidden
|
|
|
|
)
|
|
|
|
# =========================
|
|
# 3. NORMAL IMAGE
|
|
# =========================
|
|
|
|
return ImagePredictionResult(
|
|
|
|
label="normal",
|
|
|
|
score=1.0,
|
|
|
|
approved=True,
|
|
|
|
reason="OK",
|
|
|
|
detected_labels=[]
|
|
|
|
)
|