Jugal-sheth commited on
Commit
af9778a
·
1 Parent(s): 20e9a65

Add application file

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from model import model, classes
4
+
5
+ # Load state dict
6
+ model.load_state_dict(torch.load('model.pth'))
7
+
8
+ # Transforms
9
+ transform = transforms.Compose([
10
+ transforms.Resize((32, 32)),
11
+ transforms.ToTensor(),
12
+ transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
13
+
14
+ def classify(img):
15
+ x = transform(img)
16
+
17
+ with torch.no_grad():
18
+ pred = model(x)
19
+
20
+ idx = torch.argmax(pred)
21
+ cls = classes[idx]
22
+
23
+ return cls
24
+
25
+
26
+ iface = gr.Interface(...)
27
+ iface.launch()