Spaces:
Runtime error
Runtime error
File size: 3,834 Bytes
73df9d6 6c83666 73df9d6 f30e3e5 73df9d6 409458e 73df9d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# image_moderator.py
from transformers import CLIPProcessor, CLIPModel, ViTForImageClassification, ViTFeatureExtractor
from PIL import Image, ImageFilter
import torch
import cv2
import os
class ImageContentModerator:
def __init__(self, nsfw_threshold=0.85, blur_radius=99):
# NSFW Setup
self.nsfw_threshold = nsfw_threshold
self.nsfw_classes = [
"porn", "nudity", "sexual activity", "explicit",
"safe", "neutral", "hentai", "suggestive", "drawing"
]
self.nsfw_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
self.nsfw_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
# Violence Setup
self.violence_model = ViTForImageClassification.from_pretrained('jaranohaal/vit-base-violence-detection')
self.violence_extractor = ViTFeatureExtractor.from_pretrained('jaranohaal/vit-base-violence-detection')
self.custom_labels = {1: "Non-Violent", 0: "Violent"}
self.blur_radius = blur_radius
print("[Moderator] Both models loaded successfully.")
def is_nsfw(self, image):
inputs = self.nsfw_processor(text=self.nsfw_classes, images=image, return_tensors="pt", padding=True)
with torch.no_grad():
outputs = self.nsfw_model(**inputs)
probs = outputs.logits_per_image.softmax(dim=1)[0]
top_class = self.nsfw_classes[probs.argmax()]
confidence = probs.max().item()
print(f"[NSFW] Predicted: {top_class} ({confidence:.2f})")
return top_class not in ['safe', 'neutral', 'drawing'] and confidence > self.nsfw_threshold
def is_violent(self, image):
inputs = self.violence_extractor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = self.violence_model(**inputs)
class_idx = outputs.logits.argmax(-1).item()
label = self.custom_labels[class_idx]
print(f"[Violence] Predicted: {label}")
return label == "Violent"
def blur_image(self, image_path, output_path):
image = Image.open(image_path)
blurred = image.filter(ImageFilter.GaussianBlur(radius=self.blur_radius))
blurred.save(output_path)
return output_path
def process_image(self, image_path, output_path="moderated_image.jpg"):
image = Image.open(image_path).convert("RGB")
if self.is_nsfw(image):
print("⚠️ NSFW Content Detected — Blurring Image")
return self.blur_image(image_path, output_path)
if self.is_violent(image):
print("⚠️ Violent Content Detected — Blurring Image")
return self.blur_image(image_path, output_path)
print("✅ Image is clean — No action taken")
image.save(output_path)
return output_path
from transformers import pipeline
from better_profanity import profanity
# Load profanity model
# pipeline("text-classification", model="unitary/toxic-bert").save_pretrained("./models/unitary-toxic-bert")
toxic_classifier = pipeline("text-classification", model="./models/unitary-toxic-bert")
profanity.load_censor_words()
custom_words = [
"idiot", "moron", "dumb", "stupid", "loser", "bastard", "retard", "scumbag",
"asshole", "jerk", "shit", "fuck", "damn", "hell", "crap", "bitch"
]
profanity.add_censor_words(custom_words)
def mask_bad_words(text):
return profanity.censor(text)
def smart_censor(text, toxic_threshold=0.85):
result = toxic_classifier(text)[0]
label = result['label'].lower()
score = result['score']
if label == "toxic" and score > toxic_threshold:
masked_text = mask_bad_words(text)
if masked_text != text:
return masked_text
else:
return "⚠️ Vulgar Content Detected"
return text
|