Update app.py

app.py

@@ -218,42 +218,44 @@ def site_details_panel():
 # =============================
 # FEATURE MODULES
 # =============================
-
-# 1. Soil Recognizer (CNN Image Model)
-# 1. Soil Recognizer (ResNet18 Model)
-
+# ----------------------------
+# Soil Recognizer Page (Integrated 6-Class ResNet18)
+# ----------------------------
 import torch
 import torch.nn as nn
 import torchvision.models as models
 import torchvision.transforms as T
 from PIL import Image
-import numpy as np
 import streamlit as st
 
 # ----------------------------
-# Load Soil Recognition Model
+# Load Soil Model (6 Classes)
 # ----------------------------
 @st.cache_resource
-def load_soil_model():
+def load_soil_model(path="/content/drive/MyDrive/soil_best_model.pth"):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     try:
-        # Load ResNet18 (pretrained base, adjust if you used different arch)
-        model = models.resnet18(weights=None)
+        model = models.resnet18(pretrained=False)
         num_ftrs = model.fc.in_features
-        model.fc = nn.Linear(num_ftrs, 5)   # 5 soil classes
-        state_dict = torch.load("soil_best_model.pth", map_location="cpu")
+        model.fc = nn.Linear(num_ftrs, 6)  # 6 soil classes
+
+        # Load checkpoint
+        state_dict = torch.load(path, map_location=device)
         model.load_state_dict(state_dict)
+        model = model.to(device)
         model.eval()
-        return model
+        return model, device
     except Exception as e:
         st.error(f"⚠️ Could not load soil model: {e}")
-        return None
+        return None, device
 
-soil_model = load_soil_model()
+soil_model, device = load_soil_model()
 
-# Soil class labels
-SOIL_CLASSES = ["Sand", "Silt", "Clay", "Gravel", "Peat"]
+# ----------------------------
+# Soil Classes & Transform
+# ----------------------------
+SOIL_CLASSES = ["Clay", "Gravel", "Loam", "Peat", "Sand", "Silt"]
 
-# Image preprocessing
 transform = T.Compose([
     T.Resize((224, 224)),
     T.ToTensor(),
@@ -261,45 +263,54 @@ transform = T.Compose([
                 [0.229, 0.224, 0.225])
 ])
 
+# ----------------------------
+# Prediction Function
+# ----------------------------
+def predict_soil(img: Image.Image):
+    if soil_model is None:
+        return "Model not loaded", {}
+
+    img = img.convert("RGB")
+    inp = transform(img).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        logits = soil_model(inp)
+        probs = torch.softmax(logits[0], dim=0)
+
+    top_idx = torch.argmax(probs).item()
+    predicted_class = SOIL_CLASSES[top_idx]
+
+    result = {SOIL_CLASSES[i]: float(probs[i]) for i in range(len(SOIL_CLASSES))}
+    return predicted_class, result
+
 # ----------------------------
 # Soil Recognizer Page
 # ----------------------------
 def soil_recognizer_page():
     st.header("🖼️ Soil Recognizer (ResNet18)")
 
-    site = get_active_site()
+    site = get_active_site()  # your existing site getter
     if site is None:
         st.warning("⚠️ No active site selected. Please add or select a site from the sidebar.")
         return
 
     uploaded = st.file_uploader("Upload soil image", type=["jpg", "jpeg", "png"])
     if uploaded is not None:
-        img = Image.open(uploaded).convert("RGB")
+        img = Image.open(uploaded)
         st.image(img, caption="Uploaded soil image", use_column_width=True)
 
-        if soil_model:
-            try:
-                inp = transform(img).unsqueeze(0)  # add batch dim
-                with torch.no_grad():
-                    logits = soil_model(inp)
-                    probs = torch.softmax(logits, dim=1)[0]
-                    conf, pred = torch.max(probs, 0)
-                    predicted_class = SOIL_CLASSES[pred.item()]
-                    confidence = conf.item()
-
-                st.success(f"✅ Predicted: **{predicted_class}** ({confidence:.2%} confidence)")
-
-                if st.button("Save to site"):
-                    site["Soil Profile"] = predicted_class
-                    site["Soil Recognizer Confidence"] = confidence
-                    save_active_site(site)
-                    st.success("Saved prediction to active site memory.")
+        predicted_class, confidence_scores = predict_soil(img)
+        st.success(f"✅ Predicted: **{predicted_class}**")
 
-            except Exception as e:
-                st.error(f"❌ Inference error: {e}")
-        else:
-            st.warning("⚠️ Soil model not loaded. Please check `soil_best_model.pth`.")
+        st.subheader("Confidence Scores")
+        for cls, score in confidence_scores.items():
+            st.write(f"{cls}: {score:.2%}")
 
+        if st.button("Save to site"):
+            site["Soil Profile"] = predicted_class
+            site["Soil Recognizer Confidence"] = confidence_scores[predicted_class]
+            save_active_site(site)
+            st.success("Saved prediction to active site memory.")
 
 
 # ----------------------------