File size: 4,038 Bytes
ad83752
 
 
7fb5eb7
ad83752
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7913754
ad83752
 
 
 
 
 
 
 
 
 
7fb5eb7
ad83752
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e1c2dc2
ad83752
 
7fb5eb7
ad83752
 
 
e1c2dc2
 
ad83752
 
 
 
 
 
 
 
 
 
 
 
e1c2dc2
ad83752
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e1c2dc2
 
 
7fb5eb7
 
 
e1c2dc2
ad83752
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6ba8390
 
 
7fb5eb7
6ba8390
 
 
 
 
 
 
 
 
ad83752
 
 
 
 
 
 
 
 
 
7fb5eb7
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""
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()