nsfw_image_detection

from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests
import torch

# Load the model and processor
model_name = "acaciabengo/nsfw_image_detection"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# Mapping IDs to human-readable labels
labels = {'drawings': 0, 'hentai': 1, 'neutral': 2, 'porn': 3, 'sexy': 4}
id_to_label = {v: k for k, v in labels.items()}

# Load an image from a URL or local path
url = "https://example.com/sample_image.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# Preprocess image and perform inference
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

# Calculate probabilities and get the top prediction
probabilities = logits.softmax(dim=-1).squeeze().tolist()
predicted_class_idx = logits.argmax(-1).item()

print(f"Predicted label: {id_to_label[predicted_class_idx]}")
print("\nProbabilities:")
for i, prob in enumerate(probabilities):
    print(f"- {id_to_label[i]}: {prob:.2%}")