| |
|
|
| import streamlit as st |
| import cv2 |
| import numpy as np |
| from ultralytics import YOLO |
| from PIL import Image |
| import matplotlib.pyplot as plt |
|
|
| st.title("YOLOv8 Image Segmentation with Webcam & Elbow Curve") |
|
|
| |
| |
| |
| @st.cache_resource |
| def load_model(model_name="yolov8n-seg.pt"): |
| return YOLO(model_name) |
|
|
| model = load_model() |
|
|
| |
| |
| |
| input_option = st.radio("Input Method", ["Upload Image", "Webcam Capture"]) |
|
|
| if input_option == "Upload Image": |
| uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) |
| if uploaded_file is not None: |
| image = Image.open(uploaded_file).convert("RGB") |
| img_array = np.array(image) |
|
|
| elif input_option == "Webcam Capture": |
| webcam_image = st.camera_input("Capture an image") |
| if webcam_image is not None: |
| image = Image.open(webcam_image).convert("RGB") |
| img_array = np.array(image) |
| else: |
| st.stop() |
|
|
| |
| |
| |
| if 'img_array' in locals(): |
| with st.spinner("Running YOLOv8 segmentation..."): |
| results = model(img_array) |
| annotated_img = results[0].plot() |
|
|
| |
| col1, col2 = st.columns(2) |
| with col1: |
| st.image(image, caption="Original Image", use_column_width=True) |
| with col2: |
| st.image(annotated_img, caption="Segmented Image", use_column_width=True) |
|
|
| |
| annotated_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR) |
| cv2.imwrite("segmented_output.jpg", annotated_bgr) |
| with open("segmented_output.jpg", "rb") as f: |
| st.download_button("Download Segmented Image", f, "segmented_output.jpg") |
|
|
| |
| |
| |
| st.subheader("Elbow Curve") |
|
|
| try: |
| from sklearn.cluster import KMeans |
|
|
| |
| masks = getattr(results[0], "masks", None) |
| centroids = [] |
|
|
| if masks is not None: |
| try: |
| mask_np = masks.data |
| except AttributeError: |
| mask_np = masks |
|
|
| |
| if hasattr(mask_np, "numpy"): |
| mask_np = mask_np.numpy() |
| elif isinstance(mask_np, list): |
| mask_np = np.array(mask_np) |
|
|
| if mask_np.ndim >= 3: |
| for m in mask_np: |
| y_idx, x_idx = np.where(m > 0.5) |
| if len(x_idx) > 0: |
| centroids.append([np.mean(x_idx), np.mean(y_idx)]) |
|
|
| |
| if len(centroids) == 0: |
| st.info("No masks detected — using random points for elbow curve.") |
| centroids = np.random.rand(10, 2) * 100 |
|
|
| centroids = np.array(centroids) |
|
|
| |
| if len(centroids) < 2: |
| st.info("Too few points for elbow curve.") |
| else: |
| max_k = min(10, len(centroids)) |
| wcss = [] |
| for k in range(1, max_k + 1): |
| kmeans = KMeans(n_clusters=k, random_state=42).fit(centroids) |
| wcss.append(kmeans.inertia_) |
|
|
| |
| selected_k = st.slider("Select k to highlight on Elbow Curve", 1, max_k, value=3) |
|
|
| |
| fig, ax = plt.subplots() |
| ax.plot(range(1, max_k + 1), wcss, marker='o', label='WCSS') |
| ax.axvline(x=selected_k, color='r', linestyle='--', label=f'Selected k = {selected_k}') |
| ax.set_xlabel("Number of clusters (k)") |
| ax.set_ylabel("WCSS") |
| ax.set_title("Elbow Curve") |
| ax.legend() |
| st.pyplot(fig) |
|
|
| except Exception as e: |
| st.info(f"Elbow curve unavailable: {e}") |
