Vrk commited on
Commit
34fc4ed
·
1 Parent(s): 70bf604

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +92 -0
app.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import tensorflow as tf
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.optim as optim
7
+ import torch.nn.functional as F
8
+ import timm
9
+ from PIL import Image
10
+ from torchvision import transforms
11
+
12
+ from Models import ResNet, EfficientNet, BaseLine
13
+
14
+ def get_model(model_name, classes, device):
15
+
16
+ if model_name == 'Inception-V3':
17
+ model = tf.lite.Interpreter(model_path='vgg.tflite')
18
+ model.allocate_tensors()
19
+
20
+ elif model_name == 'VGG':
21
+ model = tf.lite.Interpreter(model_path='vgg.tflite')
22
+ model.allocate_tensors()
23
+
24
+ elif model_name == 'EfficientNet-B0':
25
+ model = EfficientNet(len(classes)).to(device)
26
+ model.load_state_dict(torch.load('EfficientNet-Model.pt'))
27
+
28
+ elif model_name == 'ResNet-50':
29
+ model = ResNet(len(classes)).to(device)
30
+ model.load_state_dict(torch.load('model-resnet50.pt'))
31
+
32
+ elif model_name == 'Base Line Model':
33
+ model = BaseLine(len(classes)).to(device)
34
+ model.load_state_dict(torch.load('BaseLine-Model.pt'))
35
+
36
+ return model
37
+
38
+ def make_predictions(input_img, model_name):
39
+ classes = ['buildings','forest', 'glacier', 'mountain', 'sea', 'street']
40
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
41
+
42
+ model = get_model(model_name, classes, device)
43
+
44
+ if model_name in ['EfficientNet-B0', 'ResNet-50', 'Base Line Model']:
45
+ model.eval()
46
+ img = get_transform(input_img, device)
47
+ pred = model(img)
48
+ if torch.cuda.is_available():
49
+ pred = F.softmax(pred).detach().cpu().numpy()
50
+ y_prob = pred.argmax(axis=1)[0]
51
+ else:
52
+ pred = F.softmax(pred).detach().numpy()
53
+ y_prob = pred.argmax(axis=1)[0]
54
+
55
+ if model_name in ['Inception-V3', 'VGG']:
56
+ input_img = np.array(input_img)
57
+ img = input_img / 255.
58
+ input_tensor= np.array(np.expand_dims(img,0), dtype=np.float32)
59
+ input_index = model.get_input_details()[0]["index"]
60
+
61
+ # setting input tensor
62
+ model.set_tensor(input_index, input_tensor)
63
+
64
+ #Run the inference
65
+ model.invoke()
66
+ output_details = model.get_output_details()
67
+
68
+ # output data of image
69
+ pred = model.get_tensor(output_details[0]['index'])
70
+ y_prob = pred.argmax()
71
+
72
+ label = classes[y_prob]
73
+ confidences = {classes[i]: float(pred[0][i]) for i in range(len(classes))}
74
+
75
+ return label, confidences
76
+
77
+ demo = gr.Interface(
78
+ fn = make_predictions,
79
+ inputs = [gr.Image(shape=(150, 150), type="pil"), gr.Dropdown(choices=['EfficientNet-B0', 'ResNet-50', 'Inception-V3', 'VGG', 'Base Line Model'], value='EfficientNet-B0', label='Choose Model')],
80
+ outputs = [gr.outputs.Textbox(label="Output Class"), gr.outputs.Label(label='Confidences')],
81
+ title = "MultiClass Classifier",
82
+ examples=[
83
+ ["Sample_Images/Buildings.jpg", 'EfficientNet-B0'],
84
+ ["Sample_Images/Forest.jpg", 'EfficientNet-B0'],
85
+ ['Sample_Images/Street.jpg', 'EfficientNet-B0'],
86
+ ['Sample_Images/glacier.jpg', 'EfficientNet-B0'],
87
+ ['Sample_Images/mountain.jpg', 'EfficientNet-B0'],
88
+ ['Sample_Images/sea.jpg', 'EfficientNet-B0']
89
+ ],
90
+ )
91
+
92
+ demo.launch(debug=True)