Update app.py

app.py

@@ -453,98 +453,73 @@ if __name__ == "__main__":
 #  - GSD Curve Page (CSV upload + plotting + parameter calc)
 # -------------------------------------------------------
 
-import numpy as np
-import pandas as pd
-from PIL import Image
+def soil_recognizer_page():
+    st.header("🖼️ Soil Recognizer")
+    idx = st.session_state["active_site_idx"]
 
-# -------------------------------------------------------
-# Soil Recognizer
-# -------------------------------------------------------
-import torch
-import torch.nn as nn
-import torchvision.models as models
-import torchvision.transforms as transforms
-import streamlit as st
+    st.write("Upload a soil sample photo. If a trained model is available, it will infer the soil class.")
 
-# ----------------------------
-# Define SoilNet Architecture
-# ----------------------------
-class SoilNet(nn.Module):
-    def __init__(self, num_classes=5):
-        super(SoilNet, self).__init__()
-        # Example: ResNet18 backbone (adjust if you trained differently)
-        self.model = models.resnet18(weights=None)   # no pretrained weights
-        in_features = self.model.fc.in_features
-        self.model.fc = nn.Linear(in_features, num_classes)
-
-    def forward(self, x):
-        return self.model(x)
-
-# ----------------------------
-# Load Soil Recognition Model
-# ----------------------------
-@st.cache_resource
-def load_soil_model():
-    try:
-        model = SoilNet(num_classes=5)  # adjust num_classes if needed
-        state_dict = torch.load("soil_best_model.pth", map_location="cpu")
-        model.load_state_dict(state_dict)
-        model.eval()
-        return model
-    except Exception as e:
-        st.error(f"⚠️ Could not load soil model: {e}")
-        return None
+    uploaded = st.file_uploader(
+        "Upload sample photo",
+        type=["png", "jpg", "jpeg"],
+        key=mk("sr_upload", idx)
+    )
+
+    if uploaded:
+        img = Image.open(uploaded).convert("RGB")
+        st.image(img, use_column_width=True)
 
-soil_model = load_soil_model()
+        if torch and os.path.exists("soil_best_model.pth"):
+            st.info("✅ Model found — running inference (CPU).")
 
-# Define soil classes (adjust to your training labels)
-SOIL_CLASSES = ["Sand", "Silt", "Clay", "Gravel", "Peat"]
+            try:
+                # --- Load model safely ---
+                model = torch.load("soil_best_model.pth", map_location="cpu")
+                if hasattr(model, "eval"):
+                    model.eval()
 
-# Image preprocessing pipeline
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),   # match training input size
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],  # ImageNet normalization
-                         [0.229, 0.224, 0.225])
-])
+                # --- Preprocess ---
+                transform = T.Compose([
+                    T.Resize((224, 224)),
+                    T.ToTensor(),
+                    T.Normalize([0.485, 0.456, 0.406],
+                                [0.229, 0.224, 0.225])
+                ])
+                inp = transform(img).unsqueeze(0)
 
-# ----------------------------
-# Soil Recognizer UI
-# ----------------------------
-def soil_recognizer_ui():
-    st.header("🖼️ Soil Recognizer (Image / OCR)")
-    site = get_active_site()
+                with st.spinner("Running model..."):
+                    logits = model(inp)
+                    probs = torch.softmax(logits, dim=-1).detach().cpu().numpy()[0]
 
-    col1, col2 = st.columns(2)
+                labels = ["Sand", "Silt", "Clay", "Gravel", "Peat"]
+                best = labels[int(np.argmax(probs))]
+                conf = float(np.max(probs))
 
-    with col1:
-        uploaded = st.file_uploader("Upload soil image", type=["jpg", "jpeg", "png"])
-        if uploaded is not None:
-            img = Image.open(uploaded).convert("RGB")
-            st.image(img, caption="Uploaded soil image", use_column_width=True)
+                st.success(f"Predicted: **{best}** (confidence {conf:.2%})")
 
-            if soil_model:
-                try:
-                    # Preprocess
-                    img_t = transform(img).unsqueeze(0)  # add batch dimension
-                    with torch.no_grad():
-                        outputs = soil_model(img_t)
-                        probs = torch.softmax(outputs, dim=1)
-                        conf, pred = torch.max(probs, 1)
-                        predicted_class = SOIL_CLASSES[pred.item()]
-                        confidence = conf.item()
-
-                    # Save to site
-                    site["Soil Profile"] = predicted_class
-                    site.setdefault("classifier_inputs", {})["image_confidence"] = confidence
-                    save_sites(SITES)
-
-                    st.success(f"✅ Predicted: **{predicted_class}** ({confidence:.2%} confidence)")
+                if st.button("Save to site", key=mk("sr_save_btn", idx)):
+                    st.session_state["sites"][idx]["Soil Profile"] = best
+                    st.success("✅ Saved soil profile to site.")
 
-                except Exception as e:
-                    st.error(f"❌ Inference error: {e}")
+            except Exception as e:
+                st.error(f"❌ Model inference failed: {e}")
+
+        else:
+            # --- Heuristic Fallback ---
+            st.warning("⚠️ No trained model file found — running heuristic fallback.")
+            arr = np.array(img.resize((50, 50))).mean(axis=(0, 1))
+            r, g, b = arr
+            if r > 120 and g > 110:
+                pred = "Sand"
             else:
-                st.warning("⚠️ Soil model not loaded. Please check `soil_best_model.pth`.")
+                pred = "Silt"
+
+            st.success(f"Fallback prediction: **{pred}**")
+
+            if st.button("Save fallback to site", key=mk("sr_save_fallback", idx)):
+                st.session_state["sites"][idx]["Soil Profile"] = pred
+                st.success("✅ Saved fallback result to site.")
+                    
 # -------------------------------------------------------
 # Soil Classifier
 # -------------------------------------------------------