Views
No views yet
1import torch
2from transformers import ViTForImageClassification, ViTFeatureExtractor
3from PIL import Image
4
5# Load the model and feature extractor
6model = ViTForImageClassification.from_pretrained('jaranohaal/vit-base-violence-detection')
7feature_extractor = ViTFeatureExtractor.from_pretrained('jaranohaal/vit-base-violence-detection')
8
9# Load an image
10image = Image.open('image.jpg')
11
12# Preprocess the image
13inputs = feature_extractor(images=image, return_tensors="pt")
14
15# Perform inference
16with torch.no_grad():
17 outputs = model(**inputs)
18 logits = outputs.logits
19 predicted_class_idx = logits.argmax(-1).item()
20
21# Print the predicted class
22print("Predicted class:", model.config.id2label[predicted_class_idx])