app.py

"""
app.py
------
Streamlit front-end for the Crop vs No-Crop SegFormer-B2 demo.
"""

import logging

import streamlit as st

import utils
from predictor import (
    INPUT_SIZE,
    MODEL_NAME,
    CropSegmenter,
    InferenceError,
    ModelLoadError,
)

logging.basicConfig(level=logging.INFO)

st.set_page_config(
    page_title="Crop vs No-Crop Segmentation (Crop Contains only Wheat and Tomato)",
    page_icon="🌾",
    layout="wide",
)


@st.cache_resource(show_spinner="Loading segmentation model (first run only)…")
def get_segmenter() -> CropSegmenter:
    return CropSegmenter()


def render_sidebar() -> None:
    st.sidebar.header("ℹ️ Model Information")
    st.sidebar.markdown(
        f"""
- **Architecture:** NVIDIA SegFormer-B2
- **Backbone:** `{MODEL_NAME}`
- **Task:** Binary semantic segmentation
- **Weights:** `best.pt`
"""
    )

    st.sidebar.header("🎨 Class Definitions")
    st.sidebar.markdown(
        """
| Class | Meaning | Mask color |
|:-----:|---------|------------|
| 0 | Background / No Crop | Dark gray |
| 1 | Crop | Bright yellow |
"""
    )


def main() -> None:
    render_sidebar()

    st.title("🌾 Crop vs No-Crop Segmentation")
    st.markdown(
        "Upload a field image and the fine-tuned **SegFormer-B2** model "
        "segments it into **Crop** and **Background (No Crop)** regions."
    )

    try:
        segmenter = get_segmenter()
    except ModelLoadError as exc:
        st.error(f"🚫 Could not load the model: {exc}")
        st.stop()
    except Exception as exc:  # noqa: BLE001
        st.error(f"🚫 Unexpected error while loading the model: {exc}")
        st.stop()

    uploaded = st.file_uploader(
        "Upload an image", type=["jpg", "jpeg", "png"], accept_multiple_files=False
    )
    if uploaded is None:
        st.info("👆 Upload a JPG / JPEG / PNG image to get started.")
        return

    try:
        image = utils.load_image(uploaded)
    except utils.ImageLoadError as exc:
        st.error(f"🚫 {exc}")
        return

    st.subheader("Uploaded image")
    st.image(image, caption="Input", use_container_width=True)

    if not st.button("🚀 Run Segmentation", type="primary"):
        return

    try:
        with st.spinner("Running segmentation… a few seconds on CPU."):
            mask = segmenter.predict(
                image,
                input_size=INPUT_SIZE,
                swap_classes=False,
                normalization="imagenet",
            )
            color_mask = utils.create_color_mask(mask)
            overlay = utils.create_overlay(image, mask)
            combined = utils.combine_panels(image, color_mask, overlay)
    except InferenceError as exc:
        st.error(f"🚫 {exc}")
        return
    except Exception as exc:  # noqa: BLE001
        st.error(f"🚫 Unexpected error during inference: {exc}")
        return

    st.subheader("Results")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.image(image, caption="Original", use_container_width=True)
    with col2:
        st.image(color_mask, caption="Predicted Mask", use_container_width=True)
    with col3:
        st.image(overlay, caption="Crop Overlay", use_container_width=True)

    crop_ratio = float((mask == 1).mean()) * 100.0
    st.metric("Estimated crop coverage", f"{crop_ratio:.1f}%")

    summary = segmenter.last_logit_summary
    if summary:
        with st.expander("📈 More Details"):
            st.write(
                f"Mean logit gap (crop − background): "
                f"{summary['mean_gap_crop_minus_bg']:.3f}  "
                f"(positive favors crop, negative favors background)"
            )
            st.write(
                f"Logit range: {summary['logit_min']:.2f} … {summary['logit_max']:.2f}"
            )

    png_bytes = utils.image_to_png_bytes(combined)
    st.download_button(
        label="⬇️ Download Result",
        data=png_bytes,
        file_name="prediction_result.png",
        mime="image/png",
    )


if __name__ == "__main__":
    main()