ayusk commited on
Commit
c93d9e3
1 Parent(s): d0059ed

food classifier with 101 classes

Browse files
app.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ # 1.
3
+ import gradio as gr
4
+ import os
5
+ import torch
6
+
7
+ from model import create_effnetb2_model
8
+ from timeit import default_timer as timer
9
+ from typing import Tuple, Dict
10
+
11
+ # Setup class names
12
+ with open("class_names.txt", "r") as f:
13
+ class_names = [food.strip('\n') for food in f.readlines()]
14
+
15
+ # 2.
16
+ effnetb2_food101, effnetb2_transforms = create_effnetb2_model(num_classes = 101)
17
+
18
+ # Load saved weights
19
+ effnetb2_food101.load_state_dict(torch.load("models/state_dict__effnetb2_food101_20_percent.pth",
20
+ map_location = torch.device('cpu')))
21
+
22
+
23
+ # 3.
24
+ def predict(img) -> Tuple[Dict, float]:
25
+
26
+ # Start a timer
27
+ start_time = timer()
28
+
29
+ # Transform the input image for use with EffNetB2
30
+ img = effnetb2_transforms(img).unsqueeze(0)
31
+
32
+ # Put the model into eval mode, make prediction
33
+ effnetb2_food101.eval()
34
+ with torch.inference_mode():
35
+ # Pass transformed image through the model and turn the prediction logits into probabilities
36
+ pred_probs = torch.softmax(effnetb2_food101(img), dim = 1)
37
+
38
+ # Create a prediction label and prediction probability dictionary
39
+ pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
40
+
41
+ # Calculate pre time
42
+ end_time = timer()
43
+ pred_time = round(end_time - start_time, 4)
44
+
45
+ # Return pred dict and pred time
46
+ return pred_labels_and_probs, pred_time
47
+
48
+ # 4.
49
+ title = 'FoodIdentifier Big (a little) 🍣🍕🥩'
50
+ description = "An EfficientNetB2 feature extractor computer vision model to classify images as pizza, sushi or steak"
51
+ article = " anything I want for the description of the description above 🤪"
52
+
53
+ # Create example list
54
+ # Get example filepaths in a list of lists
55
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
56
+
57
+ # Create the Gradio demo
58
+ demo = gr.Interface(fn = predict, # maps input to output
59
+ inputs = gr.Image(type = 'pil'),
60
+ outputs = [gr.Label(num_top_classes = 5, label = "Predictions"),
61
+ gr.Number(label = "Prediction time (s)")],
62
+ examples = example_list,
63
+ title = title,
64
+ description = description,
65
+ article = article
66
+ )
67
+
68
+ # Launch the demo
69
+ demo.launch(debug = False, # print errors locally?
70
+ share = True) # generate a publically shareable URL
class_names.txt ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ apple_pie
2
+ baby_back_ribs
3
+ baklava
4
+ beef_carpaccio
5
+ beef_tartare
6
+ beet_salad
7
+ beignets
8
+ bibimbap
9
+ bread_pudding
10
+ breakfast_burrito
11
+ bruschetta
12
+ caesar_salad
13
+ cannoli
14
+ caprese_salad
15
+ carrot_cake
16
+ ceviche
17
+ cheese_plate
18
+ cheesecake
19
+ chicken_curry
20
+ chicken_quesadilla
21
+ chicken_wings
22
+ chocolate_cake
23
+ chocolate_mousse
24
+ churros
25
+ clam_chowder
26
+ club_sandwich
27
+ crab_cakes
28
+ creme_brulee
29
+ croque_madame
30
+ cup_cakes
31
+ deviled_eggs
32
+ donuts
33
+ dumplings
34
+ edamame
35
+ eggs_benedict
36
+ escargots
37
+ falafel
38
+ filet_mignon
39
+ fish_and_chips
40
+ foie_gras
41
+ french_fries
42
+ french_onion_soup
43
+ french_toast
44
+ fried_calamari
45
+ fried_rice
46
+ frozen_yogurt
47
+ garlic_bread
48
+ gnocchi
49
+ greek_salad
50
+ grilled_cheese_sandwich
51
+ grilled_salmon
52
+ guacamole
53
+ gyoza
54
+ hamburger
55
+ hot_and_sour_soup
56
+ hot_dog
57
+ huevos_rancheros
58
+ hummus
59
+ ice_cream
60
+ lasagna
61
+ lobster_bisque
62
+ lobster_roll_sandwich
63
+ macaroni_and_cheese
64
+ macarons
65
+ miso_soup
66
+ mussels
67
+ nachos
68
+ omelette
69
+ onion_rings
70
+ oysters
71
+ pad_thai
72
+ paella
73
+ pancakes
74
+ panna_cotta
75
+ peking_duck
76
+ pho
77
+ pizza
78
+ pork_chop
79
+ poutine
80
+ prime_rib
81
+ pulled_pork_sandwich
82
+ ramen
83
+ ravioli
84
+ red_velvet_cake
85
+ risotto
86
+ samosa
87
+ sashimi
88
+ scallops
89
+ seaweed_salad
90
+ shrimp_and_grits
91
+ spaghetti_bolognese
92
+ spaghetti_carbonara
93
+ spring_rolls
94
+ steak
95
+ strawberry_shortcake
96
+ sushi
97
+ tacos
98
+ takoyaki
99
+ tiramisu
100
+ tuna_tartare
101
+ waffles
examples/2250611.jpg ADDED
examples/33860.jpg ADDED
examples/911808.jpg ADDED
examples/steak.jpg ADDED
model.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torchvision
3
+
4
+ from torch import nn
5
+
6
+
7
+ def create_effnetb2_model(num_classes:int=3,
8
+ seed:int=42):
9
+ """Creates an EfficientNetB2 feature extractor model and transforms.
10
+
11
+ Args:
12
+ num_classes (int, optional): number of classes in the classifier head.
13
+ Defaults to 3.
14
+ seed (int, optional): random seed value. Defaults to 42.
15
+
16
+ Returns:
17
+ model (torch.nn.Module): EffNetB2 feature extractor model.
18
+ transforms (torchvision.transforms): EffNetB2 image transforms.
19
+ """
20
+ # Create EffNetB2 pretrained weights, transforms and model
21
+ weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
22
+ transforms = weights.transforms()
23
+ model = torchvision.models.efficientnet_b2(weights=weights)
24
+
25
+ # Freeze all layers in base model
26
+ for param in model.parameters():
27
+ param.requires_grad = False
28
+
29
+ # Change classifier head with random seed for reproducibility
30
+ torch.manual_seed(seed)
31
+ model.classifier = nn.Sequential(
32
+ nn.Dropout(p=0.3, inplace=True),
33
+ nn.Linear(in_features=1408, out_features=num_classes),
34
+ )
35
+
36
+ return model, transforms
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+
2
+ torch>1.12.0
3
+ torchvision>0.13.0
4
+ gradio>3.1.4