Spaces:
Runtime error
Runtime error
File size: 3,471 Bytes
ed5cf1f |
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 |
#!streamlit/bin/python
import streamlit as st
from pathlib import Path
import pandas as pd
import numpy as np
import tensorflow as tf
from PIL import Image
from io import BytesIO
import json
#from GDownload import download_file_from_google_drive
@st.cache(allow_output_mutation=True)
def load_model():
# if selected_model == 'PVAN-Stanford':
# model_location = '1-q1R5dLfIFW7BbzKuYTjolAoqpjVClsb'
# save_dest = Path('saved_model')
# save_dest.mkdir(exist_ok=True)
# saved_model = Path("saved_model/FerNet_EfficientNet.h5")
# elif selected_model == 'PVAN-Tsinghua':
# model_location = '1-q1R5dLfIFW7BbzKuYTjolAoqpjVClsb'
# save_dest = Path('saved_model')
# save_dest.mkdir(exist_ok=True)
# saved_model = Path("saved_model/FerNet_EfficientNet.h5")
# if not saved_model.exists():
# download_file_from_google_drive(model_location, saved_model)
saved_model = str(Path().parent.absolute())+"/saved_models/FerNetEfficientNetB2"
saved_model = tf.keras.models.load_model(saved_model)
return saved_model
@st.cache
def load_classes():
with open(str(Path().parent.absolute())+'/App/classes_dict.json') as classes:
class_names = json.load(classes)
return class_names
def load_and_prep_image(filename, img_shape=260):
#img = tf.io.read_file(filename)
img = np.array(filename)#tf.io.decode_image(filename, channels=3)
# Resize our image
img = tf.image.resize(img, [img_shape,img_shape])
# Scale
return img # don't need to resclae images for EfficientNet models in Tensorflow
if __name__ == '__main__':
hide_st_style = """
<style>
footer {visibility: hidden;}
header {visibility: hidden;}
</style>
"""
st.markdown(hide_st_style, unsafe_allow_html=True)
st.title("Dog Breeds Detector")
options = ['PVAN-Stanford', 'PVAN-Tsinghua']
selected_model = st.selectbox('Select a model to use (Default: PVAN-Stanford):', options)
saved_model = load_model()
class_names = load_classes()
st.write("Choose any dog image and get the corresponding breed:")
uploaded_image = st.file_uploader("Choose an image...")
if uploaded_image:
uploaded_image = Image.open(uploaded_image)
# try:
uploaded_image = uploaded_image.convert("RGB")
membuf = BytesIO()
uploaded_image.save(membuf, format="jpeg")
uploaded_image = Image.open(membuf)
# finally:
image_for_the_model = load_and_prep_image(uploaded_image)
prediction = saved_model.predict(tf.expand_dims(image_for_the_model, axis=0), verbose=0)
top_k_proba, top_k_indices = tf.nn.top_k(prediction,k=5)
top_5_classes = {top_n+1:class_names[str(top_k)] for top_n, top_k in enumerate(list(tf.squeeze(top_k_indices).numpy()))}
top_k_proba = tf.squeeze(top_k_proba).numpy()
top_5_classes = pd.DataFrame({"Top-k":top_5_classes.keys(), "Dog Breed": top_5_classes.values(), "Probability": top_k_proba})
#top_5_classes.set_index("Top-k", inplace=True)
print(tf.argmax(prediction, axis=1).numpy())
predicted_breed = class_names[str(tf.argmax(prediction, axis=1).numpy()[0])]
predicted_breed = ' '.join(predicted_breed.split('_'))
predicted_breed = predicted_breed.title()
st.header(f'This dog looks like a {predicted_breed}')
col1, col2 = st.columns([1,2])
col1.image(uploaded_image,use_column_width=True)
col2.bar_chart(top_5_classes, x="Dog Breed", y="Probability")
|