Update app.py

app.py

@@ -3,12 +3,7 @@ import tensorflow as tf
 import numpy as np
 import cv2
 from PIL import Image
-from tensorflow.keras import layers, models
-from tensorflow.keras.applications import EfficientNetB0
-from tensorflow.keras.applications.efficientnet import preprocess_input
-import joblib
 import io
-import os
 
 # Set page config
 st.set_page_config(
@@ -17,7 +12,7 @@ st.set_page_config(
     layout="wide"
 )
 
-# Custom CSS with improved styling
+# Custom CSS to improve the appearance
 st.markdown("""
     <style>
     .main {
@@ -27,10 +22,14 @@ st.markdown("""
         width: 100%;
         margin-top: 1rem;
     }
+    .upload-text {
+        text-align: center;
+        padding: 2rem;
+    }
     .prediction-card {
         padding: 2rem;
         border-radius: 0.5rem;
-        background-color: #d7d7d9;
+        background-color: #f0f2f6;
         margin: 1rem 0;
     }
     .top-predictions {
@@ -40,154 +39,150 @@ st.markdown("""
         border-radius: 0.5rem;
         box-shadow: 0 1px 3px rgba(0,0,0,0.12);
     }
-    .survey-card {
-        padding: 1rem;
-        background-color: #f0f2f6;
-        border-radius: 0.5rem;
-        margin-top: 1rem;
+    .prediction-bar {
+        display: flex;
+        align-items: center;
+        margin: 0.5rem 0;
+    }
+    .prediction-label {
+        width: 100px;
+        font-weight: 500;
     }
     </style>
     """, unsafe_allow_html=True)
 
-# from mega import Mega
-
-# # Đăng nhập vào tài khoản Mega
-# def upload_to_mega(file_path, folder_name):
-#     """
-#     Upload file to a specific folder on Mega.nz
-#     """
-#     try:
-#         # Đăng nhập vào tài khoản Mega
-#         mega = Mega()
-#         m = mega.login(os.getenv('EMAIL'), os.getenv('PASSWORD'))
-
-#         # Tìm thư mục đích
-#         folder = m.find(folder_name)
-
-#         if not folder:
-#             # Nếu thư mục không tồn tại, hiển thị thông báo lỗi
-#             return f"Thư mục '{folder_name}' không tồn tại!"
-
-#         # Tải tệp lên thư mục
-#         file = m.upload(file_path, folder[0])
-#         return f"Upload thành công! Link: {m.get_upload_link(file)}"
-
-#     except Exception as e:
-#         return f"Lỗi khi tải lên Mega: {str(e)}"
-
+@st.cache_resource
+def load_model():
+    """Load the trained model"""
+    return tf.keras.models.load_model('custom_model.h5')
+
+def preprocess_image(image):
+    """Preprocess the uploaded image"""
+    # # Convert to RGB if needed
+    # if image.mode != 'RGB':
+    #     image = image.convert('RGB')
+    
+    # Convert to numpy array
+    img_array = np.array(image)
+    
+    # # Convert to RGB if needed
+    # if len(img_array.shape) == 2:  # Grayscale
+    #     img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
+    # elif img_array.shape[2] == 4:  # RGBA
+    #     img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
+    
+    # # Preprocess image similar to training
+    # img_hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
+    # img_hsv[:, :, 2] = cv2.equalizeHist(img_hsv[:, :, 2])
+    # img_array = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
+    
+    # # Adjust brightness
+    # target_brightness = 150
+    # current_brightness = np.mean(img_array)
+    # alpha = target_brightness / (current_brightness + 1e-5)
+    # img_array = cv2.convertScaleAbs(img_array, alpha=alpha, beta=0)
+    
+    # # Apply Gaussian blur
+    # img_array = cv2.GaussianBlur(img_array, (5, 5), 0)
+    
+    # Resize
+    img_array = cv2.resize(img_array, (256, 256))
+    
+    # Normalize
+    img_array = img_array.astype('float32') / 255.0
+    
+    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")
-
-    # Load model and scaler
-    model, scaler = load_model_and_scaler()
-    if model is None or scaler is None:
-        st.error("Failed to load model or scaler. Please ensure both files exist.")
-        return
-
-    # Initialize session state
+    
+    # Initialize session state for prediction if not exists
     if 'predictions' not in st.session_state:
         st.session_state.predictions = None
-    if 'uploaded_image' not in st.session_state:
-        st.session_state.uploaded_image = None
-
+    
+    # Create two columns
     col1, col2 = st.columns(2)
-
+    
     with col1:
         st.subheader("Upload Image")
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
+        
         if uploaded_file is not None:
-            try:
-                image = Image.open(uploaded_file)
-                st.image(image, caption="Uploaded Image", use_column_width=True)
-                st.session_state.uploaded_image = image
-
-                with st.spinner('Analyzing image...'):
-                    processed_image = preprocess_image(image, scaler)
-                    prediction = model.predict(processed_image, verbose=0)
-
+            # Display uploaded image
+            image = Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+
+            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', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
-                    st.session_state.predictions = get_top_predictions(prediction, class_names)
-
-            except Exception as e:
-                st.error(f"Error processing image: {str(e)}")
-
+                    
+                    # 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.predictions:
-            # Display main prediction
-            top_class, top_confidence = st.session_state.predictions[0]
-            st.markdown(
-                f"""
-                <div class='prediction-card'>
-                    <h3>Primary Prediction: Grade {top_class}</h3>
-                    <h3>Confidence: {top_confidence:.2f}%</h3>
-                </div>
-                """,
-                unsafe_allow_html=True
-            )
-
-            # Display confidence bar
-            st.progress(top_confidence / 100)
-
-            # Display top 5 predictions
-            st.markdown("### Top 5 Predictions")
-            st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
-
-            for class_name, confidence in st.session_state.predictions:
-                cols = st.columns([2, 6, 2])
-                with cols[0]:
-                    st.write(f"Grade {class_name}")
-                with cols[1]:
-                    st.progress(confidence / 100)
-                with cols[2]:
-                    st.write(f"{confidence:.2f}%")
-
-            st.markdown("</div>", unsafe_allow_html=True)
-
-            # User Survey
-            st.markdown("<div class='survey-card'>", unsafe_allow_html=True)
-            st.markdown("### Model Accuracy Survey")
-            st.write("Mô hình có dự đoán chính xác màu sắc của đá trong ảnh này không?")
-
-            # Accuracy Confirmation
-            accuracy = st.radio(
-                "Đánh giá độ chính xác",
-                ["Chọn", "Chính xác", "Không chính xác"],
-                index=0
-            )
-
-            if accuracy == "Không chính xác":
-              # Color input for incorrect prediction
-              correct_color = st.text_input(
-                  "Vui lòng nhập màu sắc chính xác của đá:",
-                  help="Ví dụ: 10, 9.7, 9.5, 9.2, v.v."
-              )
-
-              # if st.button("Gửi phản hồi và tải ảnh"):
-              #     if correct_color and st.session_state.uploaded_image:
-              #         # Save the image temporarily
-              #         temp_image_path = f"temp_image_{hash(uploaded_file.name)}.png"
-              #         st.session_state.uploaded_image.save(temp_image_path)
-
-              #         # Upload to Mega.nz
-              #         upload_result = upload_to_mega(temp_image_path, correct_color)
-              #         if "Upload thành công" in upload_result:
-              #             st.success(upload_result)
-              #         else:
-              #             st.error(upload_result)
-
-              #         # Clean up temporary file
-              #         os.remove(temp_image_path)
-              #     else:
-              #         st.warning("Vui lòng nhập màu sắc chính xác")
-
-            st.markdown("</div>", unsafe_allow_html=True)
+        if st.session_state.predictions is not None:
+            # Create a card-like container for results
+            results_container = st.container()
+            with results_container:
+                # Display main prediction
+                st.markdown("<div class='prediction-card'>", unsafe_allow_html=True)
+                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)
+                
+                # 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 to see the predictions")
-
+            st.info("Upload an image and click 'Predict' to see the results")
+    
+    # Footer
     st.markdown("---")
     st.markdown("Made with ❤️ using Streamlit")