import os

import streamlit as st
from PIL import Image

from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.preprocessing.image import img_to_array


def predict(file):
    st.markdown(f'{file.name} をアップロードしました.')
    img_path = os.path.join(file.name)

    # 画像を保存する
    with open(img_path, 'wb') as f:
        f.write(file.read())

    # 保存した画像を表示
    img = Image.open(img_path)
    st.image(img)

    # 画像をArray形式に変換
    img = load_img(img_path, target_size=(256, 256))
    img_array = img_to_array(img)
    img_array = img_array.reshape((1, 256, 256, 3))
    img_array = img_array / 255

    # 保存したモデルを呼び出し
    model_path = os.path.join('model.h5')
    model = load_model(model_path)
    result = model.predict(img_array)

    if result[0][0] > result[0][1]:
        prediction = '猫'
    else:
        prediction = '犬'
    return prediction


def show_result(prediction):
    if st.button("判定"):
        st.text_area("判定結果:", prediction, height=20)


def main():
    st.title("AI MNIST")
    file = st.file_uploader('画像をアップロードしてください.', type=['jpg', 'jpeg', 'png'])
    if file:
        prediction = predict(file)
        show_result(prediction)


if __name__ == '__main__':
    main()