app.py

# app.py

import os
import math
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
from PIL import Image
import gradio as gr
from groq import Groq
from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas

# Load API Key
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
    raise ValueError("GROQ_API_KEY not found in environment. Please add it in HF Space Secrets.")
client = Groq(api_key=api_key)

# ------------------------------
# Model: Soil Classifier
# ------------------------------
class SoilClassifier(nn.Module):
    def __init__(self):
        super(SoilClassifier, self).__init__()
        self.base_model = models.resnet18(weights=None)
        num_features = self.base_model.fc.in_features
        self.base_model.fc = nn.Linear(num_features, 1)

    def forward(self, x):
        return self.base_model(x)

model = SoilClassifier()
model.base_model.load_state_dict(torch.load('soil_model.pth', map_location=torch.device('cpu')))
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# ------------------------------
# Shared report log
# ------------------------------
report_summary = []

# ------------------------------
# Soil Prediction Function
# ------------------------------
def predict_soil_type(image):
    image = image.convert("RGB")
    img_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img_tensor)
        raw_output = outputs.item()
        prediction = torch.sigmoid(outputs).item()
        result = f"Model Raw Output: {raw_output:.4f}"
        report_summary.append(f"Soil Image Prediction: {result}")
        return result


# ------------------------------
# AI Advisor
# ------------------------------
def ask_soil_region(query):
    prompt = f"""You are GeoMate, a world-class geotechnical expert.
Answer the following query using global soil knowledge, latest construction practices, and foundation design standards.
Query: {query}"""
    response = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# ------------------------------
# Classification System
# ------------------------------
def classify_soil(system, liquid_limit, plasticity_index, grain_size):
    try:
        result = ""
        if system == "USCS":
            if grain_size > 50:
                result = "Gravel"
            elif grain_size > 0.075:
                result = "Sand"
            else:
                result = "Clay" if plasticity_index > 7 else "Silt"
        elif system == "AASHTO":
            if liquid_limit < 40 and plasticity_index < 10:
                result = "A-1 or A-2 (Granular Soil)"
            elif plasticity_index > 10:
                result = "A-5 to A-7 (Silty/Clayey Soil)"
            else:
                result = "A-4 (Silt)"
        else:
            result = "Invalid system"
        report_summary.append(f"Soil Classification: {result}")
        return result
    except Exception as e:
        return f"Error: {e}"

# ------------------------------
# Engineering Calculations
# ------------------------------
def convert_pressure(val, unit):
    val = float(val)
    if unit == "psf":
        return val * 0.04788
    return val

def bearing_capacity_solver(q, Nq, S, B):
    try:
        q_converted = convert_pressure(q, S)
        result = q_converted * float(Nq) * float(B)
        report_summary.append(f"Bearing Capacity: {round(result, 2)} kN/m²")
        return f"{round(result, 2)} kN/m²"
    except Exception as e:
        return f"Error: {e}"

def slope_stability_solver(c, phi, gamma, height):
    try:
        phi = math.radians(float(phi))
        fs = (float(c) + float(gamma) * float(height) * math.tan(phi)) / (float(gamma) * float(height))
        report_summary.append(f"Slope Stability Factor of Safety: {round(fs, 3)}")
        return f"{round(fs, 3)} (Factor of Safety)"
    except Exception as e:
        return f"Error: {e}"

def consolidation_solver(delta_sigma, mv, H):
    try:
        settlement = float(mv) * float(delta_sigma) * float(H)
        report_summary.append(f"Settlement: {round(settlement, 3)} m")
        return f"{round(settlement, 3)} m"
    except Exception as e:
        return f"Error: {e}"

def seepage_solver(k, i, A):
    try:
        q = float(k) * float(i) * float(A)
        report_summary.append(f"Seepage Discharge: {round(q, 4)} m³/s")
        return f"{round(q, 4)} m³/s"
    except Exception as e:
        return f"Error: {e}"

def compaction_solver(W, V):
    try:
        dry_density = float(W) / float(V)
        report_summary.append(f"Dry Density: {round(dry_density, 2)} kN/m³")
        return f"{round(dry_density, 2)} kN/m³"
    except Exception as e:
        return f"Error: {e}"

# ------------------------------
# PDF Report Generator
# ------------------------------
def export_full_report():
    try:
        file_path = "/tmp/GeoMate_Report.pdf"
        c = canvas.Canvas(file_path, pagesize=A4)
        width, height = A4
        c.setFont("Helvetica", 12)
        y = height - 50
        for line in report_summary:
            c.drawString(40, y, line)
            y -= 20
            if y < 50:
                c.showPage()
                c.setFont("Helvetica", 12)
                y = height - 50
        c.save()
        return file_path
    except Exception as e:
        return f"Error: {e}"

# ------------------------------
# Gradio Interface
# ------------------------------
with gr.Blocks(title="GeoMate 🌍 - Soil Engineering Toolkit") as demo:
    gr.Markdown("""
    <div style='text-align:center; font-size:26px; font-weight:bold; color:#ff6600;'>🌍 GeoMate - Soil Engineering Toolkit</div>
    <p style='color:#333;'>Perform all major soil-related geotechnical calculations in one place!</p>
    """)

    with gr.Tab("📷 Soil Recognizer"):
        img_input = gr.Image(type="pil", label="Upload Soil Image")
        img_output = gr.Textbox(label="Prediction Output")
        img_input.change(fn=predict_soil_type, inputs=img_input, outputs=img_output)

    with gr.Tab("🤖 Ask GeoMate"):
        gr.Interface(fn=ask_soil_region,
                     inputs=gr.Textbox(placeholder="e.g., What foundation is suitable in Karachi?", lines=2),
                     outputs="text").render()

    with gr.Tab("🧪 Soil Classification"):
        system = gr.Dropdown(["USCS", "AASHTO"], label="Classification System", value="USCS")
        ll = gr.Number(label="Liquid Limit (%)")
        pi = gr.Number(label="Plasticity Index (%)")
        gs = gr.Number(label="Grain Size (mm)")
        classify_btn = gr.Button("Classify Soil")
        classification = gr.Textbox(label="Soil Type")
        classify_btn.click(classify_soil, [system, ll, pi, gs], classification)

    with gr.Tab("🏗️ Bearing Capacity"):
        q = gr.Number(label="Overburden Pressure")
        q_unit = gr.Dropdown(["kN/m²", "psf", "kPa"], label="Unit", value="kN/m²")
        nq = gr.Number(label="Nq (Bearing Capacity Factor)")
        B = gr.Number(label="Width of Foundation (m)")
        bc_result = gr.Textbox(label="Ultimate Bearing Capacity")
        gr.Button("Calculate").click(bearing_capacity_solver, [q, nq, q_unit, B], bc_result)

    with gr.Tab("⛰️ Slope Stability"):
        c = gr.Number(label="Cohesion (kN/m²)")
        phi = gr.Number(label="Friction Angle (°)")
        gamma = gr.Number(label="Unit Weight (kN/m³)")
        h = gr.Number(label="Height of Slope (m)")
        fs_result = gr.Textbox(label="Factor of Safety")
        gr.Button("Check Stability").click(slope_stability_solver, [c, phi, gamma, h], fs_result)

    with gr.Tab("📉 Consolidation Settlement"):
        ds = gr.Number(label="Change in Stress (kN/m²)")
        mv = gr.Number(label="Volume Compressibility (m²/kN)")
        H = gr.Number(label="Soil Layer Thickness (m)")
        s_result = gr.Textbox(label="Settlement")
        gr.Button("Compute Settlement").click(consolidation_solver, [ds, mv, H], s_result)

    with gr.Tab("💧 Seepage"):
        k = gr.Number(label="Permeability (m/s)")
        i = gr.Number(label="Hydraulic Gradient")
        A = gr.Number(label="Flow Area (m²)")
        seep_result = gr.Textbox(label="Seepage Discharge")
        gr.Button("Calculate Seepage").click(seepage_solver, [k, i, A], seep_result)

    with gr.Tab("🔩 Compaction Test"):
        W = gr.Number(label="Dry Weight of Soil (kN)")
        V = gr.Number(label="Volume of Mold (m³)")
        comp_result = gr.Textbox(label="Dry Density")
        gr.Button("Calculate Dry Density").click(compaction_solver, [W, V], comp_result)

    with gr.Tab("📄 Generate Report"):
        pdf_output = gr.File(label="Download PDF")
        gr.Button("Export Full Report").click(fn=export_full_report, inputs=[], outputs=pdf_output)

# Launch App
demo.launch()