vapit commited on
Commit
32e94d8
Β·
1 Parent(s): 47a054c

add the final scripts

Browse files
.gitattributes CHANGED
@@ -33,3 +33,5 @@ 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_vit_feature_extractor_pizza_steak_sushi_20_percent.pth filter=lfs diff=lfs merge=lfs -text
37
+ .pth filter=lfs diff=lfs merge=lfs -text
09_pretrained_vit_feature_extractor_pizza_steak_sushi_20_percent.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e3c46b8bce9fea760b9e2891199c70a41e577c40dba131904599c1ed9d43e757
3
+ size 343273342
app.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### 1. Imports and class names setup ###
2
+ from model import create_vitB16_model
3
+ import torch
4
+ from typing import Tuple, Dict
5
+ from timeit import default_timer as timer
6
+ import gradio as gr
7
+
8
+ # Setup class names
9
+ class_names = ['pizza', 'steak', 'sushi']
10
+
11
+ ### 2. Model and transforms perparation ###
12
+ model, model_transforms = create_vitB16_model(num_classes=len(class_names))
13
+
14
+ # Load save weights
15
+ model.load_state_dict(torch.load(f='09_pretrained_vit_feature_extractor_pizza_steak_sushi_20_percent.pth',
16
+ map_location='cpu'))
17
+
18
+ # 3. Predict Function
19
+
20
+ def predict(img) -> Tuple[Dict, float]:
21
+ # Start a timer
22
+ start_time = timer()
23
+
24
+ # Transform the input image for use with vitB16
25
+ img = model_transforms(img).unsqueeze(dim=0)
26
+
27
+ # Put model into eval mode, make prediction
28
+ model.eval()
29
+ with torch.inference_mode():
30
+ # Pass transformed image through the model and turn the prediction logits into probabilities
31
+ pred_logit = model(img)
32
+ pred_prob = torch.softmax(pred_logit, dim=1)
33
+
34
+ # Create a prediction label and prediction probability dictionary
35
+ pred_labels_and_probs = {class_names[i]: float(pred_prob[0][i]) for i in range(len(class_names))}
36
+
37
+ # Calculate pred time
38
+ end_time = timer()
39
+ pred_time = round(end_time - start_time, 4)
40
+
41
+ # Return pred dict and pred time
42
+ return pred_labels_and_probs, pred_time
43
+
44
+ ### 4. Gradio app ###
45
+
46
+
47
+ # Create title, description and article
48
+ title = "FoodVision Mini πŸ•πŸ₯©πŸ£"
49
+ description = "A [vision Transformer B16 feature extractor](https://pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_16.html) computer vision model to classify images as pizza, steak or sushi."
50
+ article = "Created with 🀎 (and a mixture of mathematics, statistics, and tons of calculations πŸ‘©πŸ½β€πŸ”¬)"
51
+
52
+ # Create example list
53
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
54
+
55
+ demo = gr.Interface(fn=predict,
56
+ inputs=gr.Image(type='pil'),
57
+ outputs=[gr.Label(num_top_classes=3, label='Predictions'),
58
+ gr.Number(label="Prediction time (s)")],
59
+ examples=example_list,
60
+ title=title,
61
+ description=description,
62
+ article=article)
63
+
64
+ # Launch the demo!
65
+ demo.launch(debug=False, # print errors locally?
66
+ share=True) # generate a publically shareable URL
examples/2582289.jpg ADDED
examples/3622237.jpg ADDED
examples/592799.jpg ADDED
model.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch, torchvision
2
+ def create_vitB16_model(num_classes: int=3, seeds: int = 42):
3
+
4
+ # 1. Setup pretrained viT Weights
5
+ weights = torchvision.models.ViT_B_16_Weights.DEFAULT
6
+
7
+ # 2. Get transforms
8
+ transforms = weights.transforms()
9
+
10
+ # 3. Setup pretrained instance
11
+ model = torchvision.models.vit_b_16(weights=weights)
12
+
13
+ # 4. Freeze the base layers in the model (this will stop all layers from training)
14
+ for params in model.parameters():
15
+ params.requires_grad = False
16
+
17
+ # Set seeds for reproducibility
18
+ torch.manual_seed(seeds)
19
+
20
+ # 5. Modify the number of output layers
21
+ model.heads = torch.nn.Sequential(
22
+ torch.nn.Linear(in_features=768, out_features=num_classes, bias=True)
23
+ )
24
+
25
+ return model, transforms
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch==2.4.0
2
+ torchvision==0.19.0
3
+ gradio==4.44.0