Coders-Nexa/nexa-api

1from transformers import AutoModelForImageClassification2import torch3import cv24import numpy as np5 6MODEL_NAME = "giacomoarienti/nsfw-classifier"7 8device = "cuda" if torch.cuda.is_available() else "cpu"9 10model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(device)11model.eval()12 13id2label = model.config.id2label14 15 16# ✅ Manual preprocessing17def preprocess(img):18    img = cv2.resize(img, (224, 224))19    img = img / 255.0  # normalize20 21    # HWC → CHW22    img = np.transpose(img, (2, 0, 1))23 24    # to tensor25    img = torch.tensor(img, dtype=torch.float32).unsqueeze(0)26 27    return img.to(device)28 29 30def scan_image(img):31    inputs = preprocess(img)32 33    with torch.no_grad():34        outputs = model(inputs)35        probs = torch.softmax(outputs.logits, dim=1)[0]36 37    scores = {id2label[i]: float(probs[i]) for i in range(len(id2label))}38 39    primary = max(scores, key=scores.get)40    safe = primary not in ["hentai", "porn", "sexy"]41 42    return primary, scores, safe43 44 45# ✅ batch version46def scan_batch(images):47    batch = []48 49    for img in images:50        img = cv2.resize(img, (224, 224))51        img = img / 255.052        img = np.transpose(img, (2, 0, 1))53        batch.append(img)54 55    batch = torch.tensor(batch, dtype=torch.float32).to(device)56 57    with torch.no_grad():58        outputs = model(batch)59        probs = torch.softmax(outputs.logits, dim=1)60 61    results = []62 63    for prob in probs:64        scores = {id2label[i]: float(prob[i]) for i in range(len(id2label))}65        primary = max(scores, key=scores.get)66        safe = primary not in ["hentai", "porn", "sexy"]67 68        results.append((primary, scores, safe))69 70    return results