Update app.py

app.py

@@ -74,6 +74,101 @@ if not GROQ_AVAILABLE:
 # 3) Global constants & helper functions
 MAX_SITES = 4
 
+
+# ----------------------------
+# 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 streamlit as st
+
+# ----------------------------
+# Load Soil Model (6 Classes)
+# ----------------------------
+@st.cache_resource
+def load_soil_model(path="soil_best_model.pth"):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    try:
+        model = models.resnet18(pretrained=False)
+        num_ftrs = model.fc.in_features
+        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, device
+    except Exception as e:
+        st.error(f"⚠️ Could not load soil model: {e}")
+        return None, device
+
+soil_model, device = load_soil_model()
+
+# ----------------------------
+# Soil Classes & Transform
+# ----------------------------
+SOIL_CLASSES = ["Clay", "Gravel", "Loam", "Peat", "Sand", "Silt"]
+
+transform = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize([0.485, 0.456, 0.406],
+                [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()  # 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)
+        st.image(img, caption="Uploaded soil image", use_column_width=True)
+
+        predicted_class, confidence_scores = predict_soil(img)
+        st.success(f"✅ Predicted: **{predicted_class}**")
+
+        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.")
+
 # Pre-defined dropdown text mappings (as you requested) — exact text with mapping numbers for logic backend
 DILATANCY_OPTIONS = [
     "1. Quick to slow",
@@ -500,6 +595,7 @@ def build_full_geotech_pdf(site: Dict[str, Any], filename: str, include_map_imag
 if "sites" not in st.session_state:
     # initialize with a default site
     st.session_state["sites"] = [{
+        "Soil Profile": None,
         "Site Name":"Home",
         "Project Name":"Demo Project",
         "Site ID": 0,
@@ -1112,6 +1208,8 @@ elif page == "RAG":
     rag_page()
 elif page == "Reports":
     reports_page()
+elif page == "Soil recognizer":
+    soil_recognizer_page()
 else:
     landing_page()