Srikanth0804/Image_Segmentation_With_Elbow_Curve

1# streamlit_yolo_seg_webcam_fixed.py2 3import streamlit as st4import cv25import numpy as np6from ultralytics import YOLO7from PIL import Image8import matplotlib.pyplot as plt9 10st.title("YOLOv8 Image Segmentation with Webcam & Elbow Curve")11 12# -------------------------------13# Load YOLO Segmentation Model14# -------------------------------15@st.cache_resource16def load_model(model_name="yolov8n-seg.pt"):17    return YOLO(model_name)18 19model = load_model()20 21# -------------------------------22# Input Method23# -------------------------------24input_option = st.radio("Input Method", ["Upload Image", "Webcam Capture"])25 26if input_option == "Upload Image":27    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])28    if uploaded_file is not None:29        image = Image.open(uploaded_file).convert("RGB")30        img_array = np.array(image)31 32elif input_option == "Webcam Capture":33    webcam_image = st.camera_input("Capture an image")34    if webcam_image is not None:35        image = Image.open(webcam_image).convert("RGB")36        img_array = np.array(image)37else:38    st.stop()39 40# -------------------------------41# Run Segmentation42# -------------------------------43if 'img_array' in locals():44    with st.spinner("Running YOLOv8 segmentation..."):45        results = model(img_array)46        annotated_img = results[0].plot()47 48    # Display images49    col1, col2 = st.columns(2)50    with col1:51        st.image(image, caption="Original Image", use_column_width=True)52    with col2:53        st.image(annotated_img, caption="Segmented Image", use_column_width=True)54 55    # Download annotated image56    annotated_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)57    cv2.imwrite("segmented_output.jpg", annotated_bgr)58    with open("segmented_output.jpg", "rb") as f:59        st.download_button("Download Segmented Image", f, "segmented_output.jpg")60 61    # -------------------------------62    # Elbow Curve (robust)63    # -------------------------------64    st.subheader("Elbow Curve")65 66    try:67        from sklearn.cluster import KMeans68 69        # Extract centroids from masks70        masks = getattr(results[0], "masks", None)71        centroids = []72 73        if masks is not None:74            try:75                mask_np = masks.data76            except AttributeError:77                mask_np = masks78 79            # Convert to numpy80            if hasattr(mask_np, "numpy"):81                mask_np = mask_np.numpy()82            elif isinstance(mask_np, list):83                mask_np = np.array(mask_np)84 85            if mask_np.ndim >= 3:86                for m in mask_np:87                    y_idx, x_idx = np.where(m > 0.5)88                    if len(x_idx) > 0:89                        centroids.append([np.mean(x_idx), np.mean(y_idx)])90 91        # Fallback if no masks found92        if len(centroids) == 0:93            st.info("No masks detected — using random points for elbow curve.")94            centroids = np.random.rand(10, 2) * 10095 96        centroids = np.array(centroids)97 98        # Only compute elbow if >= 2 points99        if len(centroids) < 2:100            st.info("Too few points for elbow curve.")101        else:102            max_k = min(10, len(centroids))103            wcss = []104            for k in range(1, max_k + 1):105                kmeans = KMeans(n_clusters=k, random_state=42).fit(centroids)106                wcss.append(kmeans.inertia_)107 108            # Slider to select k109            selected_k = st.slider("Select k to highlight on Elbow Curve", 1, max_k, value=3)110 111            # Plot elbow curve112            fig, ax = plt.subplots()113            ax.plot(range(1, max_k + 1), wcss, marker='o', label='WCSS')114            ax.axvline(x=selected_k, color='r', linestyle='--', label=f'Selected k = {selected_k}')115            ax.set_xlabel("Number of clusters (k)")116            ax.set_ylabel("WCSS")117            ax.set_title("Elbow Curve")118            ax.legend()119            st.pyplot(fig)120 121    except Exception as e:122        st.info(f"Elbow curve unavailable: {e}")