GeorgeImmanuel commited on
Commit
4b1b927
1 Parent(s): eaad07f

things are added

Browse files
Files changed (4) hide show
  1. app.py +81 -0
  2. model.py +19 -0
  3. requirements.txt +3 -0
  4. vit_b_16_20_percent_data.pth +3 -0
app.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ # import the essentials
3
+ from demos.foodvision_mini.model import create_vit_b_16_model
4
+ import torch
5
+ import torchvision
6
+ import time
7
+ import gradio as gr
8
+ import numpy as np
9
+ from pathlib import Path
10
+
11
+ class_names = ['pizza','steak','sushi']
12
+ device = 'cuda' if torch.cuda.is_available else 'cpu'
13
+
14
+ # creating the vit_b_16_model and loading it with state_dict of our trained model
15
+ vit_b_16_model,vit_b_16_transform = create_vit_b_16_model(num_classes=3)
16
+ vit_b_16_model.load_state_dict(torch.load(f='vit_b_16_20_percent_data.pth'))
17
+
18
+ # create the predict function
19
+ def predict(img):
20
+
21
+ """
22
+ args:
23
+ img: is an image
24
+
25
+ returns: prediction class, prediction probability, and time taken to make the prediction
26
+
27
+ """
28
+
29
+ # transforming the image
30
+ tr_img = vit_b_16_transform(img).unsqueeze(dim=0).to(device)
31
+
32
+ # make prediction with vit_b_16
33
+ model = vit_b_16_model.to(device)
34
+
35
+ # starting the time
36
+ start_time = time.perf_counter()
37
+
38
+ model.eval()
39
+ with torch.inference_mode():
40
+ pred_logit = model(tr_img)
41
+ pred_label = torch.argmax(pred_logit,dim=1).cpu()
42
+ pred_prob = torch.max(torch.softmax(pred_logit,dim=1)).cpu().item()
43
+
44
+ # ending the time
45
+ end_time = time.perf_counter()
46
+ # pred_dict = {str(class_names[i]):float(pred_prob[0][i].item()) for i in range(len(class_names))}
47
+ pred_prob = float(np.round(pred_prob,3))
48
+ pred_class = class_names[pred_label]
49
+ time_taken = float(np.round(end_time-start_time,3))
50
+
51
+
52
+
53
+ return pred_class,pred_prob,time_taken
54
+
55
+
56
+ # create example list
57
+ example_dir = Path('demos/foodvision_mini/examples')
58
+ example_list = [['examples/' + str(filepath)] for filepath in os.listdir(example_dir)]
59
+
60
+ # create Gradio interface
61
+ description = 'A machine learning model to classify images into pizza,steak and sushi appropriately'
62
+ title = 'Image Classifier'
63
+
64
+
65
+ demo = gr.Interface(fn=predict, # this function maps the inputs to the output
66
+ inputs=gr.Image(type='pil'), # pillow image
67
+ outputs=[gr.Label(num_top_classes=1,label='Prediction'),
68
+ gr.Number(label='prediction probability'),
69
+ gr.Number(label='prediction time(s)')],
70
+ examples=example_list,
71
+ description=description,
72
+ title=title
73
+ )
74
+
75
+ demo.launch(debug=False, # print errors locally?
76
+ share=True) # share to the public?
77
+
78
+
79
+
80
+
81
+
model.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torchvision
3
+ import torch
4
+ def create_vit_b_16_model(num_classes=3):
5
+
6
+ weights = torchvision.models.ViT_B_16_Weights.DEFAULT
7
+ transform = torchvision.models.ViT_B_16_Weights.DEFAULT.transforms()
8
+ model = torchvision.models.vit_b_16(weights=weights)
9
+
10
+ # freeze the layers
11
+ for param in model.parameters():
12
+ param.requires_grad = False
13
+
14
+ # modify the heads layer
15
+ model.heads = torch.nn.Sequential(
16
+ torch.nn.Linear(in_features=768,out_features=num_classes)
17
+ )
18
+
19
+ return model,transform
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch>=1.12.0
2
+ torchvision>=0.13.0
3
+ gradio>=3.1.4
vit_b_16_20_percent_data.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:72dbf7ad12f2bf4ce1aedab9105ba670c9a7ba6fae54aba670b77700e954a8ae
3
+ size 343266910