Views
No views yet
google/vit-base-patch16-224| Label | ID | Meaning |
|---|---|---|
| reject | 0 | Image not suitable for catalog |
| accept | 1 | Image suitable for catalog |
pip install transformers torch1import torch
2from transformers import AutoImageProcessor, AutoModelForImageClassification
3
4# Load base ViT model
5model = AutoModelForImageClassification.from_pretrained(
6 "google/vit-base-patch16-224",
7 num_labels=2,
8 id2label={0: "reject", 1: "accept"},
9 label2id={"reject": 0, "accept": 1},
10 ignore_mismatched_sizes=True
11)
12
13# Load and apply head weights
14head_weights = torch.load("head.pt", map_location="cpu")
15model.classifier.weight = head_weights["classifier.weight"]
16model.classifier.bias = head_weights["classifier.bias"]
17
18model.eval()
19print("Model ready!")1from PIL import Image
2from transformers import AutoImageProcessor
3
4processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
5
6# Load your image
7image = Image.open("your_bag_image.jpg").convert("RGB")
8
9# Preprocess
10inputs = processor(images=image, return_tensors="pt")
11
12# Predict
13import torch
14with torch.no_grad():
15 outputs = model(**inputs)
16 predicted_id = outputs.logits.argmax(-1).item()
17 label = {0: "reject", 1: "accept"}[predicted_id]
18 confidence = torch.softmax(outputs.logits, dim=-1).max().item()
19
20print(f"Prediction: {label}")
21print(f"Confidence: {confidence:.2%}")| Parameter | Value |
|---|---|
| Base model | google/vit-base-patch16-224 |
| Epochs | 5 |
| Batch size | 16 |
| Image size | 600x600 (thumbnail) |
| Optimizer | AdamW (default) |
| Test accuracy | 98% |
malaika16/bag-labels