app.py

import sys
from pathlib import Path

import streamlit as st
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import os

ROOT_DIR = Path(__file__).resolve().parents[1]
if str(ROOT_DIR) not in sys.path:
    sys.path.append(str(ROOT_DIR))

from model_loader import load_model, predict

st.set_page_config(page_title="Xplendia Off-Road AI", layout="wide")

st.title("🚙 Xplendia Off-Road Autonomy Vision")
st.write("Upload an off-road image and view AI terrain segmentation output.")


@st.cache_resource
def get_model():
    return load_model()

uploaded = st.file_uploader("Upload vehicle camera image", type=["png", "jpg", "jpeg"])

st.sidebar.header("Color Legend")
st.sidebar.write("""
Cyan / Blue → Sky  
Orange → Soil / terrain  
Green → Vegetation / grass  
Olive → Bushes / dry grass  
Purple → Ground clutter  
Gray / Brown → Rocks / rough terrain  
""")

if uploaded:
    image = Image.open(uploaded)

    st.subheader("Vehicle Camera Input")
    st.image(image, use_container_width=True)

    st.subheader("AI Segmentation Output")

    try:
        model = get_model()
        prediction = predict(model, image)

        fig, ax = plt.subplots()
        ax.imshow(np.array(prediction))
        ax.axis("off")
        st.pyplot(fig)

        st.success("Terrain map generated successfully.")
        st.info("Driving Decision: Proceed carefully. Avoid dense rocks / vegetation-heavy regions.")
    except FileNotFoundError as exc:
        st.error(f"Model file not found: {exc}")
    except Exception as exc:
        st.error(f"Prediction failed: {exc}")