Update app.py

app.py

@@ -7,12 +7,18 @@ import gradio as gr
 import joblib
 import os
 
-# ✅ Load label map
+# ✅ Label maps
 label_map = {
     'G': 0, 'DR': 1, 'WD': 2, 'ND': 3, 'OTHER': 4
 }
 reverse_label_map = {v: k for k, v in label_map.items()}
-severity_map = {0: 'Low', 1: 'Medium', 2: 'High'}  # Adjust if needed
+
+# ⚠️ Adjust this based on cluster-to-severity relation (see notes below)
+severity_map = {
+    0: 'Low',     # e.g., yellow cluster
+    1: 'Medium',  # e.g., pink cluster
+    2: 'High'     # e.g., blue cluster
+}
 
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -37,13 +43,13 @@ def predict(img):
     try:
         img_tensor = transform(img).unsqueeze(0).to(device)
 
-        # Damage classification
+        # 🔍 Predict damage class
         with torch.no_grad():
             output = model(img_tensor)
             _, predicted = torch.max(output, 1)
             damage_class = reverse_label_map[predicted.item()]
 
-        # Feature extraction
+        # 🔍 Feature extraction for severity
         x = model.conv1(img_tensor)
         x = model.bn1(x)
         x = model.relu(x)
@@ -55,9 +61,11 @@ def predict(img):
         x = model.avgpool(x)
         x = torch.flatten(x, 1).detach().cpu().numpy()
 
-        # Severity clustering
+        # 🔍 PCA + Clustering
         x_pca = pca.transform(x)
         severity_cluster = kmeans.predict(x_pca)[0]
+        print(f"Predicted Severity Cluster: {severity_cluster}")  # Debug log
+
         severity_label = severity_map.get(severity_cluster, f"Cluster {severity_cluster}")
 
         return f"Damage: {damage_class}", f"Severity: {severity_label}"
@@ -65,7 +73,7 @@ def predict(img):
     except Exception as e:
         return f"Error: {str(e)}", ""
 
-# ✅ Gradio UI
+# ✅ Gradio Interface
 gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),