Update app.py

app.py

@@ -26,6 +26,28 @@ st.markdown("""
         text-align: center;
         padding: 2rem;
     }
+    .prediction-card {
+        padding: 2rem;
+        border-radius: 0.5rem;
+        background-color: #f0f2f6;
+        margin: 1rem 0;
+    }
+    .top-predictions {
+        margin-top: 2rem;
+        padding: 1rem;
+        background-color: white;
+        border-radius: 0.5rem;
+        box-shadow: 0 1px 3px rgba(0,0,0,0.12);
+    }
+    .prediction-bar {
+        display: flex;
+        align-items: center;
+        margin: 0.5rem 0;
+    }
+    .prediction-label {
+        width: 100px;
+        font-weight: 500;
+    }
     </style>
     """, unsafe_allow_html=True)
 
@@ -71,16 +93,27 @@ def preprocess_image(image):
     
     return img_array
 
+def get_top_predictions(prediction, class_names, top_k=5):
+    """Get top k predictions with their probabilities"""
+    # Get indices of top k predictions
+    top_indices = prediction.argsort()[0][-top_k:][::-1]
+    
+    # Get corresponding class names and probabilities
+    top_predictions = [
+        (class_names[i], float(prediction[0][i]) * 100)
+        for i in top_indices
+    ]
+    
+    return top_predictions
+
 def main():
     # Title
     st.title("🪨 Stone Classification")
     st.write("Upload an image of a stone to classify its type")
     
     # Initialize session state for prediction if not exists
-    if 'prediction' not in st.session_state:
-        st.session_state.prediction = None
-    if 'confidence' not in st.session_state:
-        st.session_state.confidence = None
+    if 'predictions' not in st.session_state:
+        st.session_state.predictions = None
     
     # Create two columns
     col1, col2 = st.columns(2)
@@ -96,52 +129,58 @@ def main():
             
             # Add predict button
             if st.button("Predict"):
-                try:
-                    # Load model
-                    model = load_model()
-                    
-                    # Preprocess image
-                    processed_image = preprocess_image(image)
-                    
-                    # Make prediction
-                    prediction = model.predict(np.expand_dims(processed_image, axis=0))
-                    class_names = ['10', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
-                    
-                    # Get prediction and confidence
-                    predicted_class = class_names[np.argmax(prediction)]
-                    confidence = float(np.max(prediction)) * 100
-                    
-                    # Store in session state
-                    st.session_state.prediction = predicted_class
-                    st.session_state.confidence = confidence
-                    
-                except Exception as e:
-                    st.error(f"Error during prediction: {str(e)}")
+                with st.spinner('Analyzing image...'):
+                    try:
+                        # Load model
+                        model = load_model()
+                        
+                        # Preprocess image
+                        processed_image = preprocess_image(image)
+                        
+                        # Make prediction
+                        prediction = model.predict(np.expand_dims(processed_image, axis=0))
+                        class_names = ['10', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
+                        
+                        # Get top 5 predictions
+                        top_predictions = get_top_predictions(prediction, class_names)
+                        
+                        # Store in session state
+                        st.session_state.predictions = top_predictions
+                        
+                    except Exception as e:
+                        st.error(f"Error during prediction: {str(e)}")
     
     with col2:
         st.subheader("Prediction Results")
-        if st.session_state.prediction is not None:
+        if st.session_state.predictions is not None:
             # Create a card-like container for results
             results_container = st.container()
             with results_container:
-                st.markdown("""
-                    <style>
-                    .prediction-card {
-                        padding: 2rem;
-                        border-radius: 0.5rem;
-                        background-color: #f0f2f6;
-                        margin: 1rem 0;
-                    }
-                    </style>
-                    """, unsafe_allow_html=True)
-                
+                # Display main prediction
                 st.markdown("<div class='prediction-card'>", unsafe_allow_html=True)
-                st.markdown(f"### Predicted Class: {st.session_state.prediction}")
-                st.markdown(f"### Confidence: {st.session_state.confidence:.2f}%")
+                top_class, top_confidence = st.session_state.predictions[0]
+                st.markdown(f"### Primary Prediction: Grade {top_class}")
+                st.markdown(f"### Confidence: {top_confidence:.2f}%")
                 st.markdown("</div>", unsafe_allow_html=True)
                 
-                # Add confidence bar
-                st.progress(st.session_state.confidence / 100)
+                # Display confidence bar for top prediction
+                st.progress(top_confidence / 100)
+                
+                # Display top 5 predictions
+                st.markdown("### Top 5 Predictions")
+                st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
+                
+                # Create a Streamlit container for the predictions
+                for class_name, confidence in st.session_state.predictions:
+                    col_label, col_bar, col_value = st.columns([2, 6, 2])
+                    with col_label:
+                        st.write(f"Grade {class_name}")
+                    with col_bar:
+                        st.progress(confidence / 100)
+                    with col_value:
+                        st.write(f"{confidence:.2f}%")
+                
+                st.markdown("</div>", unsafe_allow_html=True)
         else:
             st.info("Upload an image and click 'Predict' to see the results")