Queue-Tip commited on
Commit
57942f6
·
verified ·
1 Parent(s): 7a1847a

Create Simple Image Classification

Browse files
Files changed (1) hide show
  1. Simple Image Classification +28 -0
Simple Image Classification ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Import necessary libraries
2
+ from transformers import ViTFeatureExtractor, ViTForImageClassification
3
+ from PIL import Image
4
+ import requests
5
+ import torch
6
+ import matplotlib.pyplot as plt
7
+
8
+ # Load pre-trained feature extractor and model
9
+ feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
10
+ model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
11
+
12
+ # Load and display the image
13
+ url = "URL_of_the_image"
14
+ image = Image.open(requests.get(url, stream=True).raw)
15
+ plt.imshow(image)
16
+ plt.show()
17
+
18
+ # Extract features from the image
19
+ inputs = feature_extractor(images=image, return_tensors="pt")
20
+
21
+ # Make predictions
22
+ outputs = model(**inputs)
23
+ logits = outputs.logits
24
+ predicted_class_idx = logits.argmax(-1).item()
25
+
26
+ # Get and print the predicted class name
27
+ predicted_class = model.config.id2label[predicted_class_idx]
28
+ print(f'Predicted class: {predicted_class}')