Update app.py

app.py

@@ -6,62 +6,122 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import joblib
-import json
-from sklearn.tree import export_text
+import os
 
 # Constants
 NUM_CLASSES = 4
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 CLASS_NAMES = ['Cracked_road', 'Flooded_muddy', 'Good_condition', 'Potholes']
 
-# Load the CNN model
-def load_cnn_model():
-    # Build model architecture (same as in train_cnn.py)
+# Build CNN model architecture (from train_cnn.py)
+def build_model(num_classes=4):
+    """Build EfficientNet-B0 based model for road surface classification"""
     model = efficientnet_b0(weights=None)
     
-    # Replace classifier
+    # Replace classifier with custom layers
     in_features = model.classifier[1].in_features
     model.classifier = nn.Sequential(
         nn.Linear(in_features, 256),
         nn.BatchNorm1d(256),
         nn.ReLU(),
         nn.Dropout(p=0.2),
-        nn.Linear(256, NUM_CLASSES)
+        nn.Linear(256, num_classes)
     )
     
-    # Load the saved weights
-    model.load_state_dict(torch.load("best_cnn_model.pth", map_location=DEVICE))
+    return model
+
+# Load the CNN model
+def load_cnn_model():
+    """Load the trained CNN model"""
+    model = build_model(NUM_CLASSES)
+    
+    # Try different possible paths for the model file
+    model_paths = [
+        "best_cnn_model.pth",
+        "models/best_cnn_model.pth", 
+        "assign2_Latest/models/best_cnn_model.pth"
+    ]
+    
+    model_loaded = False
+    for path in model_paths:
+        if os.path.exists(path):
+            try:
+                model.load_state_dict(torch.load(path, map_location=DEVICE))
+                model_loaded = True
+                print(f"CNN model loaded from {path}")
+                break
+            except Exception as e:
+                print(f"Failed to load from {path}: {e}")
+                continue
+    
+    if not model_loaded:
+        raise FileNotFoundError("Could not find or load CNN model file")
+    
     model.to(DEVICE)
     model.eval()
     return model
 
 # Load the Random Forest model
 def load_rf_model():
-    with open("rf_model.pkl", "rb") as f:
-        return joblib.load(f)
+    """Load the trained Random Forest model"""
+    rf_paths = [
+        "rf_model.pkl",
+        "models/rf_model.pkl",
+        "assign2_Latest/models/rf_model.pkl"
+    ]
+    
+    for path in rf_paths:
+        if os.path.exists(path):
+            try:
+                return joblib.load(path)
+            except Exception as e:
+                print(f"Failed to load RF model from {path}: {e}")
+                continue
+    
+    raise FileNotFoundError("Could not find or load Random Forest model file")
+
+# Extract features using CNN for Random Forest
+def extract_features_single(cnn_model, image_tensor):
+    """Extract features from a single image using CNN for Random Forest"""
+    cnn_model.eval()
+    with torch.no_grad():
+        # Get features from the layer before final classification
+        features = cnn_model.features(image_tensor)
+        features = cnn_model.avgpool(features)
+        features = torch.flatten(features, 1)
+        
+        # Pass through the first part of classifier (before final layer)
+        if hasattr(cnn_model.classifier, '__iter__'):
+            for i, layer in enumerate(cnn_model.classifier[:-1]):  # Exclude final classification layer
+                features = layer(features)
+        
+        return features.cpu().numpy()
 
 # Preprocess image for CNN
 def preprocess_image_cnn(image):
+    """Preprocess image for CNN prediction"""
     preprocess = transforms.Compose([
         transforms.Resize((224, 224)),
+        transforms.CenterCrop(224),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])
     return preprocess(image).unsqueeze(0)
 
 # Preprocess image for RF
-def preprocess_image_rf(image):
-    # Convert to numpy, resize, and flatten
-    image = image.resize((64, 64))
-    image_array = np.array(image)
-    # If grayscale, expand to RGB
-    if len(image_array.shape) == 2:
-        image_array = np.stack([image_array] * 3, axis=2)
-    # Flatten the image
-    return image_array.reshape(1, -1)
+def preprocess_image_rf(image, cnn_model):
+    """Preprocess image for Random Forest prediction using CNN features"""
+    # First preprocess for CNN
+    img_tensor = preprocess_image_cnn(image)
+    img_tensor = img_tensor.to(DEVICE)
+    
+    # Extract features using CNN
+    features = extract_features_single(cnn_model, img_tensor)
+    return features
 
 # Prediction function
 def predict(img, model_choice, cnn_model, rf_model):
+    """Make prediction using selected model"""
     if img is None:
         return "No image provided", "Please upload an image"
     
@@ -85,8 +145,8 @@ def predict(img, model_choice, cnn_model, rf_model):
             return CLASS_NAMES[predicted_class], confidence_text
         
         else:  # Random Forest
-            # Preprocess for RF
-            img_features = preprocess_image_rf(img)
+            # Preprocess for RF using CNN features
+            img_features = preprocess_image_rf(img, cnn_model)
             
             # Get prediction
             predicted_class = rf_model.predict(img_features)[0]
@@ -102,76 +162,55 @@ def predict(img, model_choice, cnn_model, rf_model):
     except Exception as e:
         return f"Error: {str(e)}", "An error occurred during prediction."
 
-# Convert Random Forest to a safer JSON format
-def convert_rf_to_safe_format():
-    """Convert Random Forest to a safer JSON format"""
-    
-    # Load the existing model
-    rf_model = joblib.load("models/rf_model.pkl")
-    
-    # Extract model parameters
-    model_params = {
-        'n_estimators': rf_model.n_estimators,
-        'max_depth': rf_model.max_depth,
-        'min_samples_split': rf_model.min_samples_split,
-        'random_state': rf_model.random_state,
-        'classes': rf_model.classes_.tolist(),
-        'n_features': rf_model.n_features_in_,
-        'feature_importances': rf_model.feature_importances_.tolist()
-    }
-    
-    # Save model parameters
-    with open("rf_model_params.json", "w") as f:
-        json.dump(model_params, f)
-    
-    # Save individual tree parameters (simplified approach)
-    trees_data = []
-    for i, tree in enumerate(rf_model.estimators_):
-        tree_data = {
-            'tree_id': i,
-            'feature': tree.tree_.feature.tolist(),
-            'threshold': tree.tree_.threshold.tolist(),
-            'children_left': tree.tree_.children_left.tolist(),
-            'children_right': tree.tree_.children_right.tolist(),
-            'value': tree.tree_.value.tolist()
-        }
-        trees_data.append(tree_data)
-    
-    with open("rf_trees_data.json", "w") as f:
-        json.dump(trees_data, f)
-    
-    print("Model converted to safe JSON format")
-
-# Load models
-cnn_model = load_cnn_model()
-rf_model = load_rf_model()
+# Initialize models
+print("Loading models...")
+try:
+    cnn_model = load_cnn_model()
+    rf_model = load_rf_model()
+    print("Models loaded successfully!")
+except Exception as e:
+    print(f"Error loading models: {e}")
+    # Create dummy models for testing
+    cnn_model = None
+    rf_model = None
 
 # Create Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("## Road Surface Classification")
+with gr.Blocks(title="Road Surface Classification") as demo:
+    gr.Markdown("# Road Surface Classification")
+    gr.Markdown("Upload an image of a road surface to classify its condition using either CNN or Random Forest models.")
     
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Road Image")
-            model_selector = gr.Radio(choices=["CNN", "Random Forest"], value="CNN", label="Select Model")
-            classify_btn = gr.Button("Classify")
+            model_selector = gr.Radio(
+                choices=["CNN", "Random Forest"], 
+                value="CNN", 
+                label="Select Model"
+            )
+            classify_btn = gr.Button("Classify", variant="primary")
+            
         with gr.Column():
-            label_output = gr.Textbox(label="Predicted Class")
-            confidence_output = gr.Textbox(label="Confidence Scores", lines=5)
+            label_output = gr.Textbox(label="Predicted Class", interactive=False)
+            confidence_output = gr.Textbox(
+                label="Confidence Scores", 
+                lines=5, 
+                interactive=False
+            )
 
+    # Event handler
     classify_btn.click(
         fn=lambda img, model_choice: predict(img, model_choice, cnn_model, rf_model),
         inputs=[image_input, model_selector],
         outputs=[label_output, confidence_output]
     )
     
-    gr.Markdown("### How to use")
-    gr.Markdown("1. Upload a road image")
-    gr.Markdown("2. Select either CNN or Random Forest model")
-    gr.Markdown("3. Click 'Classify' to get the prediction")
+    # Add information about classes
+    gr.Markdown("### Classes")
+    gr.Markdown("- **Cracked_road**: Roads with visible cracks")
+    gr.Markdown("- **Flooded_muddy**: Roads affected by flooding or mud")
+    gr.Markdown("- **Good_condition**: Roads in good condition")
+    gr.Markdown("- **Potholes**: Roads with potholes")
 
 # Launch the app
-demo.launch()
-
 if __name__ == "__main__":
-    convert_rf_to_safe_format()
+    demo.launch()