Update app.py

app.py

@@ -3,21 +3,12 @@ import tensorflow as tf
 import numpy as np
 import cv2
 from PIL import Image
-import io
-import os
-import cv2
-import numpy as np
-from tensorflow import keras
-from tensorflow.keras import layers, models
-from sklearn.preprocessing import StandardScaler
-from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
-from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
-import matplotlib.pyplot as plt
-import random
 from tensorflow.keras import layers, models
 from tensorflow.keras.applications import EfficientNetB0
 from tensorflow.keras.applications.efficientnet import preprocess_input
-from tensorflow.keras.layers import Lambda  # Đảm bảo nhập Lambda từ tensorflow.keras.layers
+from sklearn.preprocessing import StandardScaler
+import io
+
 # Set page config
 st.set_page_config(
     page_title="Stone Classification",
@@ -25,7 +16,7 @@ st.set_page_config(
     layout="wide"
 )
 
-# Custom CSS to improve the appearance
+# Custom CSS with improved styling
 st.markdown("""
     <style>
     .main {
@@ -35,10 +26,6 @@ st.markdown("""
         width: 100%;
         margin-top: 1rem;
     }
-    .upload-text {
-        text-align: center;
-        padding: 2rem;
-    }
     .prediction-card {
         padding: 2rem;
         border-radius: 0.5rem;
@@ -52,202 +39,159 @@ 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)
 
 @st.cache_resource
 def load_model():
     """Load the trained model"""
-    return tf.keras.models.load_model('mlp_model.h5')
-    
+    try:
+        return tf.keras.models.load_model('mlp_model.h5')
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None
+
 def color_histogram(image, bins=16):
-    # (Previous implementation remains the same)
+    """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()
     
-    hist_r = hist_r / np.sum(hist_r)
-    hist_g = hist_g / np.sum(hist_g)
-    hist_b = hist_b / np.sum(hist_b)
+    # Normalize histograms
+    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)
     
     return np.concatenate([hist_r, hist_g, hist_b])
 
 def color_moments(image):
-    # (Previous implementation remains the same)
+    """Calculate color moments features"""
     img = image.astype(np.float32) / 255.0
-    
     moments = []
-    for i in range(3):  # For each color channel
+    
+    for i in range(3):
         channel = img[:,:,i]
-        
         mean = np.mean(channel)
-        std = np.std(channel)
-        skewness = np.mean(((channel - mean) / std) ** 3)
-        
+        std = np.std(channel) + 1e-7  # Avoid division by zero
+        skewness = np.mean(((channel - mean) / std) ** 3) if std != 0 else 0
         moments.extend([mean, std, skewness])
     
     return np.array(moments)
 
 def dominant_color_descriptor(image, k=3):
-    # (Previous implementation remains the same)
-    pixels = image.reshape(-1, 3)
+    """Calculate dominant color descriptor"""
+    pixels = image.reshape(-1, 3).astype(np.float32)
     
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
     flags = cv2.KMEANS_RANDOM_CENTERS
     
     try:
-        _, labels, centers = cv2.kmeans(pixels.astype(np.float32), k, None, criteria, 10, flags)
-        
+        _, labels, centers = cv2.kmeans(pixels, k, None, criteria, 10, flags)
         unique, counts = np.unique(labels, return_counts=True)
         percentages = counts / len(labels)
-        
-        dominant_colors = centers.flatten()
-        color_percentages = percentages
-        
-        return np.concatenate([dominant_colors, color_percentages])
-    except:
-        return np.zeros(2 * k)
+        return np.concatenate([centers.flatten(), percentages])
+    except Exception:
+        return np.zeros(k * 4)  # Return zero vector if clustering fails
 
 def color_coherence_vector(image, k=3):
-    # Convert to grayscale
+    """Calculate color coherence vector"""
     gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-    
-    # Convert the grayscale image to 8-bit format before applying threshold
     gray = np.uint8(gray)
     
-    # Apply Otsu's thresholding method
     _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-    
-    # Perform connected components analysis
     num_labels, labels = cv2.connectedComponents(binary)
     
     ccv = []
     for i in range(1, min(k+1, num_labels)):
         region_mask = (labels == i)
         total_pixels = np.sum(region_mask)
-        coherent_pixels = total_pixels
-        
-        ccv.extend([coherent_pixels, total_pixels])
-    
-    while len(ccv) < 2 * k:
-        ccv.append(0)
-    
-    return np.array(ccv)
-
-
-# ViT and Feature Extraction Functions (from previous implementation)
-# (Keeping the Patches, PatchEncoder, and create_vit_feature_extractor functions)
-
-def extract_features(image):
-    """
-    Extract multiple features from an image
-    """
-    color_hist = color_histogram(image)
-    color_mom = color_moments(image)
-    dom_color = dominant_color_descriptor(image)
-    ccv = color_coherence_vector(image)
+        ccv.extend([total_pixels, total_pixels])
     
-    return np.concatenate([color_hist, color_mom, dom_color, ccv])
-
-from transformers import ViTFeatureExtractor, ViTModel
-import torch
-from tensorflow.keras import layers, models
+    # Pad with zeros if needed
+    ccv.extend([0] * (2 * k - len(ccv)))
+    return np.array(ccv[:2*k])
 
-def create_vit_feature_extractor(input_shape=(256, 256, 3), num_classes=None):
-    # Xây dựng mô hình ViT đã huấn luyện sẵn từ TensorFlow
+@st.cache_resource
+def create_feature_extractor():
+    """Create and cache the feature extractor model"""
+    input_shape = (256, 256, 3)
     inputs = layers.Input(shape=input_shape)
+    x = layers.Lambda(preprocess_input)(inputs)
     
-    # Thêm lớp Lambda để tiền xử lý ảnh
-    x = Lambda(preprocess_input, output_shape=input_shape)(inputs)  # Xử lý ảnh đầu vào
-    
-    # Bạn có thể thay thế phần này bằng một mô hình ViT đã được huấn luyện sẵn.
-    # Dưới đây là ví dụ dùng EfficientNetB0 thay vì ViT.
-    # Tạo mô hình ViT hoặc sử dụng mô hình khác đã được huấn luyện sẵn
-    vit_model = EfficientNetB0(include_top=False, weights='imagenet', input_tensor=x)
-    
-    # Trích xuất đặc trưng từ mô hình ViT
-    x = layers.GlobalAveragePooling2D()(vit_model.output)
-    
-    if num_classes:
-        x = layers.Dense(num_classes, activation='softmax')(x)  # Thêm lớp phân loại (nếu có)
+    base_model = EfficientNetB0(
+        include_top=False,
+        weights='imagenet',
+        input_tensor=x
+    )
     
+    x = layers.GlobalAveragePooling2D()(base_model.output)
     return models.Model(inputs=inputs, outputs=x)
+
+def extract_features(image):
+    """Extract all features from an image"""
+    # Convert image to uint8 for OpenCV operations
+    image_uint8 = (image * 255).astype(np.uint8)
     
+    # Extract traditional features
+    hist_features = color_histogram(image_uint8)
+    moment_features = color_moments(image_uint8)
+    dominant_features = dominant_color_descriptor(image_uint8)
+    ccv_features = color_coherence_vector(image_uint8)
+    
+    return np.concatenate([
+        hist_features,
+        moment_features,
+        dominant_features,
+        ccv_features
+    ])
+
 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
+    # Convert to numpy array and resize
     img_array = np.array(image)
-    
-    # Ensure RGB format
-    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)
-    
-    # Resize
     img_array = cv2.resize(img_array, (256, 256))
-    
-    # Normalize
     img_array = img_array.astype('float32') / 255.0
     
     # Extract traditional features
-    image_features = extract_features(img_array)
+    traditional_features = extract_features(img_array)
     
-    # Create and process ViT features
-    vit_extractor = create_vit_feature_extractor()
-    
-    # Reshape image for ViT processing - THIS IS THE KEY FIX
-    img_for_vit = np.expand_dims(img_array, axis=0)  # Add batch dimension
-    image_vit = vit_extractor.predict(img_for_vit)  
-    
-    # Flatten ViT features if needed
-    image_vit = image_vit.reshape(1, -1)  # Ensure 2D shape
+    # Extract deep features
+    feature_extractor = create_feature_extractor()
+    deep_features = feature_extractor.predict(
+        np.expand_dims(img_array, axis=0),
+        verbose=0
+    )
     
     # Combine features
-    image_combined = np.concatenate([image_features.reshape(1, -1), image_vit], axis=1)
+    combined_features = np.concatenate([
+        traditional_features.reshape(1, -1),
+        deep_features.reshape(1, -1)
+    ], axis=1)
     
     # Scale features
     scaler = StandardScaler()
-    image_scaled = scaler.fit_transform(image_combined)
-    
-    return image_scaled.squeeze()  # Remove any unnecessary dimensions
+    return scaler.fit_transform(combined_features)
 
 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
+    # 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:
@@ -255,69 +199,60 @@ 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
+            try:
+                image = Image.open(uploaded_file)
+                st.image(image, caption="Uploaded Image", use_column_width=True)
+                
+                with st.spinner('Analyzing image...'):
                     model = load_model()
+                    if model is None:
+                        st.error("Failed to load model")
+                        return
                     
-                    # Preprocess image
                     processed_image = preprocess_image(image)
+                    prediction = model.predict(processed_image, verbose=0)
                     
-                    # Ensure correct shape for prediction
-                    processed_image = np.expand_dims(processed_image, axis=0)
-                    
-                    # Make prediction
-                    prediction = model.predict(processed_image)
                     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")
-        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)
+        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)
         else:
-            st.info("Upload an image and click 'Predict' to see the results")
+            st.info("Upload an image to see the predictions")
     
-    # Footer
     st.markdown("---")
     st.markdown("Made with ❤️ using Streamlit")