Update app.py

app.py

@@ -3,6 +3,10 @@ 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
 
 # Set page config
@@ -12,7 +16,7 @@ st.set_page_config(
     layout="wide"
 )
 
-# Custom CSS to improve the appearance
+# Custom CSS with improved styling
 st.markdown("""
     <style>
     .main {
@@ -22,14 +26,10 @@ st.markdown("""
         width: 100%;
         margin-top: 1rem;
     }
-    .upload-text {
-        text-align: center;
-        padding: 2rem;
-    }
     .prediction-card {
         padding: 2rem;
         border-radius: 0.5rem;
-        background-color: #f0f2f6;
+        background-color: #d7d7d9;
         margin: 1rem 0;
     }
     .top-predictions {
@@ -39,109 +39,162 @@ st.markdown("""
         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)
 
+# Cache the model loading
 @st.cache_resource
-def load_model():
-    """Load the trained model"""
-    return tf.keras.models.load_model('custom_model.h5')
+def load_model_and_scaler():
+    """Load the trained model and scaler"""
+    try:
+        model = tf.keras.models.load_model('mlp_model.h5')
+        # Tải scaler đã lưu
+        scaler = joblib.load('standard_scaler.pkl')
+        return model, scaler
+    except Exception as e:
+        st.error(f"Error loading model or scaler: {str(e)}")
+        return None, None
 
-def preprocess_image(image):
-    """Preprocess the uploaded image"""
-    # # Convert to RGB if needed
-    # if image.mode != 'RGB':
-    #     image = image.convert('RGB')
+def color_histogram(image, bins=16):
+    """Calculate color histogram features"""
+    hist_r = cv2.calcHist([image], [0], None, [bins], [0, 256]).flatten()
+    hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten()
+    hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten()
     
-    # Convert to numpy array
-    img_array = np.array(image)
+    hist_r = hist_r / (np.sum(hist_r) + 1e-7)
+    hist_g = hist_g / (np.sum(hist_g) + 1e-7)
+    hist_b = hist_b / (np.sum(hist_b) + 1e-7)
     
-    # # 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)
+    return np.concatenate([hist_r, hist_g, hist_b])
+
+def color_moments(image):
+    """Calculate color moments features"""
+    img = image.astype(np.float32) / 255.0
+    moments = []
     
-    # # 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)
+    for i in range(3):
+        channel = img[:,:,i]
+        mean = np.mean(channel)
+        std = np.std(channel) + 1e-7
+        skewness = np.mean(((channel - mean) / std) ** 3) if std != 0 else 0
+        moments.extend([mean, std, skewness])
     
-    # # 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)
+    return np.array(moments)
+
+def dominant_color_descriptor(image, k=3):
+    """Calculate dominant color descriptor"""
+    pixels = image.reshape(-1, 3).astype(np.float32)
     
-    # # Apply Gaussian blur
-    # img_array = cv2.GaussianBlur(img_array, (5, 5), 0)
+    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
+    flags = cv2.KMEANS_RANDOM_CENTERS
     
-    # Resize
-    img_array = cv2.resize(img_array, (256, 256))
+    try:
+        _, labels, centers = cv2.kmeans(pixels, k, None, criteria, 10, flags)
+        unique, counts = np.unique(labels, return_counts=True)
+        percentages = counts / len(labels)
+        return np.concatenate([centers.flatten(), percentages])
+    except Exception:
+        return np.zeros(k * 4)
+
+def color_coherence_vector(image, k=3):
+    """Calculate color coherence vector"""
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    gray = np.uint8(gray)
     
-    # Normalize
-    img_array = img_array.astype('float32') / 255.0
+    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    num_labels, labels = cv2.connectedComponents(binary)
     
-    return img_array
-
-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('[email protected]', '01283315889')
+    ccv = []
+    for i in range(1, min(k+1, num_labels)):
+        region_mask = (labels == i)
+        total_pixels = np.sum(region_mask)
+        ccv.extend([total_pixels, total_pixels])
+    
+    ccv.extend([0] * (2 * k - len(ccv)))
+    return np.array(ccv[:2*k])
 
-        # Tìm thư mục đích
-        folder = m.find(folder_name)
+@st.cache_resource
+def create_vit_feature_extractor():
+    """Create and cache the ViT feature extractor"""
+    input_shape = (256, 256, 3)
+    inputs = layers.Input(shape=input_shape)
+    x = layers.Lambda(preprocess_input)(inputs)
+    
+    base_model = EfficientNetB0(
+        include_top=False,
+        weights='imagenet',
+        input_tensor=x
+    )
+    
+    x = layers.GlobalAveragePooling2D()(base_model.output)
+    return models.Model(inputs=inputs, outputs=x)
 
-        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!"
+def extract_features(image):
+    """Extract all features from an image"""
+    # Traditional features
+    hist_features = color_histogram(image)
+    moment_features = color_moments(image)
+    dominant_features = dominant_color_descriptor(image)
+    ccv_features = color_coherence_vector(image)
+    
+    traditional_features = np.concatenate([
+        hist_features,
+        moment_features,
+        dominant_features,
+        ccv_features
+    ])
+    
+    # Deep features using ViT
+    feature_extractor = create_vit_feature_extractor()
+    vit_features = feature_extractor.predict(
+        np.expand_dims(image, axis=0),
+        verbose=0
+    )
+    
+    # Combine all features
+    return np.concatenate([traditional_features, vit_features.flatten()])
 
-        # 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)}"
+def preprocess_image(image, scaler):
+    """Preprocess the uploaded image"""
+    # Convert to RGB if needed
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
+    # Convert to numpy array and resize
+    img_array = np.array(image)
+    img_array = cv2.resize(img_array, (256, 256))
+    img_array = img_array.astype('float32') / 255.0
+    
+    # Extract all features
+    features = extract_features(img_array)
+    
+    # Scale features using the provided scaler
+    scaled_features = scaler.transform(features.reshape(1, -1))
+    
+    return scaled_features
 
-    except Exception as e:
-        return f"Lỗi khi tải lên Mega: {str(e)}"
-        
 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 = [
+    return [
         (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
+    # 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
     if 'predictions' not in st.session_state:
         st.session_state.predictions = None
     
-    # Create two columns
     col1, col2 = st.columns(2)
     
     with col1:
@@ -149,30 +202,19 @@ def main():
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
         
         if uploaded_file is not None:
-            # 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)
+            try:
+                image = Image.open(uploaded_file)
+                st.image(image, caption="Uploaded Image", use_column_width=True)
+                
+                with st.spinner('Analyzing image...'):
+                    processed_image = preprocess_image(image, scaler)
+                    prediction = model.predict(processed_image, verbose=0)
                     
-                    # 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)
                     
-                    # 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)}")
+            except Exception as e:
+                st.error(f"Error processing image: {str(e)}")
     
     with col2:
         st.subheader("Prediction Results")
@@ -188,14 +230,14 @@ def main():
                 """,
                 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]:
@@ -204,50 +246,11 @@ def main():
                     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)
         else:
             st.info("Upload an image to see the predictions")
-
+    
     st.markdown("---")
     st.markdown("Made with ❤️ using Streamlit")