Update app.py

app.py

@@ -97,8 +97,8 @@ def model_description():
 def prediction():
 
     def load_model(model_name):
-        num_classes = 7  # Pastikan sesuai dengan jumlah kelas yang digunakan saat training
-    
+        num_classes = 7  # Sesuaikan dengan jumlah kelas mineral
+        
         if model_name == "DenseNet":
             model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
                                          filename="densenet_finetuned.pth")
@@ -109,7 +109,7 @@ def prediction():
             model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
                                          filename="mobileNetV2_finetuned.pth")
             model = models.mobilenet_v2(pretrained=False)
-            model.classifier = torch.nn.Linear(model.classifier[0].in_features, num_classes)  # Fix in_features
+            model.classifier = torch.nn.Linear(model.last_channel, num_classes)  # Perbaikan classifier
     
         elif model_name == "SqueezeNet":
             model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
@@ -164,35 +164,45 @@ def prediction():
             confidence, predicted = torch.max(probabilities, 1)
     
         return predicted.item(), confidence.item()
-
-
-    st.markdown("## 🍎 Fruit and Vegetable Classifier")
-    st.markdown("Upload an image of a fruit or vegetable, and I will classify it for you!")
-
+    
+    
+    # ====== Streamlit UI ======
+    
+    st.markdown("## 🔍 Mineral Classifier")
+    st.markdown("Upload an image of a mineral, and I will classify it for you!")
+    
     st.sidebar.header("Model Settings")
     model_name = st.sidebar.selectbox("Select Model", ("DenseNet", "SqueezeNet", "MobileNet"))
-
-    confidence_threshold = st.sidebar.number_input("Set Confidence Threshold", min_value=0.0, max_value=1.0, value=0.5,
-                                                   step=0.01)
-
+    
+    confidence_threshold = st.sidebar.number_input("Set Confidence Threshold", min_value=0.0, max_value=1.0, value=0.5, step=0.01)
+    
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
+    
     predict_button = st.sidebar.button("Predict")
-
+    
+    # Daftar label kelas mineral
+    class_labels = ["quartz", "pyrite", "muscovite", "malachite", "chrysocolla", "bornite", "biotite"]
+    
     if uploaded_file is not None and predict_button:
         with st.spinner("Processing the image..."):
             image = Image.open(uploaded_file)
-
+    
             model = load_model(model_name)
-            label, confidence = classify_image(model, image)
-
+            predicted_index, confidence = classify_image(model, image)
+    
+            # Konversi index prediksi ke nama mineral
+            if 0 <= predicted_index < len(class_labels):
+                predicted_label = class_labels[predicted_index]
+            else:
+                predicted_label = "Unknown"
+    
             col1, col2 = st.columns([1, 2])
-
+    
             with col1:
                 st.image(image, caption='Uploaded Image', width=250)
-
+    
             with col2:
-                st.write(f"**Prediction**: {predicted.capitalize()}")
+                st.write(f"**Prediction**: {predicted_label.capitalize()}")  # Perbaiki penggunaan capitalize()
                 st.write(f"**Confidence**: {confidence:.4f}")
 
 def contact():