from pathlib import Path
import pickle
import json
import zipfile

import keras
import numpy as np
from PIL import Image
import streamlit as st


ROOT = Path(__file__).resolve().parents[1]
SATELLITE_LABELS = [
    "Sea or Lake", "Highway", "River", "Pasture", "Industrial",
    "Residential", "Permanent Crop", "Annual Crop", "Forest",
    "Herbaceous Vegetation",
]
st.set_page_config(page_title="Traffic & Satellite Vision Hub", page_icon="🛰️", layout="wide")


@st.cache_resource
def load_traffic_assets():
    model = keras.models.load_model(ROOT / "traffic_sign.h5", compile=False)
    with open(ROOT / "labels.pkl", "rb") as file:
        labels = pickle.load(file)
    return model, labels


@st.cache_resource
def load_satellite_model():
    path = ROOT / "satellite_image_tf.keras"
    try:
        return keras.models.load_model(path, compile=False)
    except ValueError:
        # Repair the nested EfficientNet output stored as a one-item list.
        with zipfile.ZipFile(path) as archive:
            config = json.loads(archive.read("config.json"))
        layers = config["config"]["layers"]
        raw_base = keras.saving.deserialize_keras_object(
            layers[1], custom_objects={"Functional": keras.Model}
        )
        base = keras.Model(raw_base.input, raw_base.output[0])
        tail = [keras.layers.deserialize(layer) for layer in layers[2:]]
        model = keras.Sequential([keras.Input((224, 224, 3)), base, *tail])
        model.load_weights(path)
        return model


def image_array(image, size):
    resized = image.convert("RGB").resize(size, Image.Resampling.LANCZOS)
    return resized, np.expand_dims(np.asarray(resized, dtype="float32") / 255.0, 0)


def show_ranked(prediction, labels, count=3):
    indices = np.argsort(prediction)[-count:][::-1]
    for rank, index in enumerate(indices, 1):
        st.write(f"{rank}. **{labels[int(index)]}** — {float(prediction[index]):.1%}")


st.title("🛰️ Traffic & Satellite Vision Hub")
st.caption("Two end-to-end computer-vision projects with real TensorFlow model inference.")
traffic_tab, satellite_tab = st.tabs(["German Traffic Signs", "Satellite Land Cover"])

with traffic_tab:
    st.subheader("German traffic sign recognition")
    st.write("Photograph or upload a close, centered traffic sign. The model recognizes 43 GTSRB classes.")
    camera = st.camera_input("Take a traffic sign photo")
    uploaded = st.file_uploader("Or upload an image", ["jpg", "jpeg", "png"], key="traffic")
    traffic_file = camera if camera is not None else uploaded
    if traffic_file:
        image = Image.open(traffic_file)
        st.image(image, width=340)
        if st.button("Classify traffic sign", type="primary", use_container_width=True):
            model, label_map = load_traffic_assets()
            _, model_input = image_array(image, (30, 30))
            prediction = np.asarray(model.predict(model_input, verbose=0))[0]
            best = int(np.argmax(prediction))
            st.metric("Prediction", label_map[best], f"{float(prediction[best]):.1%} confidence")
            st.subheader("Top 3 predictions")
            show_ranked(prediction, label_map)
            if float(prediction[best]) < 0.65:
                st.warning("Low confidence: crop the image so the sign is clear and fills most of the frame.")

with satellite_tab:
    st.subheader("EuroSAT land-cover classification")
    st.write("Upload a top-down satellite tile. The EfficientNet-based model distinguishes 10 land-cover types.")
    satellite_file = st.file_uploader("Upload satellite image", ["jpg", "jpeg", "png"], key="satellite")
    if satellite_file:
        image = Image.open(satellite_file)
        st.image(image, width=340)
        if st.button("Classify land cover", type="primary", use_container_width=True):
            model = load_satellite_model()
            _, model_input = image_array(image, (224, 224))
            prediction = np.asarray(model.predict(model_input, verbose=0))[0]
            best = int(np.argmax(prediction))
            st.metric("Predicted land cover", SATELLITE_LABELS[best], f"{float(prediction[best]):.1%} confidence")
            st.subheader("Top 3 predictions")
            show_ranked(prediction, SATELLITE_LABELS)

with st.expander("About these projects"):
    st.write("The traffic model uses the 43-class German Traffic Sign Recognition Benchmark. The satellite model uses the 10-class EuroSAT Sentinel-2 land-cover benchmark. Predictions are portfolio demonstrations and should not be used for driving or operational mapping decisions.")