Update app.py

app.py

@@ -3,6 +3,7 @@ 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'):
@@ -17,6 +18,7 @@ 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)
@@ -24,22 +26,31 @@ def f1(inp, targ, average='macro'):
 而在上傳圖片之後，模型將會為您分類出適合的類別與其對應的機率
 """
 
+# 顯示示例圖片
+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(learn_inf.predict(img))
-
-    f"""
+    st.write(f"""
     ### 預測的類別為： {pred}
-    ### 符合程度： {probs[key].item()*100: .3f}%
-    """
-
+    ### 符合程度： {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.")
 
@@ -47,15 +58,15 @@ if option == "Upload Image":
         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)
-            
+            st.text("Problem reading image from", url)