Spaces:
Sleeping
Sleeping
File size: 2,506 Bytes
71a443b 22bb2bd 71a443b 61f8391 71a443b 61f8391 71a443b 84f78ed 71a443b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import gradio as gr
import os
import torch
from model import create_vit_model
from timeit import default_timer as timer
from typing import Tuple, Dict
# setup class names
class_names = ['daisy', 'dandelion', 'rose', 'sunflower', 'tulip']
### model and transforms preparation
#create vit model
vit, vit_transforms = create_vit_model(num_classes=5)
#load saved weights
vit.load_state_dict(torch.load(f="pretrained_vit_festure_extractor_flower_classification.pth",
map_location=torch.device('cpu')))
### predict function
# create predict function
def predict(img) -> Tuple[Dict, float]:
"""transforms and perfroms a prediction on img and returns prediction and time taken"""
#start the time
start_time = timer()
#transform the target image and add a batch dim
img = vit_transforms(img).unsqueeze(0)
#put the model in eval mode and turn on inference
vit.eval()
with torch.inference_mode():
# pass the transformed imag thru the model and turn the prediction logits into prediction probabilities
pred_probs = torch.softmax(vit(img), dim =1)
# create a prediction label and prediction probability
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
# calculate prediction time
end_time = timer()
pred_time = round(end_time - start_time, 5)
# return the prediction dictionary and prediction time
return pred_labels_and_probs, pred_time
### gradio app
# create title, description and article strings
title = "flower classification"
description = "a vit_16 feature extractor computer vision model to classify images of flowers as: daisy, dandelion, rose, sunflower, and tulip"
article = "created at [flower classification](https://github.com/geitta/flower_classification)"
# create examples list from examples dir
example_list = [["examples/" + example] for example in os.listdir("examples")]
# create gradio demo
demo = gr.Interface(fn=predict, #mapping function from input to output
inputs=gr.Image(type="pil"), # what are the inputs
outputs=[gr.Label(num_top_classes=5, label="predictions"), #what are the outputs
gr.Number(label="prediction time (s)")], #our fn has two outputs, therefore we have 2 outputs
examples = example_list,
title=title,
description=description,
article=article
)
#launch the demo
demo.launch()
|