from fastai.vision.all import *
from sklearn.metrics import precision_score, recall_score, f1_score
from io import BytesIO
import requests
import streamlit as st
import os

# 定義自定義度量函數
def precision(inp, targ, average='macro'):
    inp, targ = inp.cpu(), targ.cpu()  # 確保張量在 CPU 上
    return precision_score(targ, inp.argmax(dim=-1), average=average)

def recall(inp, targ, average='macro'):
    inp, targ = inp.cpu(), targ.cpu()  # 確保張量在 CPU 上
    return recall_score(targ, inp.argmax(dim=-1), average=average)

def f1(inp, targ, average='macro'):
    inp, targ = inp.cpu(), targ.cpu()  # 確保張量在 CPU 上
    return f1_score(targ, inp.argmax(dim=-1), average=average)

# 標題和描述
"""
# 層積岩的種類分類
此模型為對層積岩進行分類，類別有：碎屑岩(Clastic Rock)、火山碎屑岩(Pyroclastic Rock)、生物岩(Biolite)和化學岩(Chemical Sedimentary Rock)
使用者可以上傳自己想查閱的圖片進行查詢
而在上傳圖片之後，模型將會為您分類出適合的類別與其對應的機率
"""

# 顯示示例圖片
st.subheader("示例圖片")
example_dir = "./example"
example_images = os.listdir(example_dir)

for image_name in example_images:
    image_path = os.path.join(example_dir, image_name)
    st.image(image_path, caption=image_name, use_column_width=True)

def predict(img):
    st.image(img, caption="Your image", use_column_width=True)
    pred, key, probs = learn_inf.predict(img)
    st.write(f"""
    ### 預測的類別為： {pred}
    ### 符合程度： {probs[key].item()*100:.3f}%
    """)

# 加載模型
path = "./"
learn_inf = load_learner(path + "rockdetect.pkl")

# 選擇上傳方式：本地上傳或URL
option = st.radio("", ["Upload Image", "Image URL"])

# 處理本地圖片上傳
if option == "Upload Image":
    uploaded_file = st.file_uploader("Please upload an image.")

    if uploaded_file is not None:
        img = PILImage.create(uploaded_file)
        predict(img)

# 處理圖片URL
else:
    url = st.text_input("Please input a url.")

    if url != "":
        try:
            # 從URL下載圖片
            response = requests.get(url)
            pil_img = PILImage.create(BytesIO(response.content))
            predict(pil_img)
        except:
            st.text("Problem reading image from", url)