Views
No views yet
1from efficientnet_pytorch import EfficientNet
2import torch
3import torchvision.transforms as transforms
4
5model = EfficientNet.from_name('efficientnet-b7')
6model._fc= torch.nn.Linear(in_features=model._fc.in_features, out_features=len(annotations_map), bias=True)
7model.load_state_dict(torch.load('/content/efficientnetb7_tyrequality_classifier.pth'))
8
9model.eval()
10img = Image.open('/content/defective-tires-cause-accidents-min.jpg')
11test_transform = transforms.Compose([
12 transforms.Resize(224),
13 transforms.CenterCrop(224),
14 transforms.ToTensor(),
15 transforms.Normalize([0.485, 0.456, 0.406],
16 [0.229, 0.224, 0.225])
17])
18input_data = test_transform(img).unsqueeze(0)
19
20with torch.no_grad():
21 output = model(input_data)
22
23_, predicted_class = torch.max(output, 1)
24
25probs = torch.nn.functional.softmax(output, dim=1)
26conf, _ = torch.max(probs, 1)
27
28print('Predicted Class:', predicted_class.item())
29print('Predicted Label:', id2label[predicted_class.item()])
30print(f'Confidence: {conf.item()*100}%')
31
32plt.title(id2label[predicted_class.item()])
33plt.axis("off")
34plt.imshow(img)
35plt.show()