File size: 1,596 Bytes
dff7e68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import torch
from transformers import AutoModel, AutoImageProcessor
from model import DinoV3LinearMultiLinear
def load_model(weights_path, device="cuda"):
"""
Load the pre-trained classifier.
Args:
weights_path: Path to the saved weights (.pt file)
device: Device to load model on ('cuda' or 'cpu')
Returns:
model: Loaded DinoV3LinearMultiLinear model in eval mode
processor: Image processor for preprocessing input images
"""
# Load config
import json
with open("config.json", "r") as f:
config = json.load(f)
# Load backbone
backbone = AutoModel.from_pretrained(config["model_name"])
hidden_size = backbone.config.hidden_size
# Instantiate classifier head
model = DinoV3LinearMultiLinear(
backbone=backbone,
num_classes=config["num_classes"],
hidden_size=hidden_size,
freeze_backbone=True
)
# Load trained weights
model.load_state_dict(torch.load(weights_path, map_location=device)["model_state_dict"])
model.to(device)
model.eval()
# Load image processor
processor = AutoImageProcessor.from_pretrained(config["model_name"])
# Load labels
with open("id2label.json", "r") as f:
id2label = json.load(f)
return model, processor, id2label
def probs_to_labels(probs, id2label):
"""
Convert probability distribution to labels.
"""
predicted_indices = probs.argmax(dim=1)
predicted_labels = [id2label[str(idx.item())] for idx in predicted_indices]
return predicted_labels |