LEGENDFTW/image-segmentation-explorer
0
1import sys2import os3sys.path.insert(0, os.path.dirname(__file__))4 5import streamlit as st6import numpy as np7from PIL import Image8import io9import json10import matplotlib11matplotlib.use("Agg")12import matplotlib.pyplot as plt13 14from segmentation import (15 segment_threshold, segment_kmeans, segment_watershed,16 compute_segment_metrics, METHOD_DESCRIPTIONS17)18from utils import load_sample_image19 20# ── Page config ───────────────────────────────────────────────────────────────21st.set_page_config(22 page_title="Image Segmentation Explorer",23 layout="wide",24)25 26# ── Cache helpers ─────────────────────────────────────────────────────────────27def arr_to_bytes(arr: np.ndarray) -> bytes:28 buf = io.BytesIO()29 Image.fromarray(arr.astype(np.uint8)).save(buf, format="PNG")30 return buf.getvalue()31 32@st.cache_data(show_spinner=False)33def cached_threshold(img_bytes, threshold, mode):34 img = np.array(Image.open(io.BytesIO(img_bytes)))35 return segment_threshold(img, threshold, mode)36 37@st.cache_data(show_spinner=False)38def cached_kmeans(img_bytes, k):39 img = np.array(Image.open(io.BytesIO(img_bytes)))40 return segment_kmeans(img, k)41 42@st.cache_data(show_spinner=False)43def cached_watershed(img_bytes, min_distance):44 img = np.array(Image.open(io.BytesIO(img_bytes)))45 return segment_watershed(img, min_distance)46 47@st.cache_data(show_spinner=False)48def cached_kmeans_chart(label_map_bytes: bytes, k: int) -> bytes:49 label_map = np.frombuffer(label_map_bytes, dtype=np.uint8).reshape(-1)50 colors_hex = ["#e74c3c","#2ecc71","#3498db","#9b59b6","#f1c40f","#e67e22","#1abc9c"]51 sizes = [int(np.sum(label_map == i)) for i in range(k)]52 labels = [f"Cluster {i+1}" for i in range(k)]53 fig, ax = plt.subplots(figsize=(5, 3), facecolor="#1a1a1a")54 ax.set_facecolor("#1a1a1a")55 ax.bar(labels, sizes, color=colors_hex[:k], edgecolor="#444")56 ax.set_ylabel("Pixel count", color="white")57 ax.set_title("Cluster sizes", color="white")58 ax.tick_params(colors="gray")59 for spine in ax.spines.values():60 spine.set_color("#444")61 fig.tight_layout()62 buf = io.BytesIO()63 fig.savefig(buf, format="png", facecolor="#1a1a1a", dpi=100)64 plt.close(fig)65 buf.seek(0)66 return buf.read()67 68# ── CSS ───────────────────────────────────────────────────────────────────────69st.markdown("""70<style>71 .main-title { font-size:2.2rem; font-weight:700; margin-bottom:0.2rem; }72 .sub-title { font-size:1rem; color:#888; margin-bottom:1.5rem; }73 .metric-card { background:rgba(255,255,255,0.06); border:1px solid rgba(255,255,255,0.1);74 border-radius:10px; padding:14px 18px; margin:6px 0; }75 .metric-label{ font-size:0.78rem; color:#aaa; font-weight:600;76 letter-spacing:.05em; text-transform:uppercase; }77 .metric-value{ font-size:1.6rem; font-weight:700; }78 .method-desc { background:rgba(255,255,255,0.06); border-radius:8px;79 padding:12px 16px; font-size:0.9rem; margin-bottom:1rem; }80 .info-box { border-left:4px solid #f4c430; border-radius:6px;81 padding:12px 16px; font-size:0.88rem; }82</style>83""", unsafe_allow_html=True)84 85# ── Header ────────────────────────────────────────────────────────────────────86st.markdown('<p class="main-title">Image Segmentation Explorer</p>', unsafe_allow_html=True)87st.markdown('<p class="sub-title">Compare three classic segmentation algorithms — '88 'thresholding, K-means, and watershed — and understand when each works best.</p>',89 unsafe_allow_html=True)90 91# ── Sidebar ───────────────────────────────────────────────────────────────────92with st.sidebar:93 st.header("Controls")94 95 st.subheader("1. Image source")96 source = st.radio("Choose input", ["Built-in sample", "Upload your own"], horizontal=True)97 98 img_array = None99 if source == "Built-in sample":100 sample_name = st.selectbox(101 "Sample image",102 ["shapes.png", "coins.png", "bands.png"],103 format_func=lambda x: x.replace(".png", "").title()104 )105 img_array = load_sample_image(sample_name)106 hints = {107 "shapes.png": "Coloured shapes — try K-means (k=4) or thresholding.",108 "coins.png": "Overlapping circles — try Watershed.",109 "bands.png": "Colour bands — try K-means (k=3).",110 }111 st.caption(hints[sample_name])112 else:113 uploaded = st.file_uploader("Upload image (JPG / PNG)", type=["jpg", "jpeg", "png"])114 if uploaded:115 pil = Image.open(uploaded).convert("RGB")116 if max(pil.size) > 600:117 pil.thumbnail((600, 600), Image.LANCZOS)118 st.caption("Image resized to 600px max for performance.")119 img_array = np.array(pil)120 else:121 st.info("Upload an image or switch to a built-in sample.")122 123 st.subheader("2. Method")124 method = st.selectbox("Algorithm", ["Thresholding", "K-means", "Watershed"])125 126 st.subheader("3. Parameters")127 params = {}128 if method == "Thresholding":129 params["mode"] = st.selectbox("Mode", ["Binary", "Otsu", "Adaptive"],130 help="Otsu automatically finds the best threshold.")131 if params["mode"] == "Binary":132 params["threshold"] = st.slider("Threshold", 0, 255, 127,133 help="Pixels above this value become foreground.")134 else:135 params["threshold"] = 127136 137 elif method == "K-means":138 params["k"] = st.slider("Number of clusters (K)", 2, 8, 3,139 help="How many colour groups to find.")140 141 elif method == "Watershed":142 params["min_distance"] = st.slider("Minimum distance between seeds", 5, 50, 15,143 help="Larger = fewer, bigger regions.")144 145# ── Guard ─────────────────────────────────────────────────────────────────────146if img_array is None:147 st.markdown('<div class="info-box">Choose a built-in sample or upload your own image using the sidebar.</div>',148 unsafe_allow_html=True)149 st.stop()150 151# ── Run segmentation (cached) ─────────────────────────────────────────────────152img_bytes = arr_to_bytes(img_array)153 154with st.spinner("Running Watershed..." if method == "Watershed" else ""):155 try:156 if method == "Thresholding":157 segmented, label_map = cached_threshold(img_bytes, params["threshold"], params["mode"])158 elif method == "K-means":159 segmented, label_map = cached_kmeans(img_bytes, params["k"])160 else:161 segmented, label_map = cached_watershed(img_bytes, params["min_distance"])162 except Exception as e:163 st.error(f"Segmentation error: {e}")164 st.stop()165 166metrics = compute_segment_metrics(label_map, method)167 168# ── Tabs ──────────────────────────────────────────────────────────────────────169tab1, tab2, tab3 = st.tabs(["Comparison", "Diagnostics", "Theory"])170 171# ── Tab 1: Comparison ─────────────────────────────────────────────────────────172with tab1:173 st.markdown(f'<div class="method-desc"><b>{method}</b>: {METHOD_DESCRIPTIONS[method]}</div>',174 unsafe_allow_html=True)175 176 col1, col2 = st.columns(2)177 with col1:178 st.image(img_array, caption="Original", use_container_width=True)179 with col2:180 st.image(segmented, caption=f"Segmented ({method})", use_container_width=True)181 182# ── Tab 2: Diagnostics ────────────────────────────────────────────────────────183with tab2:184 num_seg = metrics["num_segments"]185 st.markdown(f"""<div class="metric-card" style="display:inline-block;min-width:180px;">186 <div class="metric-label">Segments found</div>187 <div class="metric-value">{num_seg}</div>188 </div>""", unsafe_allow_html=True)189 190 if method == "Thresholding":191 c1, c2 = st.columns(2)192 with c1:193 st.markdown(f"""<div class="metric-card">194 <div class="metric-label">Foreground</div>195 <div class="metric-value">{metrics['foreground_pct']}%</div>196 <div style="font-size:.75rem;color:#aaa">of total pixels</div>197 </div>""", unsafe_allow_html=True)198 with c2:199 st.markdown(f"""<div class="metric-card">200 <div class="metric-label">Background</div>201 <div class="metric-value">{metrics['background_pct']}%</div>202 <div style="font-size:.75rem;color:#aaa">of total pixels</div>203 </div>""", unsafe_allow_html=True)204 205 elif method == "K-means":206 st.divider()207 st.subheader("Cluster size distribution")208 k = params["k"]209 flat = label_map.flatten().astype(np.uint8)210 chart = cached_kmeans_chart(flat.tobytes(), k)211 st.image(chart, use_container_width=True)212 st.caption("Each bar shows how many pixels belong to that colour cluster.")213 214 elif method == "Watershed":215 st.markdown(f"""<div class="metric-card" style="margin-top:8px;">216 <div class="metric-label">Average region size</div>217 <div class="metric-value">{metrics['avg_region_size']} px</div>218 </div>""", unsafe_allow_html=True)219 220 st.divider()221 st.subheader("Pixel intensity comparison")222 223 fig, axes = plt.subplots(1, 2, figsize=(10, 3), facecolor="#1a1a1a")224 for ax, arr, title in zip(axes, [img_array, segmented], ["Original", "Segmented"]):225 ax.set_facecolor("#1a1a1a")226 colors = ["#e74c3c", "#2ecc71", "#3498db"]227 for c, color in enumerate(colors):228 hist, edges = np.histogram(arr[:, :, c].ravel(), bins=64, range=(0, 255))229 ax.bar(edges[:-1], hist, width=4, color=color, alpha=0.6)230 ax.set_title(title, color="white", fontsize=10)231 ax.tick_params(colors="gray")232 for spine in ax.spines.values():233 spine.set_color("#444")234 fig.tight_layout()235 buf = io.BytesIO()236 fig.savefig(buf, format="png", facecolor="#1a1a1a", dpi=100)237 plt.close(fig)238 buf.seek(0)239 st.image(buf.read(), caption="Left: original pixel distribution. Right: after segmentation.",240 use_container_width=True)241 st.caption("K-means and thresholding collapse the histogram into sharp peaks — one per segment.")242 243# ── Tab 3: Theory ─────────────────────────────────────────────────────────────244with tab3:245 st.subheader("How image segmentation works")246 st.markdown("""247Image segmentation divides an image into meaningful regions — grouping pixels that belong together.248The goal is to simplify the image for further analysis, like object detection or measurement.249 250---251### Method comparison252 253| Method | Approach | Best for | Weakness |254|---|---|---|---|255| **Thresholding** | Pixel intensity cutoff | Simple bright/dark separation | Fails with uneven lighting |256| **K-means** | Colour cluster assignment | Colour-based regions | Ignores spatial structure |257| **Watershed** | Flood-fill from seed points | Touching/overlapping objects | Sensitive to noise |258 259---260### Thresholding261 262The simplest method. Convert to greyscale, then apply a cutoff value:263- Pixels **above** the threshold become foreground264- Pixels **below** the threshold become background265 266**Otsu's method** automatically finds the best threshold by minimising within-class variance.267**Adaptive thresholding** uses a local threshold for each region — handles uneven lighting.268 269---270### K-means clustering271 272Treats each pixel as a point in 3D colour space (R, G, B).273Assigns pixels to K clusters based on colour similarity.274Each cluster gets one representative colour — the mean of all its pixels.275 276The algorithm iterates:2771. Assign each pixel to the nearest cluster centre2782. Recompute cluster centres2793. Repeat until stable280 281K-means does not consider where pixels are in the image — only their colour.282 283---284### Watershed285 286Inspired by geography. Think of the greyscale image as a topographic surface:287- Dark areas are valleys (objects)288- Bright areas are ridges (boundaries)289 290The algorithm floods the surface from seed points. Where two floods meet, a boundary forms.291The distance transform finds seeds automatically — peaks of the distance-to-background map.292This makes it very good at separating objects that touch or overlap.293 294---295### When each method fails296 297- **Thresholding**: fails when objects and background have similar brightness, or lighting is uneven.298- **K-means**: merges objects of the same colour even if they are spatially separate.299- **Watershed**: noisy images create too many false seed points, causing over-segmentation.300 """)