ditobprasetio commited on
Commit
6bda815
·
1 Parent(s): 775a60f
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -12,7 +12,7 @@ def predict(image):
12
 
13
  model = EfficientNet.from_name('efficientnet-b7', num_classes=2)
14
  model_weights_path = 'efficientnetb7_tyrequality_classifier.pth'
15
- model.load_state_dict(torch.load(model_weights_path))
16
  model.eval()
17
 
18
  transform = transforms.Compose([
@@ -34,7 +34,7 @@ def predict(image):
34
  _, predicted_class = torch.max(output, 1)
35
 
36
  probs = torch.nn.functional.softmax(output, dim=1)
37
- print(probs, "probs")
38
  conf, _ = torch.max(probs, 1)
39
 
40
  result = "Tire status is {} with confidence level in {}%".format(class_to_label[predicted_class.item()], conf.item()*100)
 
12
 
13
  model = EfficientNet.from_name('efficientnet-b7', num_classes=2)
14
  model_weights_path = 'efficientnetb7_tyrequality_classifier.pth'
15
+ model.load_state_dict(torch.load(model_weights_path, map_location=torch.device('cpu')))
16
  model.eval()
17
 
18
  transform = transforms.Compose([
 
34
  _, predicted_class = torch.max(output, 1)
35
 
36
  probs = torch.nn.functional.softmax(output, dim=1)
37
+ # print(probs, "probs")
38
  conf, _ = torch.max(probs, 1)
39
 
40
  result = "Tire status is {} with confidence level in {}%".format(class_to_label[predicted_class.item()], conf.item()*100)