DogBreed / app.py
JbIPS
Load model directly from hub
22535ad
raw history blame
No virus
2.08 kB
import json
import requests
import numpy as np
import streamlit as st
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
from PIL import Image
from io import BytesIO
def make_prediction(url, model, race_names):
response = requests.get(url)
img = Image.open(BytesIO(response.content)).resize((180, 180))
img_array = img_to_array(img)
img_array = tf.expand_dims(img_array, 0) # Create batch axis
predictions = model.predict(img_array)
top5 = predictions.argsort()[0, -1:-6:-1]
breakdown = []
for race, acc in zip(np.array(race_names)[top5], predictions[0, top5]):
breakdown.append(f'{race} at {acc:.2%}')
return breakdown
def main():
race_file = open('race_names.json', 'r')
race_names = json.load(race_file)
# Load model
model = from_pretrained_keras('JbIPS/DogRace')
st.set_page_config("Who let's the dogs out")
st.title('Quelle est ta race de chien totem ?')
st.text('''
Pour découvrir ta race de chien, colle l'adresse d'une photo.
''')
url = st.text_input('URL de la photo')
predict_btn = st.button('Prédire')
if predict_btn:
pred = make_prediction(url, model, race_names)
main_race = pred[0].split(' at')[0].lower()
main_race = '/'.join(reversed(main_race.replace('-', '').split(' ')))
if main_race.startswith('husky'):
main_race = main_race.split('/')[0]
col1, col2 = st.columns(2)
with col1:
st.image(url)
with col2:
response = requests.get(url=f'https://dog.ceo/api/breed/{main_race}/images/random').json()
if response['status'] == 'success':
st.image(response['message'])
else:
st.text(main_race)
st.text(response)
st.write('Les races qui te correspondent le plus sont :')
for race in pred:
st.write(race)
if __name__ == '__main__':
main()