Update app.py

app.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from torchvision import transforms
-from torchvision.models import efficientnet_b0
+from torchvision.models import efficientnet_b0, mobilenet_v2, resnet18
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -13,12 +13,36 @@ NUM_CLASSES = 4
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 CLASS_NAMES = ['Cracked_road', 'Flooded_muddy', 'Good_condition', 'Potholes']
 
-# Build CNN model architecture (from train_cnn.py)
-def build_model(num_classes=4):
-    """Build EfficientNet-B0 based model for road surface classification"""
+# Model paths with fallbacks for different deployment environments
+MODEL_PATHS = {
+    "efficientnet": [
+        "best_efficientnet_model.pth",
+        "models/best_efficientnet_model.pth",
+        "assign2_Latest/models/best_efficientnet_model.pth"
+    ],
+    "mobilenet": [
+        "best_mobilenetv2_model.pth",
+        "models/best_mobilenetv2_model.pth",
+        "assign2_Latest/models/best_mobilenetv2_model.pth"
+    ],
+    "resnet": [
+        "best_resnet18_model.pth",
+        "models/best_resnet18_model.pth",
+        "assign2_Latest/models/best_resnet18_model.pth"
+    ],
+    "rf": [
+        "rf_model.pkl",
+        "models/rf_model.pkl",
+        "assign2_Latest/models/rf_model.pkl"
+    ]
+}
+
+# Build EfficientNet model
+def build_efficientnet(num_classes=NUM_CLASSES):
+    """Build EfficientNet-B0 based model"""
     model = efficientnet_b0(weights=None)
     
-    # Replace classifier with custom layers
+    # Replace classifier
     in_features = model.classifier[1].in_features
     model.classifier = nn.Sequential(
         nn.Linear(in_features, 256),
@@ -30,71 +54,118 @@ def build_model(num_classes=4):
     
     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"
-    ]
+# Build MobileNetV2 model
+def build_mobilenet(num_classes=NUM_CLASSES):
+    """Build MobileNetV2 model"""
+    model = mobilenet_v2(weights=None)
     
+    # Replace classifier
+    last_channel = model.last_channel
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.2),
+        nn.Linear(last_channel, 512),
+        nn.BatchNorm1d(512),
+        nn.LeakyReLU(0.2),
+        nn.Dropout(p=0.2),
+        nn.Linear(512, 256),
+        nn.BatchNorm1d(256),
+        nn.LeakyReLU(0.2),
+        nn.Dropout(p=0.1),
+        nn.Linear(256, num_classes)
+    )
+    
+    return model
+
+# Build ResNet18 model
+def build_resnet(num_classes=NUM_CLASSES):
+    """Build ResNet18 model"""
+    model = resnet18(weights=None)
+    
+    # Replace classifier
+    in_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Linear(in_features, 256),
+        nn.ReLU(),
+        nn.Dropout(0.3),
+        nn.Linear(256, num_classes)
+    )
+    
+    return model
+
+# Load model with fallback paths
+def load_model(model_type):
+    """Load trained model with fallback paths"""
+    if model_type == "efficientnet":
+        model = build_efficientnet()
+    elif model_type == "mobilenet":
+        model = build_mobilenet()
+    elif model_type == "resnet":
+        model = build_resnet()
+    elif model_type == "rf":
+        # For Random Forest, handled differently below
+        return load_rf_model()
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+    
+    # Try loading from various possible paths
+    paths = MODEL_PATHS[model_type]
     model_loaded = False
-    for path in model_paths:
+    
+    for path in 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}")
+                print(f"{model_type} model loaded from {path}")
                 break
             except Exception as e:
-                print(f"Failed to load from {path}: {e}")
+                print(f"Failed to load {model_type} from {path}: {e}")
                 continue
     
     if not model_loaded:
-        raise FileNotFoundError("Could not find or load CNN model file")
+        raise FileNotFoundError(f"Could not find or load {model_type} model file")
     
     model.to(DEVICE)
     model.eval()
     return model
 
-# Load the Random Forest model
+# Load Random Forest model
 def load_rf_model():
-    """Load the trained Random Forest model"""
-    rf_paths = [
-        "rf_model.pkl",
-        "models/rf_model.pkl",
-        "assign2_Latest/models/rf_model.pkl"
-    ]
+    """Load Random Forest model"""
+    paths = MODEL_PATHS["rf"]
     
-    for path in rf_paths:
+    for path in paths:
         if os.path.exists(path):
             try:
-                return joblib.load(path)
+                model = joblib.load(path)
+                print(f"Random Forest model loaded from {path}")
+                return model
             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()
+# Extract features for Random Forest
+def extract_features_single(model, image_tensor):
+    """Extract features from a single image for Random Forest"""
+    model.eval()
+    
+    # Create feature extractor (using EfficientNet)
+    feature_extractor = nn.Sequential(
+        model.features,
+        model.avgpool,
+        nn.Flatten()
+    ).to(DEVICE)
+    
     with torch.no_grad():
-        # Extract raw 1280-dim features (matching how the RF model was trained)
-        features = cnn_model.features(image_tensor)
-        features = cnn_model.avgpool(features)
-        features = torch.flatten(features, 1)
-        
-        return features.cpu().numpy()
+        features = feature_extractor(image_tensor).cpu().numpy()
+    
+    return features
 
-# Preprocess image for CNN
-def preprocess_image_cnn(image):
-    """Preprocess image for CNN prediction"""
+# Preprocess image for neural networks
+def preprocess_image(image):
+    """Preprocess image for model prediction"""
     preprocess = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.CenterCrop(224),
@@ -103,32 +174,26 @@ def preprocess_image_cnn(image):
     ])
     return preprocess(image).unsqueeze(0)
 
-# Preprocess image for RF
-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):
+def predict(img, model_choice, models):
     """Make prediction using selected model"""
     if img is None:
         return "No image provided", "Please upload an image"
     
     try:
-        if model_choice == "CNN":
-            # Preprocess for CNN
-            img_tensor = preprocess_image_cnn(img)
-            img_tensor = img_tensor.to(DEVICE)
+        # Convert to RGB if needed
+        if img.mode != "RGB":
+            img = img.convert("RGB")
             
-            # Get prediction
+        # Preprocess image
+        img_tensor = preprocess_image(img)
+        img_tensor = img_tensor.to(DEVICE)
+        
+        if model_choice == "EfficientNet":
+            model = models["efficientnet"]
+            # Make prediction
             with torch.no_grad():
-                outputs = cnn_model(img_tensor)
+                outputs = model(img_tensor)
                 probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
                 predicted_class = torch.argmax(probabilities).item()
             
@@ -138,14 +203,45 @@ def predict(img, model_choice, cnn_model, rf_model):
             confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
             
             return CLASS_NAMES[predicted_class], confidence_text
-        
-        else:  # Random Forest
-            # Preprocess for RF using CNN features
-            img_features = preprocess_image_rf(img, cnn_model)
             
-            # Get prediction
-            predicted_class = rf_model.predict(img_features)[0]
-            probabilities = rf_model.predict_proba(img_features)[0]
+        elif model_choice == "MobileNetV2":
+            model = models["mobilenet"]
+            # Make prediction
+            with torch.no_grad():
+                outputs = model(img_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
+                predicted_class = torch.argmax(probabilities).item()
+            
+            # Format confidence scores
+            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i].item()*100:.2f}%" 
+                               for i in range(len(CLASS_NAMES))}
+            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
+            
+            return CLASS_NAMES[predicted_class], confidence_text
+            
+        elif model_choice == "ResNet18":
+            model = models["resnet"]
+            # Make prediction
+            with torch.no_grad():
+                outputs = model(img_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
+                predicted_class = torch.argmax(probabilities).item()
+            
+            # Format confidence scores
+            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i].item()*100:.2f}%" 
+                               for i in range(len(CLASS_NAMES))}
+            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
+            
+            return CLASS_NAMES[predicted_class], confidence_text
+            
+        elif model_choice == "Random Forest":
+            # Extract features using EfficientNet
+            features = extract_features_single(models["efficientnet"], img_tensor)
+            
+            # Make prediction with Random Forest
+            rf_model = models["rf"]
+            predicted_class = rf_model.predict(features)[0]
+            probabilities = rf_model.predict_proba(features)[0]
             
             # Format confidence scores
             confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i]*100:.2f}%" 
@@ -153,33 +249,36 @@ def predict(img, model_choice, cnn_model, rf_model):
             confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
             
             return CLASS_NAMES[predicted_class], confidence_text
+        else:
+            return "Invalid model selection", "Please select a valid model"
             
     except Exception as e:
         return f"Error: {str(e)}", "An error occurred during prediction."
 
 # Initialize models
 print("Loading models...")
+models = {}
 try:
-    cnn_model = load_cnn_model()
-    rf_model = load_rf_model()
-    print("Models loaded successfully!")
+    models["efficientnet"] = load_model("efficientnet")
+    models["mobilenet"] = load_model("mobilenet")
+    models["resnet"] = load_model("resnet")
+    models["rf"] = load_model("rf")
+    print("All models loaded successfully!")
 except Exception as e:
     print(f"Error loading models: {e}")
-    # Create dummy models for testing
-    cnn_model = None
-    rf_model = None
+    # In case of error, models will be empty or partial
 
 # Create Gradio interface
 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.")
+    gr.Markdown("Upload an image of a road surface to classify its condition using various 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", 
+                choices=["EfficientNet", "MobileNetV2", "ResNet18", "Random Forest"], 
+                value="EfficientNet", 
                 label="Select Model"
             )
             classify_btn = gr.Button("Classify", variant="primary")
@@ -194,17 +293,26 @@ with gr.Blocks(title="Road Surface Classification") as demo:
 
     # Event handler
     classify_btn.click(
-        fn=lambda img, model_choice: predict(img, model_choice, cnn_model, rf_model),
+        fn=lambda img, model_choice: predict(img, model_choice, models),
         inputs=[image_input, model_selector],
         outputs=[label_output, confidence_output]
     )
     
-    # 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")
+    # Add information about classes and models
+    with gr.Row():
+        with gr.Column():
+            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")
+        
+        with gr.Column():
+            gr.Markdown("### Models")
+            gr.Markdown("- **EfficientNet**: Efficient deep learning model")
+            gr.Markdown("- **MobileNetV2**: Lightweight mobile-friendly model")
+            gr.Markdown("- **ResNet18**: Residual network architecture")
+            gr.Markdown("- **Random Forest**: Traditional ML using CNN features")
 
 # Launch the app
 if __name__ == "__main__":