back0p commited on
Commit
5d3d2ca
·
1 Parent(s): 4324dc8

intial commit

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
+ pretrained_effnetb2_pizza_steak_sushi_20_percent.pt filter=lfs diff=lfs merge=lfs -text
37
+ .pt filter=lfs diff=lfs merge=lfs -text
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import torch
4
+ import gradio as gr
5
+
6
+ from model import create_effnetb2_model
7
+ from timeit import default_timer as timer
8
+ from typing import List, Dict, Tuple
9
+
10
+ class_names = ["pizza", "steak", "sushi"]
11
+
12
+ effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=len(class_names))
13
+
14
+ effnetb2.load_state_dict(torch.load(f="pretrained_effnetb2_pizza_steak_sushi_20_percent.pt"),
15
+ map_location=torch.device("cpu")) # Load model to the CPU
16
+
17
+ def predict(img) -> Tuple[Dict, float]:
18
+ start_time = timer()
19
+
20
+ img = effnetb2_transform(img).unsqueeze(0)
21
+
22
+ effnetb2.eval()
23
+ with torch.inference_mode():
24
+ pred_probs = torch.softmax(effnetb2(img), dim=1)
25
+
26
+ pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
27
+
28
+ end_time = timer()
29
+
30
+ return pred_labels_and_probs, round(end_time-start_time, 4)
31
+
32
+ title = "FoodVision Mini"
33
+ description = "An [EfficientNetB2 feature extractor](www.google.com) computer vision model to classify images as pizza, steak, sushi"
34
+
35
+ example_list = [["examples/" + example] for example in os.listdir("examples")]
36
+
37
+ demo = gr.Interface(fn=predict,
38
+ inputs=gr.Image(type="pil"),
39
+ outputs=[gr.Label(num_top_classes=3, label="Predictions"),
40
+ gr.Number(label="Prediction time (s)")],
41
+ examples=example_list,
42
+ title=title,
43
+ description=description)
44
+
45
+ demo.launch(debug=False)
examples/2582289.jpg ADDED
examples/3622237.jpg ADDED
examples/592799.jpg ADDED
model.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+ import torchvision
4
+ from torch import nn
5
+
6
+ def create_effnetb2_model(num_classes: int = 3,
7
+ seed: int = 42):
8
+ weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
9
+ transforms = weights.transforms()
10
+ model = torchvision.models.efficientnet_b2(weights=weights).to(device)
11
+
12
+ for param in model.parameters():
13
+ param.requires_grad = False
14
+
15
+ torch.manual_seed(seed)
16
+ torch.cuda.manual_seed(seed)
17
+ model.classifier = nn.Sequential(
18
+ nn.Dropout(p=0.3, inplace=True),
19
+ nn.Linear(in_features=1408, out_features=num_classes, bias=True)
20
+ ).to(device)
21
+ return model, transforms
pretrained_effnetb2_pizza_steak_sushi_20_percent.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:05fec3d6ac9082ddfed66b4ed747ddaf4779e21b50186dd6636ceb13168528b0
3
+ size 31299161
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+
2
+ torch==2.9.0
3
+ torchvision==0.24.0
4
+ gradio==6.0.0