Spaces:
Running
Running
import json | |
from PIL import Image | |
import os | |
import streamlit as st | |
import pandas as pd | |
import numpy as np | |
import tensorflow as tf | |
import tensorflow_hub as hub | |
from tensorflow.keras.utils import img_to_array, load_img | |
def load_model(path='models/inceptionV3_model'): | |
"""Retrieves the trained model and maps it to the CPU by default, | |
can also specify GPU here.""" | |
model = tf.keras.Sequential([hub.KerasLayer(path)]) | |
model.build([None, 224, 224, 3]) | |
return model | |
def load_index_to_label_dict( | |
path: str = 'index_to_class_label.json' | |
) -> dict: | |
"""Retrieves and formats the | |
index to class label | |
lookup dictionary needed to | |
make sense of the predictions. | |
When loaded in, the keys are strings, this also | |
processes those keys to integers.""" | |
with open(path, 'r') as f: | |
index_to_class_label_dict = json.load(f) | |
index_to_class_label_dict = { | |
int(k): v for k, v in index_to_class_label_dict.items()} | |
return index_to_class_label_dict | |
def load_files( | |
keys: list, | |
path: str = 'images' | |
) -> list: | |
"""Retrieves files from images folder""" | |
files = [] | |
for file in os.listdir(path): | |
files.append(Image.open(path + '/' + file)) | |
return files | |
def load_image( | |
filename: str = '0.jpg', | |
path: str = 'images' | |
) -> list: | |
"""return image with path and filename""" | |
return Image.open(path + '/' + filename) | |
def predict( | |
img: Image.Image, | |
index_to_label_dict: dict, | |
model, | |
k: int | |
) -> list: | |
""" | |
This function transforms the image accordingly, | |
puts it to the necessary device (cpu by default here), | |
feeds the image through the model getting the output tensor, | |
converts that output tensor to probabilities using Softmax, | |
and then extracts and formats the top k predictions. | |
""" | |
img = img.resize((224, 224)) | |
image = img_to_array(img) | |
image = np.expand_dims(image, axis=0) | |
preds = model.predict(image) | |
formatted_predictions = list() | |
for idx, x in zip(range(0,8), preds[0]): | |
formatted_predictions.append([index_to_class_label_dict[idx], np.round(x*100,3)]) | |
formatted_predictions = sorted(formatted_predictions, key=lambda x:(x[1]), reverse=True) | |
return formatted_predictions | |
if __name__ == '__main__': | |
model = load_model() | |
index_to_class_label_dict = load_index_to_label_dict() | |
st.title('Rate My Photo!') | |
instructions = """ | |
Either upload your own image or select from | |
the sidebar to get a preconfigured image. | |
The image you select or upload will be fed | |
through the Deep Neural Network in real-time | |
and the output will be displayed to the screen. | |
""" | |
st.write(instructions) | |
file = st.file_uploader('Upload An Image') | |
images_type = { | |
'Bright': '0.jpg', | |
'Closed Eyes': '1.jpg', | |
'Dark': '2.jpg', | |
'Good': '3.jpg', | |
'Lens Flare': '4.jpg', | |
'Loss': '5.jpg', | |
'Motion Blur': '6.jpg', | |
'People Blur': '7.jpg' | |
} | |
data_split_names = list(images_type.keys()) | |
if file: # if user uploaded file | |
img = Image.open(file) | |
prediction = predict(img, index_to_class_label_dict, model, k=7) | |
else: | |
image_type = st.sidebar.selectbox("Examples", data_split_names) | |
image_file = images_type[image_type] | |
example_images = load_files(keys=images_type) | |
img = load_image(image_file) | |
prediction = predict(img, index_to_class_label_dict, model, 7) | |
st.title("Here is the image you've selected") | |
resized_image = img.resize((336, 336)) | |
st.image(resized_image) | |
if prediction[0][0] == "Good": | |
text = "good" | |
else: | |
text = "bad" | |
st.title("Your photo is in " + text + " quality") | |
df = pd.DataFrame(data=np.zeros((7, 2)), | |
columns=['Reason', 'Confidence Level'], | |
index=np.linspace(1, 7, 7, dtype=int)) | |
for idx, p in enumerate(prediction[:7]): | |
df.iloc[idx, 0] = p[0] | |
df.iloc[idx, 1] = str(p[1]) + '%' | |
st.write(df.to_html(escape=False), unsafe_allow_html=True) | |
st.write('\n') | |
credits = """ | |
This project was developed by <a href="https://www.linkedin.com/in/eran-perelman/" style="text-decoration: none">Eran Perelman</a>, <a href="https://www.linkedin.com/in/asi-sheratzki/" style="text-decoration: none">Asi Sheratzki</a> and <a href="https://www.linkedin.com/in/arykorenvais/" style="text-decoration: none">Ary Korenvais</a> with the guidance of <a href="https://www.linkedin.com/in/morris-alper/" style="text-decoration: none">Morris Alper</a>. | |
""" | |
st.write(credits, unsafe_allow_html=True) | |