File size: 2,566 Bytes
abe3536 | 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 | import streamlit as st
import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import matplotlib.ticker as ticker
from io import BytesIO
from PIL import Image
# -----------------------------
# Streamlit App
# -----------------------------
st.title("🖼️ Image Segmentation using K-Means Clustering")
# Sidebar options
st.sidebar.header("Controls")
uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
max_k = st.sidebar.slider("Max K for Elbow Curve", 5, 15, 10)
chosen_k = st.sidebar.slider("Choose K (Clusters)", 2, 10, 4)
if uploaded_file is not None:
# -----------------------------
# Step 1: Load Image
# -----------------------------
image = Image.open(uploaded_file)
image = np.array(image)
st.subheader("Original Image")
st.image(image, use_container_width=True)
# -----------------------------
# Step 2: Preprocess Image
# -----------------------------
pixels = image.reshape((-1, 3))
pixels = np.float32(pixels)
# -----------------------------
# Step 3: Elbow Curve
# -----------------------------
wcss = []
K = range(1, max_k+1)
for k in K:
kmeans = KMeans(n_clusters=k, random_state=42)
kmeans.fit(pixels)
wcss.append(kmeans.inertia_)
# Plot elbow curve
fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(K, wcss, 'bo-')
ax.set_xlabel("Number of clusters (k)")
ax.set_ylabel("WCSS (Inertia)")
ax.set_title("Elbow Curve")
# Format y-axis in millions
ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, _: f'{int(x/1e6)}M'))
# Add vertical dotted line for chosen k
ax.axvline(x=chosen_k, color='r', linestyle='--', label=f'k={chosen_k}')
ax.legend()
st.subheader("Elbow Curve")
st.pyplot(fig)
# -----------------------------
# Step 4: Apply KMeans with chosen k
# -----------------------------
kmeans = KMeans(n_clusters=chosen_k, random_state=42)
labels = kmeans.fit_predict(pixels)
segmented_img = kmeans.cluster_centers_[labels]
segmented_img = segmented_img.reshape(image.shape)
segmented_img = np.uint8(segmented_img)
# -----------------------------
# Step 5: Show Results
# -----------------------------
st.subheader(f"Segmented Image (k={chosen_k})")
st.image(segmented_img, use_container_width=True)
else:
st.info("👈 Upload an image from the sidebar to start.")
|