Spaces:
Sleeping
Sleeping
unknown
commited on
Commit
•
c77b8da
1
Parent(s):
10d2191
fix bugs
Browse files
main.py
CHANGED
@@ -92,12 +92,13 @@ def predict_image(image_url):
|
|
92 |
img_cp = np.float32(img_cp) / 255
|
93 |
input_tensor = preprocess_image(img_cp, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
94 |
input_tensor = torch.Tensor(input_tensor)
|
95 |
-
input_tensor.cuda()
|
96 |
|
97 |
output = model(input_tensor)
|
98 |
# print(torch.max(output, 1))
|
99 |
_, preds_tensor = torch.max(output, 1)
|
100 |
-
preds = np.squeeze(preds, preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
|
|
|
101 |
print(preds)
|
102 |
|
103 |
class_name = classes[preds]
|
@@ -108,7 +109,7 @@ def predict_image(image_url):
|
|
108 |
img = cv2.resize(img, (224, 224))
|
109 |
img = np.float32(img) / 255
|
110 |
input_tensor = torch.Tensor(input_tensor)
|
111 |
-
input_tensor.cuda()
|
112 |
|
113 |
input_tensor = preprocess_image(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
114 |
targets = [ClassifierOutputTarget(0)]
|
|
|
92 |
img_cp = np.float32(img_cp) / 255
|
93 |
input_tensor = preprocess_image(img_cp, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
94 |
input_tensor = torch.Tensor(input_tensor)
|
95 |
+
# input_tensor.cuda()
|
96 |
|
97 |
output = model(input_tensor)
|
98 |
# print(torch.max(output, 1))
|
99 |
_, preds_tensor = torch.max(output, 1)
|
100 |
+
# preds = np.squeeze(preds, preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
|
101 |
+
preds = np.squeeze(preds_tensor.cpu().numpy())
|
102 |
print(preds)
|
103 |
|
104 |
class_name = classes[preds]
|
|
|
109 |
img = cv2.resize(img, (224, 224))
|
110 |
img = np.float32(img) / 255
|
111 |
input_tensor = torch.Tensor(input_tensor)
|
112 |
+
# input_tensor.cuda()
|
113 |
|
114 |
input_tensor = preprocess_image(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
115 |
targets = [ClassifierOutputTarget(0)]
|