import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from PIL import Image

# Título do aplicativo
st.title("Classificador dos Chapéus de Palha - One Piece")

# Carregar o modelo treinado
@st.cache_resource  # Cache do modelo para melhorar performance
def load_my_model():
    return load_model('one_piece_v3.keras')

model = load_my_model()

# Lista de personagens
characters = ['Luffy', 'Zoro', 'Nami', 'Usopp', 'Sanji', 'Chopper', 'Robin', 'Franky', 'Brook', 'Jinbe']

# Função para prever o personagem em uma nova imagem
def predict_character(img):
    img = img.resize((128, 128))  # Redimensionar a imagem para o tamanho esperado pelo modelo
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array / 255.0  # Normalizar a imagem
    
    prediction = model.predict(img_array)
    predicted_class = np.argmax(prediction, axis=1)
    
    return characters[predicted_class[0]], prediction[0]

# Upload da imagem
uploaded_file = st.file_uploader("Escolha uma imagem...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Carregar a imagem
    img = Image.open(uploaded_file)
    
    # Exibir a imagem
    st.image(img, caption='Imagem carregada', use_column_width=True)
    
    # Fazer a previsão
    st.write("Classificando...")
    predicted_character, confidence_scores = predict_character(img)
    
    # Exibir o resultado
    st.success(f"O personagem na imagem é: **{predicted_character}**")