Rick-29 commited on
Commit
53b2f95
1 Parent(s): 7e0c76f

initial commit

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ 09_pretrained_effnetb2_feature_extractor_20_percent.pth filter=lfs diff=lfs merge=lfs -text
09_pretrained_effnetb2_feature_extractor_20_percent.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:082946cc6f15df586b385e2c361e0f7384511deddeabcea83dbd04c2c8c71414
3
+ size 31300218
app.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ ### 1. Imports and class names setup
3
+ import gradio as gr
4
+ import torch
5
+ import os
6
+
7
+ from timeit import default_timer as timer
8
+ from model import create_effnetb2_model
9
+ from typing import Tuple, Dict
10
+
11
+ # Setup class names
12
+ class_names = ["pizza", "steak", "sushi"]
13
+
14
+ ### 2. Model and transforms preparation
15
+ effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes = len(class_names))
16
+
17
+ # Load the saved weights
18
+ effnetb2.load_state_dict(torch.load(f="09_pretrained_effnetb2_feature_extractor_20_percent.pth",
19
+ map_location=torch.device("cpu")))
20
+
21
+ ### 3. Predict function
22
+ def predict(img) -> Tuple[Dict, float]:
23
+ # Start a timer
24
+ start_time = timer()
25
+
26
+ # Transform the input image for use with EffNetB2
27
+ img = effnetb2_transforms(img).unsqueeze(0)
28
+
29
+ # Put model into eval mode to make prediction
30
+ effnetb2.eval()
31
+ with torch.inference_mode():
32
+ # Pass transformed image through the model
33
+ pred_probs = torch.softmax(effnetb2(img), dim=1).squeeze()
34
+
35
+ # Create a prediction label and prediction probability dictionary
36
+ pred_labels_and_probs = {food: float(pred_probs[i]) for i, food in enumerate(class_names)}
37
+
38
+ # Calculate pred time
39
+ pred_time = round(timer() - start_time, 4)
40
+
41
+ # Return pred dict and pred time
42
+ return pred_labels_and_probs, pred_time
43
+
44
+ ### 4. Create the Gradio app
45
+ title = "FoodVision Mini🍕🥩🍣"
46
+ description = "An [EfficientNetB2 Feature Extractor](https://pytorch.org/vision/main/models/efficientnet.html#efficientnet_b2) computer vision model to classify images as pizza, steak and sushi."
47
+ article = "Created at [09. Pytorch Model Deployment](https://www.learnpytorch.io/09_pytorch_model_deployment)"
48
+
49
+ # Create example list
50
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
51
+
52
+ # Create the gradio demo
53
+ demo = gr.Interface(fn=predict,
54
+ inputs=gr.Image(type="pil"),
55
+ outputs=[gr.Label(num_top_classes=3, label="Predictions"),
56
+ gr.Number(label="Prediction time (s)")],
57
+ examples=example_list,
58
+ title=title,
59
+ description=description,
60
+ article=article)
61
+
62
+ # Launch the demo
63
+ demo.launch(debug=False, # Print errors locally?
64
+ share=True) # generate a publically available URL
examples/2582289.jpg ADDED
examples/3622237.jpg ADDED
examples/592799.jpg ADDED
model.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+ import torchvision
4
+
5
+ from torch import nn
6
+
7
+ def create_effnetb2_model(num_classes: int = 3,
8
+ seed: int = 42):
9
+ # 1, 2, 3 Create EffNetB2 pretrained weights, transforms and model
10
+ weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
11
+ transforms = weights.transforms()
12
+ model = torchvision.models.efficientnet_b2(weights=weights)
13
+
14
+ # 4. Freeze all layers in the base model
15
+ for param in model.parameters():
16
+ param.requires_grad = False
17
+
18
+ # 5. Change classifier head with random seed for reproducibility
19
+ torch.manual_seed(seed)
20
+ model.classifier = nn.Sequential(
21
+ nn.Dropout(p= .3, inplace=True),
22
+ nn.Linear(in_features=1408, out_features=3, bias=True)
23
+ )
24
+ return model, transforms
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch=2.1.0
2
+ torchvision=0.16.0
3
+ gradio=3.41.0