Update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 from tensorflow.keras import layers, models
 from tensorflow.keras.applications import EfficientNetB0
 from tensorflow.keras.applications.efficientnet import preprocess_input
-from sklearn.preprocessing import StandardScaler
+import joblib
 import io
 
 # Set page config
@@ -42,14 +42,17 @@ st.markdown("""
     </style>
     """, unsafe_allow_html=True)
 
+# Cache the model loading
 @st.cache_resource
-def load_model():
-    """Load the trained model"""
+def load_model_and_scaler():
+    """Load the trained model and scaler"""
     try:
-        return tf.keras.models.load_model('mlp_model.h5')
+        model = tf.keras.models.load_model('mlp_model.h5')
+        scaler = joblib.load('scaler.save')
+        return model, scaler
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None
+        st.error(f"Error loading model or scaler: {str(e)}")
+        return None, None
 
 def color_histogram(image, bins=16):
     """Calculate color histogram features"""
@@ -57,7 +60,6 @@ def color_histogram(image, bins=16):
     hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten()
     hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten()
     
-    # 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)
@@ -72,7 +74,7 @@ def color_moments(image):
     for i in range(3):
         channel = img[:,:,i]
         mean = np.mean(channel)
-        std = np.std(channel) + 1e-7  # Avoid division by zero
+        std = np.std(channel) + 1e-7
         skewness = np.mean(((channel - mean) / std) ** 3) if std != 0 else 0
         moments.extend([mean, std, skewness])
     
@@ -91,7 +93,7 @@ def dominant_color_descriptor(image, k=3):
         percentages = counts / len(labels)
         return np.concatenate([centers.flatten(), percentages])
     except Exception:
-        return np.zeros(k * 4)  # Return zero vector if clustering fails
+        return np.zeros(k * 4)
 
 def color_coherence_vector(image, k=3):
     """Calculate color coherence vector"""
@@ -107,13 +109,12 @@ def color_coherence_vector(image, k=3):
         total_pixels = np.sum(region_mask)
         ccv.extend([total_pixels, total_pixels])
     
-    # Pad with zeros if needed
     ccv.extend([0] * (2 * k - len(ccv)))
     return np.array(ccv[:2*k])
 
 @st.cache_resource
-def create_feature_extractor():
-    """Create and cache the feature extractor model"""
+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)
@@ -129,23 +130,30 @@ def create_feature_extractor():
 
 def extract_features(image):
     """Extract all features from an image"""
-    # Convert image to uint8 for OpenCV operations
-    image_uint8 = (image * 255).astype(np.uint8)
+    # Traditional features
+    hist_features = color_histogram(image)
+    moment_features = color_moments(image)
+    dominant_features = dominant_color_descriptor(image)
+    ccv_features = color_coherence_vector(image)
     
-    # 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([
+    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()])
 
-def preprocess_image(image):
+def preprocess_image(image, scaler):
     """Preprocess the uploaded image"""
     # Convert to RGB if needed
     if image.mode != 'RGB':
@@ -156,25 +164,13 @@ def preprocess_image(image):
     img_array = cv2.resize(img_array, (256, 256))
     img_array = img_array.astype('float32') / 255.0
     
-    # Extract traditional features
-    traditional_features = extract_features(img_array)
+    # Extract all features
+    features = extract_features(img_array)
     
-    # Extract deep features
-    feature_extractor = create_feature_extractor()
-    deep_features = feature_extractor.predict(
-        np.expand_dims(img_array, axis=0),
-        verbose=0
-    )
-    
-    # Combine features
-    combined_features = np.concatenate([
-        traditional_features.reshape(1, -1),
-        deep_features.reshape(1, -1)
-    ], axis=1)
+    # Scale features using the provided scaler
+    scaled_features = scaler.transform(features.reshape(1, -1))
     
-    # Scale features
-    scaler = StandardScaler()
-    return scaler.fit_transform(combined_features)
+    return scaled_features
 
 def get_top_predictions(prediction, class_names, top_k=5):
     """Get top k predictions with their probabilities"""
@@ -188,6 +184,12 @@ def main():
     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
     if 'predictions' not in st.session_state:
         st.session_state.predictions = None
@@ -204,12 +206,7 @@ def main():
                 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
-                    
-                    processed_image = preprocess_image(image)
+                    processed_image = preprocess_image(image, scaler)
                     prediction = model.predict(processed_image, verbose=0)
                     
                     class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']