Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -9,9 +9,9 @@ from typing import List, Dict, Tuple
|
|
| 9 |
|
| 10 |
class_names = ["pizza", "steak", "sushi"]
|
| 11 |
|
| 12 |
-
effnetb2,
|
| 13 |
|
| 14 |
-
effnetb2.load_state_dict(torch.load(
|
| 15 |
map_location=torch.device("cpu")) # Load model to the CPU
|
| 16 |
|
| 17 |
def predict(img) -> Tuple[Dict, float]:
|
|
|
|
| 9 |
|
| 10 |
class_names = ["pizza", "steak", "sushi"]
|
| 11 |
|
| 12 |
+
effnetb2, effnetb2_transform = create_effnetb2_model(num_classes=len(class_names))
|
| 13 |
|
| 14 |
+
effnetb2.load_state_dict(torch.load("pretrained_effnetb2_pizza_steak_sushi_20_percent.pt"),
|
| 15 |
map_location=torch.device("cpu")) # Load model to the CPU
|
| 16 |
|
| 17 |
def predict(img) -> Tuple[Dict, float]:
|