Spaces:
Running
Running
| import sys | |
| import traceback | |
| from pathlib import Path | |
| from PIL import Image | |
| from transformers import pipeline | |
| # Setup models exactly as in main.py | |
| print("Loading models...") | |
| models = {} | |
| device = -1 | |
| models['nsfw_image'] = pipeline( | |
| "image-classification", | |
| model="Falconsai/nsfw_image_detection", | |
| device=device | |
| ) | |
| models['nsfw_image_robust'] = pipeline( | |
| "image-classification", | |
| model="AdamCodd/vit-base-nsfw-detector", | |
| device=device | |
| ) | |
| file_path = Path("uploads/1779226485_Screenshot 2026-05-19 000650.png") | |
| try: | |
| print("Running moderation...") | |
| scores = [] | |
| # 1. Falconsai | |
| if 'nsfw_image' in models: | |
| try: | |
| image = Image.open(file_path) | |
| res = models['nsfw_image'](image) | |
| raw_score = float(res[0]['score']) | |
| label = res[0]['label'] | |
| nsfw_score = raw_score if label != "normal" else (1.0 - raw_score) | |
| scores.append(nsfw_score) | |
| print("Falconsai score:", nsfw_score) | |
| except Exception as e: | |
| print("Falconsai failed:") | |
| traceback.print_exc() | |
| # 2. AdamCodd | |
| if 'nsfw_image_robust' in models: | |
| try: | |
| image = Image.open(file_path) | |
| res = models['nsfw_image_robust'](image) | |
| raw_score = float(res[0]['score']) | |
| label = res[0]['label'] | |
| nsfw_score = raw_score if label == "nsfw" else (1.0 - raw_score) | |
| scores.append(nsfw_score) | |
| print("AdamCodd score:", nsfw_score) | |
| except Exception as e: | |
| print("AdamCodd failed:") | |
| traceback.print_exc() | |
| except Exception as e: | |
| print("Error:", e) | |
| traceback.print_exc() | |