Spaces:
Paused
Paused
Commit
•
0188de7
1
Parent(s):
6d36038
Create new file
Browse files
app2.py
ADDED
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
from torch import nn
|
4 |
+
import torchvision
|
5 |
+
import torchvision.models as models
|
6 |
+
print(torchvision.__version__)
|
7 |
+
from PIL import Image
|
8 |
+
import json
|
9 |
+
import skimage
|
10 |
+
import numpy as np
|
11 |
+
import math
|
12 |
+
from typing import List
|
13 |
+
print(models.AlexNet_Weights.DEFAULT)
|
14 |
+
|
15 |
+
def pred_and_plot_image(
|
16 |
+
model: torch.nn.Module,
|
17 |
+
image_path: str,
|
18 |
+
class_names: List[str] = None,
|
19 |
+
transform=None,
|
20 |
+
device: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
|
21 |
+
):
|
22 |
+
|
23 |
+
# 1. Load in image and convert the tensor values to float32
|
24 |
+
target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)
|
25 |
+
|
26 |
+
# 2. Divide the image pixel values by 255 to get them between [0, 1]
|
27 |
+
target_image = target_image / 255.0
|
28 |
+
|
29 |
+
# 3. Transform if necessary
|
30 |
+
if transform:
|
31 |
+
target_image = transform(target_image)
|
32 |
+
|
33 |
+
# 4. Make sure the model is on the target device
|
34 |
+
model.to(device)
|
35 |
+
|
36 |
+
# 5. Turn on model evaluation mode and inference mode
|
37 |
+
model.eval()
|
38 |
+
with torch.inference_mode():
|
39 |
+
# Add an extra dimension to the image
|
40 |
+
target_image = target_image.unsqueeze(dim=0)
|
41 |
+
|
42 |
+
# Make a prediction on image with an extra dimension and send it to the target device
|
43 |
+
target_image_pred = model(target_image.to(device))
|
44 |
+
|
45 |
+
# 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
|
46 |
+
target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
|
47 |
+
|
48 |
+
# 7. Convert prediction probabilities -> prediction labels
|
49 |
+
target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
|
50 |
+
|
51 |
+
if class_names:
|
52 |
+
idxs= class_names[target_image_pred_label.cpu()]
|
53 |
+
else:
|
54 |
+
idxs = target_image_pred_label
|
55 |
+
|
56 |
+
ps = target_image_pred_probs.max().cpu()
|
57 |
+
print(ps)
|
58 |
+
print(idxs)
|
59 |
+
|
60 |
+
return (ps, idxs)
|
61 |
+
|
62 |
+
|
63 |
+
def set_parameter_requires_grad(model, feature_extracting):
|
64 |
+
if feature_extracting:
|
65 |
+
for param in model.parameters():
|
66 |
+
param.requires_grad = False
|
67 |
+
|
68 |
+
def update_last_layer_pretrained_model(pretrained_model, num_classes, feature_extract):
|
69 |
+
set_parameter_requires_grad(pretrained_model, feature_extract)
|
70 |
+
if hasattr(pretrained_model, 'fc') and 'resnet' in pretrained_model.__class__.__name__.lower(): #resnet
|
71 |
+
num_ftrs = pretrained_model.fc.in_features
|
72 |
+
pretrained_model.fc = nn.Linear(num_ftrs, num_classes, bias = True)
|
73 |
+
elif hasattr(pretrained_model, 'classifier') and ('alexnet' in pretrained_model.__class__.__name__.lower() or 'vgg' in pretrained_model.__class__.__name__.lower()): #alexNet, vgg
|
74 |
+
num_ftrs = pretrained_model.classifier[6].in_features
|
75 |
+
pretrained_model.classifier[6] = nn.Linear(num_ftrs, num_classes, bias = True)
|
76 |
+
elif hasattr(pretrained_model, 'classifier') and 'squeezenet' in pretrained_model.__class__.__name__.lower(): #squeezenet
|
77 |
+
pretrained_model.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
|
78 |
+
pretrained_model.num_classes = num_classes
|
79 |
+
elif hasattr(pretrained_model, 'classifier') and ('efficientnet' in pretrained_model.__class__.__name__.lower() or 'mobilenet' in pretrained_model.__class__.__name__.lower()): #efficientnet, mobilenet
|
80 |
+
num_ftrs = pretrained_model.classifier[1].in_features
|
81 |
+
pretrained_model.classifier[1] = nn.Linear(num_ftrs, num_classes, bias = True)
|
82 |
+
elif hasattr(pretrained_model, 'AuxLogits') and 'inception' in pretrained_model.__class__.__name__.lower(): #inception
|
83 |
+
num_ftrs = pretrained_model.AuxLogits.fc.in_features
|
84 |
+
pretrained_model.AuxLogits.fc = nn.Linear(num_ftrs, num_classes) #Auxilary net
|
85 |
+
num_ftrs = pretrained_model.fc.in_features
|
86 |
+
pretrained_model.fc = nn.Linear(num_ftrs,num_classes) #Primary net
|
87 |
+
elif hasattr(pretrained_model, 'classifier') and 'densenet' in pretrained_model.__class__.__name__.lower(): #densenet
|
88 |
+
num_ftrs = pretrained_model.classifier.in_features
|
89 |
+
pretrained_model.classifier = nn.Linear(num_ftrs, num_classes, bias = True)
|
90 |
+
elif hasattr(pretrained_model, 'heads') and 'visiontransformer' in pretrained_model.__class__.__name__.lower(): #vit transformer
|
91 |
+
num_ftrs = pretrained_model.heads.head.in_features
|
92 |
+
pretrained_model.heads.head = nn.Linear(num_ftrs, num_classes, bias = True)
|
93 |
+
elif hasattr(pretrained_model, 'head') and 'swin' in pretrained_model.__class__.__name__.lower(): #swin transformer
|
94 |
+
num_ftrs = pretrained_model.head.in_features
|
95 |
+
pretrained_model.head = nn.Linear(num_ftrs, num_classes, bias = True)
|
96 |
+
return pretrained_model
|
97 |
+
|
98 |
+
def process_image(image_path):
|
99 |
+
|
100 |
+
im = Image.open(image_path)
|
101 |
+
# Resize
|
102 |
+
if im.size[1] < im.size[0]:
|
103 |
+
im.thumbnail((255, math.pow(255, 2)))
|
104 |
+
else:
|
105 |
+
im.thumbnail((math.pow(255, 2), 255))
|
106 |
+
|
107 |
+
#Crop
|
108 |
+
width, height = im.size
|
109 |
+
left = (width - 224)/2
|
110 |
+
top = (height - 224)/2
|
111 |
+
right = (width + 224)/2
|
112 |
+
bottom = (height + 224)/2
|
113 |
+
im = im.crop((left, top, right, bottom))
|
114 |
+
|
115 |
+
#Convert to np.array
|
116 |
+
np_image = np.array(im)/255
|
117 |
+
|
118 |
+
#Undo Mean, Standard Deviation and Transpose
|
119 |
+
mean = np.array([0.485, 0.456, 0.406])
|
120 |
+
std = np.array([0.229, 0.224, 0.225])
|
121 |
+
np_image = (np_image - mean)/ std
|
122 |
+
np_image = np.transpose(np_image, (2, 0, 1))
|
123 |
+
|
124 |
+
return np_image
|
125 |
+
|
126 |
+
def predict(image_path, model, topk=5):
|
127 |
+
|
128 |
+
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
129 |
+
model.to(device)
|
130 |
+
model.eval()
|
131 |
+
|
132 |
+
tensor_img = torch.FloatTensor(process_image(image_path))
|
133 |
+
tensor_img = tensor_img.unsqueeze(0)
|
134 |
+
tensor_img = tensor_img.to(device)
|
135 |
+
log_ps = model(tensor_img)
|
136 |
+
result = log_ps.topk(topk)
|
137 |
+
|
138 |
+
#target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
|
139 |
+
#target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
|
140 |
+
print(result)
|
141 |
+
if torch.cuda.is_available(): #gpu Move it from gpu to cpu for numpy
|
142 |
+
ps = torch.exp(result[0].data).cpu().numpy()[0]
|
143 |
+
idxs = result[1].data.cpu().numpy()[0]
|
144 |
+
else: #cpu Keep it on cpu for nump
|
145 |
+
ps = torch.exp(result[0].data).numpy()[0]
|
146 |
+
idxs = result[1].data.numpy()[0]
|
147 |
+
|
148 |
+
return (ps, idxs)
|
149 |
+
|
150 |
+
def process_input(image_path):
|
151 |
+
|
152 |
+
#list_of_models_weights_paths = [('densenet121','DenseNet121_Weights','checkpoint-densenet121.pth')] #[('swin_b','Swin_B'), ('vit_b_16', 'ViT_B_16')]
|
153 |
+
#list_of_models_weights_paths = [('alexnet','AlexNet_Weights','flowers_alexnet_model.pth'), ('resnet18','ResNet18_Weights','flowers_resnet18_model.pth')] #[('swin_b','Swin_B'), ('vit_b_16', 'ViT_B_16')]
|
154 |
+
#list_of_predictions = []
|
155 |
+
#for model in list_of_models_weights_paths:
|
156 |
+
#Load saved model
|
157 |
+
model_name, model_weights, model_path = ('alexnet','AlexNet_Weights','flowers_alexnet_model.pth')
|
158 |
+
checkpoint = torch.load(model_path, map_location='cpu')
|
159 |
+
pretrained_weights = eval(f"torchvision.models.{model_weights}.DEFAULT")
|
160 |
+
auto_transforms = pretrained_weights.transforms()
|
161 |
+
#pretrained_model = eval(f"torchvision.models.{model_name}(weights = pretrained_weights)")
|
162 |
+
pretrained_model = eval(f"models.{model_name}(pretrained = True)")
|
163 |
+
pretrained_model = update_last_layer_pretrained_model(pretrained_model, 102, True)
|
164 |
+
#pretrained_model.class_to_idx = checkpoint['class_to_idx']
|
165 |
+
pretrained_model.class_names = checkpoint['class_names']
|
166 |
+
pretrained_model.load_state_dict(checkpoint['state_dict'])
|
167 |
+
pretrained_model.to('cpu')
|
168 |
+
#probs, idxs = predict(image_path, model = pretrained_model, topk = 5)
|
169 |
+
#
|
170 |
+
probs, idxs = pred_and_plot_image(model=pretrained_model,
|
171 |
+
image_path=image_path,
|
172 |
+
class_names=pretrained_model.class_names,
|
173 |
+
transform=auto_transforms)
|
174 |
+
print(probs)
|
175 |
+
print(idxs)
|
176 |
+
|
177 |
+
return { idxs : float(probs) }
|
178 |
+
|
179 |
+
#append to list to be returned at end
|
180 |
+
#response = {names[i]: float(probs[i]) for i in range(len(names))}
|
181 |
+
#list_of_predictions.append(response)
|
182 |
+
|
183 |
+
#return list_of_predictions
|
184 |
+
|
185 |
+
examples = ['16_image_06670.jpg','33_image_06460.jpg','80_image_02020.jpg', 'Flowers.png','inference_example.png']
|
186 |
+
title = "Image Classifier - Species of Flower predicted by different Models"
|
187 |
+
description = "Image classifiers to recognize different species of flowers trained on 102 Category Flower Dataset"
|
188 |
+
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>"
|
189 |
+
interpretation = 'default'
|
190 |
+
enable_queue = True
|
191 |
+
iface = gr.Interface(fn=process_input, inputs=gr.inputs.Image(type='filepath'), outputs=gr.outputs.Label(num_top_classes=3), examples = examples,
|
192 |
+
title=title, description=description,article=article,interpretation=interpretation, enable_queue=enable_queue
|
193 |
+
)
|
194 |
+
iface.launch()
|
195 |
+
|
196 |
+
# Swap class to index mapping with index to class mapping and then map the classes to flower category labels using the json file
|
197 |
+
#idx_to_class = {v: k for k, v in pretrained_model.class_to_idx.items()}
|
198 |
+
#with open('cat_to_name.json','r') as f:
|
199 |
+
# cat_to_name = json.load(f)
|
200 |
+
#names = list(map(lambda x: cat_to_name[f"{idx_to_class[x]}"],idxs))
|
201 |
+
#return names, probs
|
202 |
+
#return { idxs[i].item() : float(probs[i].item()) for i in range(len(idxs))}
|