Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import numpy as np | |
| import tensorflow as tf | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| import torch.nn.functional as F | |
| import timm | |
| from PIL import Image | |
| from torchvision import transforms | |
| from Models import ResNet, EfficientNet, BaseLine | |
| def get_model(model_name, classes, device): | |
| if model_name == 'Inception-V3': | |
| model = tf.lite.Interpreter(model_path='vgg.tflite') | |
| model.allocate_tensors() | |
| elif model_name == 'VGG': | |
| model = tf.lite.Interpreter(model_path='vgg.tflite') | |
| model.allocate_tensors() | |
| elif model_name == 'EfficientNet-B0': | |
| model = EfficientNet(len(classes)).to(device) | |
| model.load_state_dict(torch.load('EfficientNet-Model.pt', map_location=torch.device(device))) | |
| elif model_name == 'ResNet-50': | |
| model = ResNet(len(classes)).to(device) | |
| model.load_state_dict(torch.load('model-resnet50.pt', map_location=torch.device(device))) | |
| elif model_name == 'Base Line Model': | |
| model = BaseLine(len(classes)).to(device) | |
| model.load_state_dict(torch.load('BaseLine-Model.pt', map_location=torch.device(device))) | |
| return model | |
| def get_transform(input_img, device): | |
| normalize = transforms.Normalize( | |
| [0.485, 0.456, 0.406], | |
| [0.229, 0.224, 0.225] | |
| ) | |
| test_transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| normalize, | |
| ]) | |
| input_img = test_transform(input_img).unsqueeze(0).to(device) | |
| return input_img | |
| def make_predictions(input_img, model_name): | |
| classes = ['buildings','forest', 'glacier', 'mountain', 'sea', 'street'] | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| model = get_model(model_name, classes, device) | |
| if model_name in ['EfficientNet-B0', 'ResNet-50', 'Base Line Model']: | |
| model.eval() | |
| img = get_transform(input_img, device) | |
| pred = model(img) | |
| if torch.cuda.is_available(): | |
| pred = F.softmax(pred).detach().cpu().numpy() | |
| y_prob = pred.argmax(axis=1)[0] | |
| else: | |
| pred = F.softmax(pred).detach().numpy() | |
| y_prob = pred.argmax(axis=1)[0] | |
| if model_name in ['Inception-V3', 'VGG']: | |
| input_img = np.array(input_img) | |
| img = input_img / 255. | |
| input_tensor= np.array(np.expand_dims(img,0), dtype=np.float32) | |
| input_index = model.get_input_details()[0]["index"] | |
| # setting input tensor | |
| model.set_tensor(input_index, input_tensor) | |
| #Run the inference | |
| model.invoke() | |
| output_details = model.get_output_details() | |
| # output data of image | |
| pred = model.get_tensor(output_details[0]['index']) | |
| y_prob = pred.argmax() | |
| label = classes[y_prob] | |
| confidences = {classes[i]: float(pred[0][i]) for i in range(len(classes))} | |
| return label, confidences | |
| demo = gr.Interface( | |
| fn = make_predictions, | |
| 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')], | |
| outputs = [gr.outputs.Textbox(label="Output Class"), gr.outputs.Label(label='Confidences')], | |
| title = "MultiClass Classifier", | |
| examples=[ | |
| ["Buildings.jpg", 'EfficientNet-B0'], | |
| ["Forest.jpg", 'EfficientNet-B0'], | |
| ['Street.jpg', 'EfficientNet-B0'], | |
| ['glacier.jpg', 'EfficientNet-B0'], | |
| ['mountain.jpg', 'EfficientNet-B0'], | |
| ['sea.jpg', 'EfficientNet-B0'] | |
| ], | |
| ) | |
| demo.launch(debug=True, inline=True) |