import streamlit as st
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import os
import tempfile
import shutil

# Coba Import YOLO
try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    
st.set_page_config(page_title="Betta Classifier")

# Periksa apakah library YOLO tersedia
def cek_library():
    if not YOLO_AVAILABLE:
        st.error("Ultralytics tidak terpasang. Silakan instal dengan perintah berikut:")
        st.code("pip install ultralytics")
        return False
    return True

st.markdown("""
<div style="background-color:#0984e3; padding: 20px; text-align: center;">
<h1 style="color: white;"> Betta Classifier Program </h1>
<h5 style="color: white;"> Betta Image Detection </h5>
</div>
""", unsafe_allow_html=True)

# Pastikan library sudah terpasang sebelum melanjutkan
if cek_library():
    uploaded_file = st.file_uploader("Upload gambar ikan cupang", type=['jpg', 'jpeg', 'png'])

    if uploaded_file:
        temp_dir = tempfile.mkdtemp()
        temp_file = os.path.join(temp_dir, "gambar.jpg")
        image = Image.open(uploaded_file)

        # Ubah Ukuran Gambar
        image = image.resize((300, 300))
        image.save(temp_file)

        # Tampilkan gambar
        st.markdown("<div style='text-align: center;'>", unsafe_allow_html=True)
        st.image(image, caption="Gambar yang diupload")
        st.markdown("</div>", unsafe_allow_html=True)

        # Deteksi Gambar
        if st.button("Deteksi Gambar"):
            with st.spinner("Sedang diproses"):
                try:
                    model = YOLO('bestt.pt')  # nama model kamu
                    hasil = model(temp_file)

                    nama_objek = hasil[0].names
                    nilai_prediksi = hasil[0].probs.data.numpy().tolist()
                    objek_terdeteksi = nama_objek[np.argmax(nilai_prediksi)]

                    fig, ax = plt.subplots()
                    ax.bar(list(nama_objek.values()), nilai_prediksi)
                    ax.set_title('Tingkat Keyakinan Prediksi')
                    ax.set_xlabel('Betta')
                    ax.set_ylabel('Keyakinan')
                    plt.xticks(rotation=45)

                    st.success(f"Betta terdeteksi: {objek_terdeteksi}")
                    st.pyplot(fig)

                except Exception as e:
                    st.error("Gambar tidak dapat terdeteksi")
                    st.error(f"Error: {e}")

                shutil.rmtree(temp_dir, ignore_errors=True)

st.markdown(
    "<div style='text-align: center;' class='footer'> Betta Detection Application Program </div>",
    unsafe_allow_html=True
)