Views
No views yet
google/vit-large-patch16-224NSFW vs Safe)transformers library.1from transformers import AutoImageProcessor, AutoModelForImageClassification
2from PIL import Image
3import torch
4
5# Load model and processor
6model_name = "theusamaaslam/vit-base-nsfw-detection-224"
7processor = AutoImageProcessor.from_pretrained(model_name)
8model = AutoModelForImageClassification.from_pretrained(model_name)
9
10# Load your image
11image = Image.open("path/to/your/image.jpg")
12
13# Preprocess
14inputs = processor(images=image, return_tensors="pt")
15
16# Inference
17with torch.no_grad():
18 outputs = model(**inputs)
19 logits = outputs.logits
20 predicted_class_idx = logits.argmax(-1).item()
21
22# Get label
23label = model.config.id2label[predicted_class_idx]
24print(f"Prediction: {label}")1
2from transformers import AutoImageProcessor, AutoModelForImageClassification
3from PIL import Image
4import requests
5import torch
6
7# URL of the image to test
8url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
9
10# Load image from URL
11image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
12
13# Load Processor and Model from your Hub repo
14repo_id = "theusamaaslam/vit-base-nsfw-detection-224"
15processor = AutoImageProcessor.from_pretrained(repo_id)
16model = AutoModelForImageClassification.from_pretrained(repo_id)
17
18# Process image and inference
19inputs = processor(images=image, return_tensors="pt")
20
21with torch.no_grad():
22 outputs = model(**inputs)
23
24# Get probabilities and predicted label
25logits = outputs.logits
26probs = logits.softmax(dim=1)
27predicted_class_idx = logits.argmax(-1).item()
28predicted_label = model.config.id2label[predicted_class_idx]
29
30print(f"Predicted class: {predicted_label}")
31print(f"Probabilities: {probs}")