app.py

import streamlit as st
from PIL import Image

from utils.model_loader import load_all_models
from utils.predictor import predict_tree, predict_species, predict_stage
from utils.gradcam import (
    generate_gradcam_pytorch_resnet,
    generate_gradcam_pytorch_mobilenet,
    generate_gradcam_keras
)

st.set_page_config(
    page_title="TreeVision AI",
    page_icon="🌳",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');

    html, body, [class*="css"] {
        font-family: 'Inter', sans-serif;
    }

    .hero-section {
        background: linear-gradient(135deg, #1a472a 0%, #2d6a4f 50%, #40916c 100%);
        border-radius: 16px;
        padding: 2.5rem 3rem;
        margin-bottom: 2rem;
        color: white;
    }

    .hero-badge {
        display: inline-block;
        background: rgba(255,255,255,0.15);
        border: 1px solid rgba(255,255,255,0.25);
        border-radius: 20px;
        padding: 0.3rem 0.9rem;
        font-size: 0.75rem;
        font-weight: 600;
        letter-spacing: 1.5px;
        text-transform: uppercase;
        color: #b7e4c7;
        margin-bottom: 1rem;
    }

    .hero-title {
        font-size: 2.6rem;
        font-weight: 800;
        letter-spacing: -1px;
        margin: 0.3rem 0;
        color: white;
        line-height: 1.2;
    }

    .hero-sub {
        font-size: 1rem;
        color: rgba(255,255,255,0.75);
        margin-top: 0.5rem;
        font-weight: 400;
    }

    .hero-stats {
        display: flex;
        gap: 2rem;
        margin-top: 1.5rem;
    }

    .hero-stat {
        text-align: center;
    }

    .hero-stat-num {
        font-size: 1.5rem;
        font-weight: 700;
        color: #b7e4c7;
    }

    .hero-stat-label {
        font-size: 0.7rem;
        color: rgba(255,255,255,0.6);
        text-transform: uppercase;
        letter-spacing: 1px;
    }

    .upload-zone {
        border: 2px dashed rgba(64, 145, 108, 0.4);
        border-radius: 12px;
        padding: 2rem;
        text-align: center;
        background: rgba(64, 145, 108, 0.03);
        margin-bottom: 1rem;
    }

    .phase-badge {
        display: inline-flex;
        align-items: center;
        gap: 0.4rem;
        background: rgba(64, 145, 108, 0.1);
        border: 1px solid rgba(64, 145, 108, 0.25);
        border-radius: 20px;
        padding: 0.3rem 0.9rem;
        font-size: 0.72rem;
        font-weight: 600;
        letter-spacing: 1px;
        text-transform: uppercase;
        color: #40916c;
        margin-bottom: 0.75rem;
    }

    .phase-title {
        font-size: 1.3rem;
        font-weight: 700;
        color: inherit;
        margin-bottom: 0.25rem;
    }

    .phase-desc {
        font-size: 0.85rem;
        opacity: 0.55;
        margin-bottom: 1.25rem;
        color: inherit;
    }

    .result-card {
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
        border: 1px solid rgba(64, 145, 108, 0.2);
        background: linear-gradient(135deg, rgba(64,145,108,0.05) 0%, rgba(45,106,79,0.08) 100%);
    }

    .result-tag {
        font-size: 0.7rem;
        font-weight: 600;
        letter-spacing: 1.5px;
        text-transform: uppercase;
        color: #40916c;
        margin-bottom: 0.4rem;
    }

    .result-main {
        font-size: 2rem;
        font-weight: 800;
        color: inherit;
        letter-spacing: -0.5px;
        line-height: 1.1;
    }

    .result-conf {
        font-size: 0.85rem;
        opacity: 0.55;
        margin-top: 0.3rem;
        color: inherit;
    }

    .conf-bar-bg {
        background: rgba(128,128,128,0.12);
        border-radius: 6px;
        height: 6px;
        margin-top: 0.75rem;
        overflow: hidden;
    }

    .conf-bar-fill {
        height: 6px;
        border-radius: 6px;
        background: linear-gradient(90deg, #40916c, #74c69d);
        transition: width 0.6s ease;
    }

    .stop-card {
        border-radius: 12px;
        padding: 1.5rem;
        border: 1px solid rgba(220,53,69,0.25);
        background: rgba(220,53,69,0.05);
        margin: 1rem 0;
    }

    .stop-title {
        font-size: 1.1rem;
        font-weight: 700;
        color: #dc3545;
        margin-bottom: 0.3rem;
    }

    .stop-desc {
        font-size: 0.85rem;
        opacity: 0.7;
        color: inherit;
    }

    .warning-card {
        border-radius: 10px;
        padding: 1rem 1.25rem;
        border: 1px solid rgba(255,193,7,0.3);
        background: rgba(255,193,7,0.05);
        margin: 0.75rem 0;
        font-size: 0.85rem;
        color: inherit;
    }

    .summary-section {
        background: linear-gradient(135deg, rgba(26,71,42,0.06) 0%, rgba(64,145,108,0.08) 100%);
        border: 1px solid rgba(64,145,108,0.15);
        border-radius: 16px;
        padding: 2rem;
        margin-top: 2rem;
    }

    .summary-title {
        font-size: 1.1rem;
        font-weight: 700;
        color: inherit;
        margin-bottom: 1.25rem;
        display: flex;
        align-items: center;
        gap: 0.5rem;
    }

    .summary-grid {
        display: grid;
        grid-template-columns: repeat(3, 1fr);
        gap: 1rem;
    }

    .summary-card {
        background: rgba(255,255,255,0.05);
        border: 1px solid rgba(128,128,128,0.15);
        border-radius: 10px;
        padding: 1.25rem;
        text-align: center;
    }

    .summary-icon {
        font-size: 1.75rem;
        margin-bottom: 0.5rem;
    }

    .summary-label {
        font-size: 0.68rem;
        font-weight: 600;
        letter-spacing: 1.5px;
        text-transform: uppercase;
        opacity: 0.45;
        color: inherit;
        margin-bottom: 0.4rem;
    }

    .summary-value {
        font-size: 1.2rem;
        font-weight: 800;
        color: inherit;
        letter-spacing: -0.3px;
    }

    .summary-conf {
        font-size: 0.75rem;
        color: #40916c;
        font-weight: 600;
        margin-top: 0.2rem;
    }

    .divider {
        border: none;
        border-top: 1px solid rgba(128,128,128,0.1);
        margin: 2rem 0;
    }

    .heatmap-label {
        font-size: 0.72rem;
        font-weight: 600;
        letter-spacing: 1px;
        text-transform: uppercase;
        opacity: 0.45;
        text-align: center;
        margin-top: 0.5rem;
        color: inherit;
    }

    .pipeline-step {
        display: flex;
        align-items: center;
        gap: 0.75rem;
        padding: 0.6rem 0;
        border-bottom: 1px solid rgba(128,128,128,0.08);
    }

    .pipeline-step:last-child {
        border-bottom: none;
    }

    .step-num {
        width: 24px;
        height: 24px;
        border-radius: 50%;
        background: rgba(64,145,108,0.15);
        border: 1px solid rgba(64,145,108,0.3);
        display: flex;
        align-items: center;
        justify-content: center;
        font-size: 0.7rem;
        font-weight: 700;
        color: #40916c;
        flex-shrink: 0;
    }

    .step-text {
        font-size: 0.82rem;
        color: inherit;
        opacity: 0.75;
    }

    .step-text strong {
        opacity: 1;
        color: inherit;
    }
</style>
""", unsafe_allow_html=True)

# ── Model Loading ────────────────────────────────────────────────────────────

@st.cache_resource(show_spinner=False)
def get_models():
    return load_all_models()

with st.spinner("🌿 Initializing TreeVision AI..."):
    try:
        models = get_models()
    except FileNotFoundError as e:
        st.error(str(e))
        st.stop()
    except Exception as e:
        st.error(f"Failed to load models: {e}")
        st.stop()

# ── Sidebar ──────────────────────────────────────────────────────────────────

with st.sidebar:
    st.markdown("## 🌳 TreeVision AI")
    st.markdown("---")

    st.markdown("#### How It Works")
    st.markdown("""
    <div class="pipeline-step">
        <div class="step-num">1</div>
        <div class="step-text"><strong>Tree Detection</strong><br>Is it a tree?</div>
    </div>
    <div class="pipeline-step">
        <div class="step-num">2</div>
        <div class="step-text"><strong>Species Detection</strong><br>Mango or White Gum?</div>
    </div>
    <div class="pipeline-step">
        <div class="step-num">3</div>
        <div class="step-text"><strong>Stage Classification</strong><br>What growth stage?</div>
    </div>
    """, unsafe_allow_html=True)

    st.markdown("---")
    st.markdown("#### Supported Species")
    st.markdown("🥭 **Mango Tree**")
    st.markdown("🌿 **White Gum (Eucalyptus)**")

    st.markdown("---")
    st.markdown("#### Growth Stages")
    stages = [("🌱", "Seedling"), ("🌿", "Sapling"), ("🌳", "Mature"), ("🍂", "Overmature")]
    for icon, stage in stages:
        st.markdown(f"{icon} {stage}")

    st.markdown("---")
    st.markdown("#### Models Used")
    st.caption("Stage 1: ResNet50 (PyTorch) · 92.86% acc")
    st.caption("Stage 2: MobileNetV2 (TensorFlow) · 87.34% acc")
    st.caption("Stage 3: MobileNetV2 (PyTorch) · 95–100% acc")

# ── Hero Section ─────────────────────────────────────────────────────────────

st.markdown("""
<div class="hero-section">
    <div class="hero-badge">🌿 AI-Powered · Computer Vision</div>
    <div class="hero-title">TreeVision AI</div>
    <div class="hero-sub">Upload a tree photograph for instant species identification and growth stage classification — powered by a 4-model deep learning pipeline with Grad-CAM visual explanations.</div>
    <div class="hero-stats">
        <div class="hero-stat">
            <div class="hero-stat-num">4</div>
            <div class="hero-stat-label">AI Models</div>
        </div>
        <div class="hero-stat">
            <div class="hero-stat-num">3</div>
            <div class="hero-stat-label">Pipeline Stages</div>
        </div>
        <div class="hero-stat">
            <div class="hero-stat-num">~95%</div>
            <div class="hero-stat-label">Avg Accuracy</div>
        </div>
        <div class="hero-stat">
            <div class="hero-stat-num">GradCAM</div>
            <div class="hero-stat-label">Explainability</div>
        </div>
    </div>
</div>
""", unsafe_allow_html=True)

# ── Upload ────────────────────────────────────────────────────────────────────

uploaded_file = st.file_uploader(
    "Drop your tree image here",
    type=["jpg", "jpeg", "png"],
    help="Upload a clear photograph of a single tree for best results."
)

if uploaded_file is None:
    col_a, col_b, col_c = st.columns(3)
    with col_a:
        st.info("📸 **Step 1** — Upload a tree photograph")
    with col_b:
        st.info("🤖 **Step 2** — AI pipeline runs automatically")
    with col_c:
        st.info("🔍 **Step 3** — View results with Grad-CAM heatmaps")
    st.stop()

try:
    image = Image.open(uploaded_file).convert("RGB")
except Exception as e:
    st.error(f"Could not open image: {e}")
    st.stop()

st.markdown("<hr class=\'divider\'>", unsafe_allow_html=True)

col_img, col_info = st.columns([1, 1], gap="large")
with col_img:
    st.markdown("<div class=\'heatmap-label\'>YOUR IMAGE</div>", unsafe_allow_html=True)
    st.image(image, use_container_width=True)
with col_info:
    st.markdown("#### Pipeline Overview")
    st.markdown("""
    This image will be analyzed in three sequential stages. Each stage produces a classification result and a Grad-CAM heatmap showing which regions of the image influenced the decision.

    **Grad-CAM** (Gradient-weighted Class Activation Mapping) highlights the areas the neural network focused on — red means high attention, blue means low attention.
    """)
    st.markdown("""
    <div class="pipeline-step">
        <div class="step-num">1</div>
        <div class="step-text"><strong>Tree vs Non-Tree</strong> — ResNet50</div>
    </div>
    <div class="pipeline-step">
        <div class="step-num">2</div>
        <div class="step-text"><strong>Mango vs White Gum</strong> — MobileNetV2</div>
    </div>
    <div class="pipeline-step">
        <div class="step-num">3</div>
        <div class="step-text"><strong>Growth Stage</strong> — Species-specific model</div>
    </div>
    """, unsafe_allow_html=True)

st.markdown("<hr class=\'divider\'>", unsafe_allow_html=True)

# ── Stage 1: Tree Detection ───────────────────────────────────────────────────

st.markdown("""
<div class="phase-badge">🔍 Stage 1</div>
<div class="phase-title">Tree Detection</div>
<div class="phase-desc">ResNet50 determines whether the uploaded image contains a tree.</div>
""", unsafe_allow_html=True)

with st.spinner("Analyzing image for tree presence..."):
    try:
        tree_label, tree_conf, tree_tensor = predict_tree(
            image, models["tree_vs_nontree"], models["device"]
        )
        tree_heatmap = generate_gradcam_pytorch_resnet(
            image, models["tree_vs_nontree"], tree_tensor
        )
    except Exception as e:
        st.error(f"Error in tree detection: {e}")
        st.stop()

col1, col2 = st.columns(2, gap="large")
with col1:
    st.image(image, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>ORIGINAL IMAGE</div>", unsafe_allow_html=True)
with col2:
    st.image(tree_heatmap, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>GRAD-CAM HEATMAP</div>", unsafe_allow_html=True)

if tree_label == "Non-Tree":
    st.markdown(f"""
    <div class="stop-card">
        <div class="stop-title">🚫 No Tree Detected</div>
        <div class="stop-desc">
            The model classified this image as <strong>Non-Tree</strong> with {tree_conf:.1%} confidence.
            The pipeline has stopped. Please upload a clear photograph of a tree.
        </div>
    </div>
    """, unsafe_allow_html=True)
    st.stop()

conf_pct = int(tree_conf * 100)
st.markdown(f"""
<div class="result-card">
    <div class="result-tag">✅ Detection Result</div>
    <div class="result-main">Tree Detected</div>
    <div class="result-conf">Model confidence: {tree_conf:.1%}</div>
    <div class="conf-bar-bg">
        <div class="conf-bar-fill" style="width:{conf_pct}%"></div>
    </div>
</div>
""", unsafe_allow_html=True)

st.markdown("<hr class=\'divider\'>", unsafe_allow_html=True)

# ── Stage 2: Species Detection ────────────────────────────────────────────────

st.markdown("""
<div class="phase-badge">🌿 Stage 2</div>
<div class="phase-title">Species Detection</div>
<div class="phase-desc">MobileNetV2 (TensorFlow) identifies whether this is a Mango or White Gum tree.</div>
""", unsafe_allow_html=True)

st.markdown("""
<div class="warning-card">
    ⚠️ <strong>Note:</strong> This system is trained exclusively on Mango and White Gum (Eucalyptus) trees.
    Accuracy may be reduced for other species.
</div>
""", unsafe_allow_html=True)

with st.spinner("Identifying tree species..."):
    try:
        species_label, species_conf, species_array = predict_species(
            image, models["species"]
        )
        species_class_idx = 1 if species_label == "Mango" else 0
        species_heatmap = generate_gradcam_keras(
            image, models["species"], species_array, species_class_idx
        )
    except Exception as e:
        st.error(f"Error in species detection: {e}")
        st.stop()

col3, col4 = st.columns(2, gap="large")
with col3:
    st.image(image, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>ORIGINAL IMAGE</div>", unsafe_allow_html=True)
with col4:
    st.image(species_heatmap, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>GRAD-CAM HEATMAP</div>", unsafe_allow_html=True)

species_icon = "🥭" if species_label == "Mango" else "🌿"
conf_pct2 = int(species_conf * 100)
st.markdown(f"""
<div class="result-card">
    <div class="result-tag">🌿 Species Result</div>
    <div class="result-main">{species_icon} {species_label} Tree</div>
    <div class="result-conf">Model confidence: {species_conf:.1%}</div>
    <div class="conf-bar-bg">
        <div class="conf-bar-fill" style="width:{conf_pct2}%"></div>
    </div>
</div>
""", unsafe_allow_html=True)

st.markdown("<hr class=\'divider\'>", unsafe_allow_html=True)

# ── Stage 3: Growth Stage ─────────────────────────────────────────────────────

st.markdown(f"""
<div class="phase-badge">🌱 Stage 3</div>
<div class="phase-title">Growth Stage Classification</div>
<div class="phase-desc">Species-specific MobileNetV2 model classifies the tree's growth stage.</div>
""", unsafe_allow_html=True)

if species_label == "Mango":
    stage_model = models["mango_stage"]
    model_used = "Mango Stage Model (95.83% accuracy)"
else:
    stage_model = models["gum_stage"]
    model_used = "White Gum Stage Model (100% accuracy)"

st.caption(f"🔧 Using: {model_used}")

with st.spinner("Classifying growth stage..."):
    try:
        stage_label, stage_conf, stage_tensor = predict_stage(
            image, stage_model, models["device"]
        )
        stage_heatmap = generate_gradcam_pytorch_mobilenet(
            image, stage_model, stage_tensor
        )
    except Exception as e:
        st.error(f"Error in stage classification: {e}")
        st.stop()

col5, col6 = st.columns(2, gap="large")
with col5:
    st.image(image, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>ORIGINAL IMAGE</div>", unsafe_allow_html=True)
with col6:
    st.image(stage_heatmap, use_container_width=True)
    st.markdown("<div class=\'heatmap-label\'>GRAD-CAM HEATMAP</div>", unsafe_allow_html=True)

stage_icons = {"Seedling": "🌱", "Sapling": "🌿", "Mature": "🌳", "Overmature": "🍂"}
stage_icon = stage_icons.get(stage_label, "🌿")
conf_pct3 = int(stage_conf * 100)
st.markdown(f"""
<div class="result-card">
    <div class="result-tag">🌱 Growth Stage</div>
    <div class="result-main">{stage_icon} {stage_label}</div>
    <div class="result-conf">Model confidence: {stage_conf:.1%}</div>
    <div class="conf-bar-bg">
        <div class="conf-bar-fill" style="width:{conf_pct3}%"></div>
    </div>
</div>
""", unsafe_allow_html=True)

# ── Summary ───────────────────────────────────────────────────────────────────

st.markdown(f"""
<div class="summary-section">
    <div class="summary-title">📊 Classification Summary</div>
    <div class="summary-grid">
        <div class="summary-card">
            <div class="summary-icon">🌳</div>
            <div class="summary-label">Detection</div>
            <div class="summary-value">Tree</div>
            <div class="summary-conf">{tree_conf:.1%} confidence</div>
        </div>
        <div class="summary-card">
            <div class="summary-icon">{species_icon}</div>
            <div class="summary-label">Species</div>
            <div class="summary-value">{species_label}</div>
            <div class="summary-conf">{species_conf:.1%} confidence</div>
        </div>
        <div class="summary-card">
            <div class="summary-icon">{stage_icon}</div>
            <div class="summary-label">Growth Stage</div>
            <div class="summary-value">{stage_label}</div>
            <div class="summary-conf">{stage_conf:.1%} confidence</div>
        </div>
    </div>
</div>
""", unsafe_allow_html=True)