virender74 commited on
Commit
75fbcd5
1 Parent(s): 4b1406d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +62 -0
app.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import copy
3
+ import torch
4
+ import gradio
5
+ import gradio as gr
6
+ from PIL import Image
7
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
8
+
9
+ os.system("wget https://www.dropbox.com/s/grcragozd4x79zc/model_ok.pth")
10
+
11
+ model = torch.load("./model_ok.pth", map_location=device)
12
+
13
+ # img = Image.open(path).convert('RGB')
14
+ from torchvision import transforms
15
+
16
+ transforms2 = transforms.Compose([
17
+ transforms.Resize(256),
18
+ transforms.ToTensor(),
19
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
20
+ ])
21
+
22
+ # img = transforms(img)
23
+ # img = img.unsqueeze(0)
24
+ model.eval()
25
+
26
+ labels = ['aunts','bees']
27
+
28
+ # with torch.no_grad():
29
+ # # preds =
30
+ # preds = model(img)
31
+ # score, indices = torch.max(preds, 1)
32
+
33
+ def recognize_digit(image):
34
+ image = transforms2(image)
35
+ image = image.unsqueeze(0)
36
+ # image = image.unsqueeze(0)
37
+ # image = image.reshape(1, -1)
38
+ # with torch.no_grad():
39
+ # preds =
40
+ # img = image.reshape((-1, 3, 256, 256))
41
+ preds = model(image).flatten()
42
+ # prediction = model.predict(image).tolist()[0]
43
+ # score, indices = torch.max(preds, 1)
44
+ # return {str(indices.item())}
45
+ return {labels[i]: float(preds[i]) for i in range(2)}
46
+
47
+
48
+ im = gradio.inputs.Image(
49
+ shape=(256, 256), image_mode="RGB", type="pil")
50
+
51
+ iface = gr.Interface(
52
+ recognize_digit,
53
+ im,
54
+ gradio.outputs.Label(num_top_classes=3),
55
+ live=True,
56
+ interpretation="default",
57
+ examples=[["images/cheetah1.jpg"], ["images/lion.jpg"]],
58
+ capture_session=True,
59
+ )
60
+
61
+ iface.test_launch()
62
+ iface.launch()