Spaces:
Sleeping
Sleeping
Fixed issues with class_names.txt in app.py
Browse files
app.py
CHANGED
@@ -8,7 +8,7 @@ from model import create_effnetb2_model
|
|
8 |
from timeit import default_timer as timer
|
9 |
|
10 |
# Setup class names
|
11 |
-
with open('
|
12 |
class_names = [food_name.strip() for food_name in f.readlines()]
|
13 |
|
14 |
### 2. Model and transforms preparation ###
|
@@ -18,7 +18,7 @@ effnetb2, effnetb2_transforms = create_effnetb2_model(
|
|
18 |
)
|
19 |
|
20 |
# Load save weights
|
21 |
-
effnetb2.load_state_dict(torch.load('
|
22 |
|
23 |
### 3. Predict function ###
|
24 |
def predict(img) -> tuple[dict, float]:
|
|
|
8 |
from timeit import default_timer as timer
|
9 |
|
10 |
# Setup class names
|
11 |
+
with open('foodvision_big/class_names.txt', 'r') as f:
|
12 |
class_names = [food_name.strip() for food_name in f.readlines()]
|
13 |
|
14 |
### 2. Model and transforms preparation ###
|
|
|
18 |
)
|
19 |
|
20 |
# Load save weights
|
21 |
+
effnetb2.load_state_dict(torch.load('foodvision_big/09_pretrained_effnetb2_feature_extractor_food101_20_percent.pth', map_location=torch.device("cpu"))) # load the model to the CPU
|
22 |
|
23 |
### 3. Predict function ###
|
24 |
def predict(img) -> tuple[dict, float]:
|