dmusingu commited on
Commit
2d3a8bd
·
verified ·
1 Parent(s): afb450b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -96,8 +96,14 @@ from torchvision import transforms
96
 
97
  model = load_model()
98
 
 
 
 
 
 
99
  def predict(inp):
100
- inp = transforms.Resize((224, 224))(inp).transforms.ToTensor()(inp).unsqueeze(0)
 
101
  with torch.no_grad():
102
  prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
103
  confidences = {labels[i]: float(prediction[i]) for i in range(3)}
 
96
 
97
  model = load_model()
98
 
99
+ augs = transforms.Compose([
100
+ transforms.Resize((224, 224)),
101
+ transforms.ToTensor()
102
+ ])
103
+
104
  def predict(inp):
105
+ # inp = transforms.Resize((224, 224))(inp).transforms.ToTensor()(inp).unsqueeze(0)
106
+ inp = augs(inp).unsqueeze(0)
107
  with torch.no_grad():
108
  prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
109
  confidences = {labels[i]: float(prediction[i]) for i in range(3)}