Update app.py

app.py

@@ -1,133 +1,137 @@
-import torch
-import torch.nn as nn
-from torchvision import transforms
-from torchvision.models import efficientnet_b0
-import gradio as gr
-import numpy as np
-import pickle
-import os
-from PIL import Image
-
-# Constants
-NUM_CLASSES = 4
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-CLASS_NAMES = ["Good_condition", "Potholes", "Cracked_road", "Flooded_muddy"]
-
-# Load the CNN model
-def load_cnn_model():
-    # Build model architecture (same as in train_cnn.py)
-    model = efficientnet_b0(weights=None)
-    
-    # Replace classifier
-    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)
-    )
-    
-    # Load the saved weights
-    model.load_state_dict(torch.load("best_cnn_model.pth", map_location=DEVICE))
-    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 pickle.load(f)
-
-# Preprocess image for CNN
-def preprocess_image_cnn(image):
-    preprocess = transforms.Compose([
-        transforms.Resize((224, 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)
-
-# Prediction function
-def predict(img, model_choice, cnn_model, rf_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)
-            
-            # Get prediction
-            with torch.no_grad():
-                outputs = cnn_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
-        
-        else:  # Random Forest
-            # Preprocess for RF
-            img_features = preprocess_image_rf(img)
-            
-            # Get prediction
-            predicted_class = rf_model.predict(img_features)[0]
-            probabilities = rf_model.predict_proba(img_features)[0]
-            
-            # Format confidence scores
-            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i]*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
-            
-    except Exception as e:
-        return f"Error: {str(e)}", "An error occurred during prediction."
-
-# Load models
-cnn_model = load_cnn_model()
-rf_model = load_rf_model()
-
-# Create Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("## Road Surface Classification")
-    
-    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")
-        with gr.Column():
-            label_output = gr.Textbox(label="Predicted Class")
-            confidence_output = gr.Textbox(label="Confidence Scores", lines=5)
-
-    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")
-
-# Launch the app
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import efficientnet_b0
+import gradio as gr
+import numpy as np
+import pickle
+import os
+import joblib
+from PIL import Image
+
+# Constants
+NUM_CLASSES = 4
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CLASS_NAMES = ["Good_condition", "Potholes", "Cracked_road", "Flooded_muddy"]
+
+# Load the CNN model
+def load_cnn_model():
+    # Build model architecture (same as in train_cnn.py)
+    model = efficientnet_b0(weights=None)
+    
+    # Replace classifier
+    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)
+    )
+    
+    # Load the saved weights
+    model.load_state_dict(torch.load("best_cnn_model.pth", map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval()
+    return model
+
+# Load the Random Forest model
+def load_rf_model():
+    try:
+        return joblib.load("rf_model.pkl")
+    except Exception as e:
+        print(f"Error loading RF model: {e}")
+        return None
+
+# Preprocess image for CNN
+def preprocess_image_cnn(image):
+    preprocess = transforms.Compose([
+        transforms.Resize((224, 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)
+
+# Prediction function
+def predict(img, model_choice, cnn_model, rf_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)
+            
+            # Get prediction
+            with torch.no_grad():
+                outputs = cnn_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
+        
+        else:  # Random Forest
+            # Preprocess for RF
+            img_features = preprocess_image_rf(img)
+            
+            # Get prediction
+            predicted_class = rf_model.predict(img_features)[0]
+            probabilities = rf_model.predict_proba(img_features)[0]
+            
+            # Format confidence scores
+            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i]*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
+            
+    except Exception as e:
+        return f"Error: {str(e)}", "An error occurred during prediction."
+
+# Load models
+cnn_model = load_cnn_model()
+rf_model = load_rf_model()
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Road Surface Classification")
+    
+    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")
+        with gr.Column():
+            label_output = gr.Textbox(label="Predicted Class")
+            confidence_output = gr.Textbox(label="Confidence Scores", lines=5)
+
+    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")
+
+# Launch the app
 demo.launch()