import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
from PIL import Image
import torch

model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Mature-Content-Detection")
feature_extractor = AutoFeatureExtractor.from_pretrained("prithivMLmods/Mature-Content-Detection")

def predict(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    labels = model.config.id2label
    return {labels[i]: float(probs[0][i]) for i in range(len(labels))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="Mature Content Detector",
    description="Detects fine-grained categories such as neutral, pornographic, sensual, hentai, etc."
)

if __name__ == "__main__":
    demo.launch()