turtlemb commited on
Commit
45d2464
·
verified ·
1 Parent(s): add2944

Uploading requirements.txt, app.py, model.py, class_names.txt, and 09_pretrained_vit_feature_extractor_food101_20_percent.pth

Browse files
09_pretrained_vit_feature_extractor_food101_20_percent.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fc8be7ecdce20b9c8f788e201fa32d9536f6d033e4ccaaa6204cf98f9e587c16
3
+ size 343571601
app.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### 1. Imports and class names setup ###
2
+ import gradio as gr
3
+ import os
4
+ import torch
5
+
6
+ from model import create_vit_model
7
+ from timeit import default_timer as timer
8
+ from typing import Tuple, Dict
9
+
10
+ # Set up class names
11
+ with open("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 ###
15
+ # Create mode and transforms
16
+ vit, vit_transforms = create_vit_model(num_classes = 101)
17
+
18
+ # Load saved weights
19
+ vit.load_state_dict(
20
+ torch.load(f = "09_pretrained_vit_feature_extractor_food101_20_percent.pth",
21
+ map_location = torch.device("cpu")) # load to CPU
22
+ )
23
+
24
+ ### 3. Predict function ###
25
+ def predict(img) -> Tuple[Dict, float]:
26
+ # Start a timer
27
+ start_time = timer()
28
+
29
+ # Transform the input image for use with ViT
30
+ img = vit_transforms(img).unsqueeze(0) # unsqueeze = add batch dimension on 0th index
31
+
32
+ # Put model into eval mode, make prediction
33
+ vit.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(vit(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 pred 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. Gradio app ###
49
+ # Create title, description, and article
50
+ title = "Big Food Image Classifier 🍔👁️💪"
51
+ description = "A [ViT transformer feature extractor](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html#vit-b-16) computer vision model to classify [101 classes](https://github.com/mrdbourke/pytorch-deep-learning/blob/main/extras/food101_class_names.txt) of food images (from Food101 dataset)."
52
+ article = "Created at [turtlemb's GitHub](https://github.com/turtlemb)."
53
+
54
+ # Create example list
55
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
56
+
57
+ # Create the Gradio demo
58
+ demo = gr.Interface(fn = predict, # maps inputs to outputs
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
+ # Launch the demo
68
+ demo.launch()
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
model.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+ import torchvision
4
+
5
+ from torch import nn
6
+
7
+ def create_vit_model(num_classes: int = 3,
8
+ seed: int = 42):
9
+ # Create ViT_B_16 pre-trained weights, transforms and model
10
+ weights = torchvision.models.ViT_B_16_Weights.DEFAULT
11
+ transforms = weights.transforms()
12
+ model = torchvision.models.vit_b_16(weights = weights)
13
+
14
+ # Freeze all of the base layers
15
+ for param in model.parameters():
16
+ param.requires_grad = False
17
+
18
+ # Change classifier head to suit our needs
19
+ torch.manual_seed(seed)
20
+ model.heads = nn.Sequential(nn.Linear(in_features = 768,
21
+ out_features = num_classes))
22
+
23
+ return model, transforms
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ torch==2.7.1
2
+ torchvision==0.22.1