Update app.py

app.py

@@ -184,46 +184,46 @@ def create_vit_feature_extractor(input_shape=(256, 256, 3), num_classes=None):
     
 def preprocess_image(image):
     """Preprocess the uploaded image"""
-    # # Convert to RGB if needed
-    # if image.mode != 'RGB':
-    #     image = image.convert('RGB')
+    # 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)
+    # 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)
+    
+    # Create and process ViT features
     vit_extractor = create_vit_feature_extractor()
-
-    # Trích xuất đặc trưng ViT từ các hình ảnh
-    image_vit = vit_extractor.predict(img_array)  # Dự đoán cho tập train
-    image_combined = np.concatenate([image_features, image_vit], axis=1)
+    
+    # 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
+    
+    # Combine features
+    image_combined = np.concatenate([image_features.reshape(1, -1), image_vit], axis=1)
+    
+    # Scale features
     scaler = StandardScaler()
     image_scaled = scaler.fit_transform(image_combined)
-    return image_scaled
+    
+    return image_scaled.squeeze()  # Remove any unnecessary dimensions
 
 def get_top_predictions(prediction, class_names, top_k=5):
     """Get top k predictions with their probabilities"""
@@ -267,8 +267,11 @@ def main():
                     # Preprocess image
                     processed_image = preprocess_image(image)
                     
+                    # Ensure correct shape for prediction
+                    processed_image = np.expand_dims(processed_image, axis=0)
+                    
                     # Make prediction
-                    prediction = model.predict(np.expand_dims(processed_image, axis=0))
+                    prediction = model.predict(processed_image)
                     class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
                     
                     # Get top 5 predictions