|
import json |
|
import requests |
|
import numpy as np |
|
import streamlit as st |
|
import tensorflow as tf |
|
from huggingface_hub import snapshot_download |
|
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) |
|
|
|
predictions = model.predict(img_array) |
|
|
|
top4 = predictions.argsort()[0, -1:-5:-1] |
|
|
|
breakdown = [] |
|
for race, acc in zip(np.array(race_names)[top4], predictions[0, top4]): |
|
breakdown.append(f'{race} at {acc:.2%}') |
|
return breakdown |
|
|
|
|
|
def main(): |
|
|
|
race_file = open('race_names.json', 'r') |
|
race_names = json.load(race_file) |
|
|
|
download_folder = snapshot_download(repo_id='JbIPS/DogRace') |
|
model = load_model(download_folder) |
|
|
|
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() |
|
|