Spaces:
Build error
Build error
1st version of the app
Browse files
app.py
CHANGED
|
@@ -1,4 +1,188 @@
|
|
|
|
|
| 1 |
import streamlit as st
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
|
| 3 |
-
|
| 4 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#importing the libraries
|
| 2 |
import streamlit as st
|
| 3 |
+
from PIL import Image
|
| 4 |
+
import torch
|
| 5 |
+
from transformers import AutoModelForImageClassification, AutoImageProcessor
|
| 6 |
+
import numpy as np
|
| 7 |
+
import pandas as pd
|
| 8 |
+
import time
|
| 9 |
+
import os
|
| 10 |
|
| 11 |
+
model_repository_id = "Dusduo/Pokemon-classification-1stGen"
|
| 12 |
+
# Loading the pokemon classifier model and its processor
|
| 13 |
+
image_processor = AutoImageProcessor.from_pretrained(model_repository_id)
|
| 14 |
+
model = AutoModelForImageClassification.from_pretrained(model_repository_id)
|
| 15 |
+
# Loading the pokemon information table
|
| 16 |
+
pokemon_info_df = pd.read_csv('pokemon_info.csv')
|
| 17 |
+
|
| 18 |
+
pokeball_image = Image.open('pokeball.png').resize((20,20))
|
| 19 |
+
|
| 20 |
+
#functions to predict image
|
| 21 |
+
def preprocess(processor: AutoImageProcessor, image):
|
| 22 |
+
return processor(image.convert("RGB").resize((200,200)), return_tensors="pt")
|
| 23 |
+
|
| 24 |
+
def predict(model: AutoModelForImageClassification, inputs, k=5):
|
| 25 |
+
|
| 26 |
+
# Forward the image to the model and retrieve the logits
|
| 27 |
+
with torch.no_grad():
|
| 28 |
+
logits = model(**inputs).logits
|
| 29 |
+
|
| 30 |
+
# Convert the retrieved logits into a vector of probabilities for each class
|
| 31 |
+
probabilities = torch.softmax(logits[0], dim=0).tolist()
|
| 32 |
+
|
| 33 |
+
# Discriminate wether or not the inputted image was an image of a Pokemon
|
| 34 |
+
# Compute the variance of the vector of probabilities
|
| 35 |
+
# The spread of the probability values is a good represent of the confusion of the model
|
| 36 |
+
# Or in other words, its confidence => the greater the spread, the lower its confidence
|
| 37 |
+
variance = np.var(probabilities)
|
| 38 |
+
|
| 39 |
+
# Too great of a spread: it is likely the image provided did not correspond to any known classes
|
| 40 |
+
if variance < 0.001: #not a pokemon
|
| 41 |
+
predicted_label = 'not a pokemon'
|
| 42 |
+
probability = -1
|
| 43 |
+
(top_k_labels, top_k_probability) = '_', '_'
|
| 44 |
+
else: # it is a pokemon
|
| 45 |
+
# Retrieve the predicted class (pokemon)
|
| 46 |
+
predicted_id = logits.argmax(-1).item()
|
| 47 |
+
predicted_label = model.config.id2label[predicted_id]
|
| 48 |
+
# Retrieve the probability for the predicted class, and format it to 2 decimals
|
| 49 |
+
probability = round(probabilities[predicted_id]*100,2)
|
| 50 |
+
# Retrieve the top 5 classes and their probabilities
|
| 51 |
+
#top_k_labels = [model.config.id2label[key] for key in np.argpartition(logits.numpy(), -k)[-k:]]
|
| 52 |
+
#top_k_probability = [round(prob*100,2) for prob in np.sort(probabilities.numpy())[-k:]]
|
| 53 |
+
|
| 54 |
+
return predicted_label, probability #, (top_k_labels, top_k_probability)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
# Designing the interface ------------------------------------------
|
| 59 |
+
|
| 60 |
+
# Use the full page instead of a narrow central column
|
| 61 |
+
st.set_page_config(layout="wide")
|
| 62 |
+
|
| 63 |
+
# Define the title
|
| 64 |
+
st.title("Gotta Classify 'Em All - 1st Generation Pokedex -")
|
| 65 |
+
# For newline
|
| 66 |
+
st.write('\n')
|
| 67 |
+
|
| 68 |
+
image = Image.open('anime1.jpeg')
|
| 69 |
+
|
| 70 |
+
col1, col2 = st.columns([3,1]) # [3,1]
|
| 71 |
+
|
| 72 |
+
with col1:
|
| 73 |
+
image = Image.open('anime1.jpeg')
|
| 74 |
+
show = st.image(image, use_column_width=True)
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
# Display Sample images ----
|
| 79 |
+
st.subheader('Sample images')
|
| 80 |
+
|
| 81 |
+
sample_imgs_dir = "sample_imgs/"
|
| 82 |
+
sample_imgs = os.listdir(sample_imgs_dir) # get the list of all sample images
|
| 83 |
+
img_idx = 0
|
| 84 |
+
|
| 85 |
+
n_cols = 4
|
| 86 |
+
groups = []
|
| 87 |
+
for i in range(0, len(sample_imgs), n_cols):
|
| 88 |
+
groups.append(sample_imgs[i:i+n_cols])
|
| 89 |
+
|
| 90 |
+
for group in groups:
|
| 91 |
+
cols = st.columns(n_cols)
|
| 92 |
+
for i,image_file in enumerate(group):
|
| 93 |
+
cols[i].image(sample_imgs_dir+image_file)
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
# Sidebar work and model outputs ---------------
|
| 98 |
+
|
| 99 |
+
st.sidebar.title("Upload Image")
|
| 100 |
+
|
| 101 |
+
#Disabling warning
|
| 102 |
+
#st.set_option('deprecation.showfileUploaderEncoding', False)
|
| 103 |
+
#Choose your own image
|
| 104 |
+
uploaded_file = st.sidebar.file_uploader("",type=['png', 'jpg', 'jpeg'], accept_multiple_files=False )
|
| 105 |
+
|
| 106 |
+
if uploaded_file is not None:
|
| 107 |
+
|
| 108 |
+
u_img = Image.open(uploaded_file)
|
| 109 |
+
show.image(u_img, 'Uploaded Image', width=400 )#use_column_width=True)
|
| 110 |
+
|
| 111 |
+
# Preprocess the image for the model
|
| 112 |
+
model_inputs = preprocess(image_processor, u_img)
|
| 113 |
+
|
| 114 |
+
# For newline
|
| 115 |
+
st.sidebar.write('\n')
|
| 116 |
+
|
| 117 |
+
if st.sidebar.button("Click Here to Classify"):
|
| 118 |
+
|
| 119 |
+
if uploaded_file is None:
|
| 120 |
+
|
| 121 |
+
st.sidebar.write("Please upload an Image to Classify")
|
| 122 |
+
|
| 123 |
+
else:
|
| 124 |
+
|
| 125 |
+
with st.spinner('Classifying ...'):
|
| 126 |
+
# Get prediction
|
| 127 |
+
prediction, probability = predict(model, model_inputs,5) #, (top_k_labels, top_k_probability)
|
| 128 |
+
time.sleep(2)
|
| 129 |
+
st.sidebar.success('Done!')
|
| 130 |
+
|
| 131 |
+
st.sidebar.header("Model predicts: ")
|
| 132 |
+
|
| 133 |
+
# Display prediction
|
| 134 |
+
|
| 135 |
+
if probability==-1:
|
| 136 |
+
|
| 137 |
+
st.sidebar.write("It seems like it is not a picture of a 1st Generation Pokemon alone.", '\n',
|
| 138 |
+
"There might be too many entities on the image." )
|
| 139 |
+
|
| 140 |
+
else:
|
| 141 |
+
st.sidebar.write(f" It's a(n) {prediction} picture.",'\n')
|
| 142 |
+
|
| 143 |
+
st.sidebar.write('Probability:',probability,'%')
|
| 144 |
+
|
| 145 |
+
# Retrieve predicted pokemon information
|
| 146 |
+
_, pokedex_number, english_name, romaji_name, katakana_name, weight_kg, height_m, type1, type2, color1, color2, classification, evolve_from, evolve_into, is_legendary = pokemon_info_df[pokemon_info_df['name']==prediction].values[0]
|
| 147 |
+
with col2:
|
| 148 |
+
# pokedex box
|
| 149 |
+
with st.container(border=True ):
|
| 150 |
+
# first row
|
| 151 |
+
with st.container():
|
| 152 |
+
pokeball_image_col,pokedex_number_col, pokemon_name_col = st.columns([1,1,8])
|
| 153 |
+
pokeball_image_col.image(pokeball_image)
|
| 154 |
+
pokedex_number_col.markdown(f'<div style="text-align: left; font-size: 1.4rem;"><b>{pokedex_number}</b></div>', unsafe_allow_html=True)
|
| 155 |
+
pokemon_name_col.markdown(f'<div style="text-align: right; font-size: 1.4rem;"><b>{english_name}</b></div>', unsafe_allow_html=True)
|
| 156 |
+
|
| 157 |
+
# second row
|
| 158 |
+
with st.container():
|
| 159 |
+
st.markdown(f'<div style="text-align: center; color: {color1}; font-size: 1.2rem;"><b>{classification}</b></div>', unsafe_allow_html=True)
|
| 160 |
+
|
| 161 |
+
# 3rd row
|
| 162 |
+
with st.container():
|
| 163 |
+
if pd.isna(type2):
|
| 164 |
+
st.write('\n')
|
| 165 |
+
st.markdown(f'<div style="display: flex; justify-content: center; align-items: center; "><div style="display: inline-block; padding: 5px; margin: 0 5px; border-radius: 5px; background-color: {color1}; color: white;">{type1}</div>', unsafe_allow_html=True)
|
| 166 |
+
else:
|
| 167 |
+
type1_col, type2_col = st.columns(2)
|
| 168 |
+
type1_col.markdown(f'<div style="display: flex; justify-content: center; align-items: center;"><div style="display: inline-block; padding: 5px; margin: 0 5px; border-radius: 5px; background-color: {color1}; color: white;">{type1}</div>', unsafe_allow_html=True)
|
| 169 |
+
type2_col.markdown(f'<div style="display: flex; justify-content: center; align-items: center;"><div style="display: inline-block; padding: 5px; margin: 0 5px; border-radius: 5px; background-color: {color2}; color: white;">{type2}</div>', unsafe_allow_html=True)
|
| 170 |
+
st.write('\n')
|
| 171 |
+
# 4th row
|
| 172 |
+
with st.container():
|
| 173 |
+
st.write(f'<div style=font-size: 1.4rem;><b>Height:</b> {height_m}m', unsafe_allow_html=True)
|
| 174 |
+
st.write('\n')
|
| 175 |
+
st.write(f'<div style=font-size: 1.4rem;><b>Weight:</b> {weight_kg}kg', unsafe_allow_html=True)
|
| 176 |
+
st.write('\n')
|
| 177 |
+
if not pd.isna(evolve_from):
|
| 178 |
+
st.markdown(f'<div style=font-size: 1.4rem;><b>Evolves from:</b> {evolve_from}', unsafe_allow_html=True)
|
| 179 |
+
#st.write(f'Evolves from: {evolve_from}')
|
| 180 |
+
st.write('\n')
|
| 181 |
+
if not pd.isna(evolve_into):
|
| 182 |
+
st.markdown(f'<div style=font-size: 1.4rem;><b>Evolves into:</b> {evolve_into}', unsafe_allow_html=True)
|
| 183 |
+
#st.write(f'Evolves into: {evolve_into}')
|
| 184 |
+
st.write('\n')
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
|