rdkulkarni commited on
Commit
bf7b758
1 Parent(s): 4ed9206

Create new file

Browse files
Files changed (1) hide show
  1. app1.py +90 -0
app1.py ADDED
@@ -0,0 +1,90 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ #Import libraries
3
+ import torch
4
+ import torchvision.models as models
5
+ import json
6
+ import skimage
7
+ #Import User Defined libraries
8
+ from neural_network_model import initialize_existing_models, build_custom_models, set_parameter_requires_grad
9
+ from utilities import process_image, get_input_args_predict
10
+
11
+ def predict(image_path, model, topk=5):
12
+
13
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
14
+ model.to(device)
15
+ model.eval()
16
+
17
+ tensor_img = torch.FloatTensor(process_image(image_path))
18
+ tensor_img = tensor_img.unsqueeze(0)
19
+ tensor_img = tensor_img.to(device)
20
+ log_ps = model(tensor_img)
21
+ result = log_ps.topk(topk)
22
+ if torch.cuda.is_available(): #gpu Move it from gpu to cpu for numpy
23
+ ps = torch.exp(result[0].data).cpu().numpy()[0]
24
+ idxs = result[1].data.cpu().numpy()[0]
25
+ else: #cpu Keep it on cpu for nump
26
+ ps = torch.exp(result[0].data).numpy()[0]
27
+ idxs = result[1].data.numpy()[0]
28
+
29
+ return (ps, idxs)
30
+
31
+ def process_input(image_path):
32
+ #0. Get user inputs
33
+ #in_arg = vars(get_input_args_predict())
34
+ #print("User arguments/hyperparameters or default used are as below")
35
+ #print(in_arg)
36
+ in_arg = {}
37
+ in_arg['gpu'] = 'gpu'
38
+ in_arg['save_dir'] = 'checkpoint-densenet121.pth'
39
+ in_arg['path'] = image_path
40
+ in_arg['top_k'] = 5
41
+
42
+ print(in_arg)
43
+ #1. Get device for prediction and Load model from checkpoint along with some other information
44
+ if in_arg['gpu'] == 'gpu' and torch.cuda.is_available():
45
+ device = torch.device("cuda")
46
+ checkpoint = torch.load(in_arg['save_dir'])
47
+ else:
48
+ device = "cpu"
49
+ checkpoint = torch.load(in_arg['save_dir'], map_location = device)
50
+ #print(f"Using {device} device for predicting/inference")
51
+
52
+ checkpoint['arch_type'] = 'existing'
53
+ if checkpoint['arch_type'] == 'existing':
54
+ model_ft, input_size = initialize_existing_models(checkpoint['arch'], checkpoint['arch_type'], len(checkpoint['class_to_idx']),
55
+ checkpoint['feature_extract'], checkpoint['hidden_units'], use_pretrained=False)
56
+ elif checkpoint['arch_type'] == 'custom':
57
+ model_ft = build_custom_models(checkpoint['arch'], checkpoint['arch_type'], len(checkpoint['class_to_idx']), checkpoint['feature_extract'],
58
+ checkpoint['hidden_units'], use_pretrained=True)
59
+ else:
60
+ #print("Nothing to predict")
61
+ exit()
62
+
63
+ model_ft.class_to_idx = checkpoint['class_to_idx']
64
+ model_ft.gpu_or_cpu = checkpoint['gpu_or_cpu']
65
+ model_ft.load_state_dict(checkpoint['state_dict'])
66
+ model_ft.to(device)
67
+
68
+ #Predict
69
+ # Get the prediction by passing image and other user preferences through the model
70
+ probs, idxs = predict(image_path = in_arg['path'], model = model_ft, topk = in_arg['top_k'])
71
+
72
+ # Swap class to index mapping with index to class mapping and then map the classes to flower category labels using the json file
73
+ idx_to_class = {v: k for k, v in model_ft.class_to_idx.items()}
74
+ with open('cat_to_name.json','r') as f:
75
+ cat_to_name = json.load(f)
76
+ names = list(map(lambda x: cat_to_name[f"{idx_to_class[x]}"],idxs))
77
+
78
+ #return names, probs
79
+ return {names[i]: float(probs[i]) for i in range(len(names))}
80
+
81
+ examples = ['16_image_06670.jpg','33_image_06460.jpg','80_image_02020.jpg', 'Flowers.png','inference_example.png']
82
+ title = "Image Classifier - Which species of Flower?"
83
+ description = "An image classifier to recognize different species of flowers trained on 102 Category Flower Dataset"
84
+ article = article="<p style='text-align: center'><a href='https://www.robots.ox.ac.uk/~vgg/data/flowers/102/index.html' target='_blank'>Source 102 Flower Dataset</a></p>"
85
+ interpretation = 'default'
86
+ enable_queue = True
87
+ iface = gr.Interface(fn=process_input, inputs=gr.inputs.Image(type='filepath'), outputs=gr.outputs.Label(num_top_classes=3), examples = examples,
88
+ title=title, description=description,article=article,interpretation=interpretation, enable_queue=enable_queue
89
+ )
90
+ iface.launch()