Views
No views yet
has_lanyardno_lanyard1import torch
2from torchvision import transforms, models
3from PIL import Image
4
5# Load model
6checkpoint = torch.load('pytorch_model.pth', map_location='cpu')
7model = models.mobilenet_v2()
8model.classifier[1] = torch.nn.Linear(1280, 2)
9model.load_state_dict(checkpoint['model_state_dict'])
10model.eval()
11
12# Preprocess
13transform = transforms.Compose([
14 transforms.Resize((224, 224)),
15 transforms.ToTensor(),
16 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
17])
18
19# Predict
20img = Image.open('test.jpg')
21img_tensor = transform(img).unsqueeze(0)
22
23with torch.no_grad():
24 output = model(img_tensor)
25 probs = torch.softmax(output, dim=1)
26 pred_class = output.argmax(1).item()
27
28classes = ['has_lanyard', 'no_lanyard']
29print(f"Prediction: {classes[pred_class]}")
30print(f"Confidence: {probs[0][pred_class]*100:.1f}%")