Spaces:
Sleeping
Sleeping
File size: 4,216 Bytes
184c4c1 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
### 1. Imports and class names setup ###
import gradio as gr
import os
import torch
from model import create_effnetb2_model
from timeit import default_timer as timer
from typing import Tuple, Dict
# Setup class names
class_names= ['alpine sea holly',
'anthurium',
'artichoke',
'azalea',
'ball moss',
'balloon flower',
'barbeton daisy',
'bearded iris',
'bee balm',
'bird of paradise',
'bishop of llandaff',
'black-eyed susan',
'blackberry lily',
'blanket flower',
'bolero deep blue',
'bougainvillea',
'bromelia',
'buttercup',
'californian poppy',
'camellia',
'canna lily',
'canterbury bells',
'cape flower',
'carnation',
'cautleya spicata',
'clematis',
"colt's foot",
'columbine',
'common dandelion',
'corn poppy',
'cyclamen',
'daffodil',
'desert-rose',
'english marigold',
'fire lily',
'foxglove',
'frangipani',
'fritillary',
'garden phlox',
'gaura',
'gazania',
'geranium',
'giant white arum lily',
'globe thistle',
'globe-flower',
'grape hyacinth',
'great masterwort',
'hard-leaved pocket orchid',
'hibiscus',
'hippeastrum',
'japanese anemone',
'king protea',
'lenten rose',
'lotus lotus',
'love in the mist',
'magnolia',
'mallow',
'marigold',
'mexican aster',
'mexican petunia',
'monkshood',
'moon orchid',
'morning glory',
'orange dahlia',
'osteospermum',
'oxeye daisy',
'passion flower',
'pelargonium',
'peruvian lily',
'petunia',
'pincushion flower',
'pink primrose',
'pink-yellow dahlia',
'poinsettia',
'primula',
'prince of wales feathers',
'purple coneflower',
'red ginger',
'rose',
'ruby-lipped cattleya',
'siam tulip',
'silverbush',
'snapdragon',
'spear thistle',
'spring crocus',
'stemless gentian',
'sunflower',
'sweet pea',
'sweet william',
'sword lily',
'thorn apple',
'tiger lily',
'toad lily',
'tree mallow',
'tree poppy',
'trumpet creeper',
'wallflower',
'water lily',
'watercress',
'wild pansy',
'windflower',
'yellow iris'
]
### 2. Model and transforms preparation ###
# Create EffNetB2 model
effnetb2, effnetb2_transforms = create_effnetb2_model(
num_classes=102, # len(class_names) would also work
)
# Load saved weights
effnetb2.load_state_dict(
torch.load(
f="pretrained_effnetb2_feature_extractor_fl102.pth",
map_location=torch.device("cpu"), # load to CPU
)
)
### 3. Predict function ###
# Create predict function
def predict(img) -> Tuple[Dict, float]:
"""Transforms and performs a prediction on img and returns prediction and time taken.
"""
# Start the timer
start_time = timer()
# Transform the target image and add a batch dimension
img = effnetb2_transforms(img).unsqueeze(0)
# Put model into evaluation mode and turn on inference mode
effnetb2.eval()
with torch.inference_mode():
# Pass the transformed image through the model and turn the prediction logits into prediction probabilities
pred_probs = torch.softmax(effnetb2(img), dim=1)
# Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
# Calculate the prediction time
pred_time = round(timer() - start_time, 5)
# Return the prediction dictionary and prediction time
return pred_labels_and_probs, pred_time
### 4. Gradio app ###
title = "Flofi Demo"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of 102 flower species."
article = "Created by Haydar Uçar."
# Create the 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=3, label="Predictions"), # what are the outputs?
gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
title=title,
description=description,
article=article)
# Launch the demo!
demo.launch() # generate a publically shareable URL?
|