vapit commited on
Commit
0c08f99
Β·
1 Parent(s): d9c763c

first 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_vit_feature_extractor_cifar_60_percent.pth filter=lfs diff=lfs merge=lfs -text
09_pretrained_vit_feature_extractor_cifar_60_percent.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7e7470fb73a7052e71ad0c8eb70352270591c619085984996e6d5b8cc6ebe072
3
+ size 343292910
app.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### 1. Imports and class names setup ###
2
+ from model import create_vitB16_model
3
+
4
+ import torch
5
+ from typing import Tuple, Dict
6
+ from timeit import default_timer as timer
7
+ import gradio as gr
8
+ import os
9
+
10
+ class_names = []
11
+
12
+ # Open Food101 class names file and read each line into a list
13
+ with open(class_names_file, 'r') as f:
14
+ for idx, class_name in enumerate(f):
15
+ class_names[idx] = class_name.strip()
16
+
17
+ ### 2. Model and transforms perparation ###
18
+ model, model_transforms = create_vitB16_model(num_classes=len(class_names))
19
+
20
+ # Load save weights
21
+ model.load_state_dict(torch.load(f='09_pretrained_vit_feature_extractor_food101_20_percent.pth',
22
+ map_location='cpu'))
23
+
24
+ # 3. Predict Function
25
+
26
+ def predict(img) -> Tuple[Dict, float]:
27
+ # Start a timer
28
+ start_time = timer()
29
+
30
+ # Transform the input image for use with vitB16
31
+ img = model_transforms(img).unsqueeze(dim=0)
32
+
33
+ # Put model into eval mode, make prediction
34
+ model.eval()
35
+ with torch.inference_mode():
36
+ # Pass transformed image through the model and turn the prediction logits into probabilities
37
+ pred_logit = model(img)
38
+ pred_prob = torch.softmax(pred_logit, dim=1)
39
+
40
+ # Create a prediction label and prediction probability dictionary
41
+ pred_labels_and_probs = {class_names[i]: float(pred_prob[0][i]) for i in range(len(class_names))}
42
+
43
+ # Calculate pred time
44
+ end_time = timer()
45
+ pred_time = round(end_time - start_time, 4)
46
+
47
+ # Return pred dict and pred time
48
+ return pred_labels_and_probs, pred_time
49
+
50
+ ### 4. Gradio app ###
51
+
52
+
53
+ # Create title, description and article
54
+ title = "Object Detection πŸ€–"
55
+ 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 10 classes of Objects from the Cifar10 dataset. [Cifar10](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html?highlight=cifar)"
56
+ article = "Created with 🀎 (and a mixture of mathematics, statistics, and tons of calculations πŸ‘©πŸ½β€πŸ”¬) by Arpit Vaghela [GitHub](https://github.com/magnifiques)"
57
+
58
+ # Create example list
59
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
60
+
61
+ demo = gr.Interface(fn=predict,
62
+ inputs=gr.Image(type='pil'),
63
+ outputs=[gr.Label(num_top_classes=10, label='Predictions'),
64
+ gr.Number(label="Prediction time (s)")],
65
+ examples=example_list,
66
+ title=title,
67
+ description=description,
68
+ article=article)
69
+
70
+ # Launch the demo!
71
+ demo.launch(debug=False, # print errors locally?
72
+ share=True) # generate a publically shareable URL
class_names.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ airplane
2
+ automobile
3
+ bird
4
+ cat
5
+ deer
6
+ dog
7
+ frog
8
+ horse
9
+ ship
10
+ truck
examples/1.png ADDED
examples/2.jpg ADDED
examples/3.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