clasificación de imágenes nsfw grande

Instale mis poderes:

pip install transformers==4.37.2 torch==2.3.1 torchvision Pillow

Invócame en tu código:

import os
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoProcessor, FocalNetForImageClassification

# Ruta a la carpeta con imágenes
image_folder = ""
# Ruta al modelo
model_path = "MichalMlodawski/nsfw-image-detection-large"

# Lista de archivos jpg en la carpeta
jpg_files = [file for file in os.listdir(image_folder) if file.lower().endswith(".jpg")]

# Verificar si hay archivos jpg en la carpeta
if not jpg_files:
    print("🚫 No se encontraron archivos jpg en la carpeta:", image_folder)
    exit()

# Cargar el modelo y el extractor de características
feature_extractor = AutoProcessor.from_pretrained(model_path)
model = FocalNetForImageClassification.from_pretrained(model_path)
model.eval()

# Transformaciones de la imagen
transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Mapeo de etiquetas del modelo a categorías NSFW
label_to_category = {
    "LABEL_0": "Seguro",
    "LABEL_1": "Cuestionable",
    "LABEL_2": "Inseguro"
}

# Procesamiento y predicción para cada imagen
results = []
for jpg_file in jpg_files:
    selected_image = os.path.join(image_folder, jpg_file)
    image = Image.open(selected_image).convert("RGB")
    image_tensor = transform(image).unsqueeze(0)

    # Procesar la imagen usando el extractor de características
    inputs = feature_extractor(images=image, return_tensors="pt")

    # Predicción usando el modelo
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        confidence, predicted = torch.max(probabilities, 1)

    # Obtener la etiqueta de la configuración del modelo
    label = model.config.id2label[predicted.item()]

    results.append((jpg_file, label, confidence.item() * 100))

# Mostrar resultados
print("🖼️ Resultados de Clasificación NSFW 🖼️")
print("=" * 40)
for jpg_file, label, confidence in results:
    category = label_to_category.get(label, "Desconocido")
    emoji = {"Seguro": "✅", "Cuestionable": "⚠️", "Inseguro": "🔞"}.get(category, "❓")
    confidence_bar = "🟩" * int(confidence // 10) + "⬜" * (10 - int(confidence // 10))

    print(f"📄 Nombre del archivo: {jpg_file}")
    print(f"🏷️ Etiqueta del modelo: {label}")
    print(f"{emoji} Categoría NSFW: {category}")
    print(f"🎯 Confianza: {confidence:.2f}% {confidence_bar}")
    print(f"{'=' * 40}")

print("🏁 ¡Clasificación completada! 🎉")