Spaces:
Runtime error
Runtime error
GeorgeImmanuel
commited on
Commit
•
4b1b927
1
Parent(s):
eaad07f
things are added
Browse files- app.py +81 -0
- model.py +19 -0
- requirements.txt +3 -0
- 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
|