Views
No views yet
1from peft import PeftConfig, PeftModel
2from transformers import AutoModelForImageClassification, AutoImageProcessor
3import torch
4from PIL import Image
5import requests
6
7repo_name = f"wtnan2003/vit-base-patch16-224-in21k-finetuned-lora-garbage_classification"
8label2id = {
9 "cardboard":0,
10 "glass":1,
11 "metal":2,
12 "paper":3,
13 "plastic":4,
14 "trash":5
15}
16id2label = {value:key for key, value in label2id.items()}
17config = PeftConfig.from_pretrained(repo_name)
18model = AutoModelForImageClassification.from_pretrained(
19 config.base_model_name_or_path,
20 label2id=label2id,
21 id2label=id2label,
22 ignore_mismatched_sizes=True,
23)
24# Load the LoRA model
25inference_model = PeftModel.from_pretrained(model, repo_name)
26url = "https://www.uky.edu/facilities/sites/www.uky.edu.facilities/files/Cardboard%20Image.png"
27# url = "https://th.bing.com/th/id/OIP.BkzhM2nwEy1edmV7WvU4EAHaJ4?pid=ImgDet&rs=1https://i.redd.it/01msg69otvl21.jpg" # glass
28image = Image.open(requests.get(url, stream=True).raw)
29image_processor = AutoImageProcessor.from_pretrained(repo_name)
30encoding = image_processor(image.convert("RGB"), return_tensors="pt")
31with torch.no_grad():
32 outputs = inference_model(**encoding)
33 logits = outputs.logits
34
35predicted_class_idx = logits.argmax(-1).item()
36print("Predicted class:", inference_model.config.id2label[predicted_class_idx])
37#Predicted class: cardboard