Falconsai

/

nsfw_image_detection

like 326

Follow

Falcons.ai 39

You need to print it out.

The provided sample code do not return scores. This one does:

import requests
from PIL import Image
from transformers import pipeline

# 1. Load the pipeline for image classification
pipe = pipeline("image-classification", model="Falconsai/nsfw_image_detection")

# 2. Load the image into memory (assuming you have the URL for the image)
image_url = 'https://dthezntil550i.cloudfront.net/v3/latest/v32301150603333930008587190/1280_960/c979ecc1-f5d7-4093-8faf-6f879b7d08e6.png'
raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')

# 3. Use the pipeline to classify the image
results = pipe(raw_image)

# 4. Print the results
print(results)
# [{'label': 'nsfw', 'score': 0.998309850692749}, {'label': 'normal', 'score': 0.001690165838226676}]

import torch
from PIL import Image
from transformers import AutoModelForImageClassification, ViTImageProcessor

img = Image.open("images1.jpeg")
model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
with torch.no_grad():
inputs = processor(images=img, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits

predicted_label = logits.argmax(-1).item()
model.config.id2label[predicted_label]
print(outputs)

Output:
ImageClassifierOutput(loss=None, logits=tensor([[-0.7965, 0.8878]]), hidden_states=None, attentions=None)

I believe 0.8878 is the score for nsfw

Please, how to further fine-tune based on this model？