app.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from PIL import Image
import io
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

app = FastAPI(title="image-moderation-api")

print("Loading Hugging Face model into RAM...")
model_name = "prithivMLmods/Nsfw_Image_Detection_OSS"

processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
model.eval()

print("Model loaded and ready!")

LABELS = {0: "SFW", 1: "NSFW"}

@app.post("/classify")
async def classify_image(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File provided is not an image.")

    try:
        image_bytes = await file.read()
        
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        inputs = processor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
        
        results = [{"label": LABELS[i], "score": round(probs[i], 4)} for i in range(len(probs))]
        
        results.sort(key=lambda x: x["score"], reverse=True)

        return {
            "filename": file.filename,
            "predictions": results
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")