File size: 4,012 Bytes
f5fd1ea dcbec55 f5fd1ea da228e3 | 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 | # streamlit_yolo_seg_webcam_fixed.py
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")
# -------------------------------
# Load YOLO Segmentation Model
# -------------------------------
@st.cache_resource
def load_model(model_name="yolov8n-seg.pt"):
return YOLO(model_name)
model = load_model()
# -------------------------------
# Input Method
# -------------------------------
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()
# -------------------------------
# Run Segmentation
# -------------------------------
if 'img_array' in locals():
with st.spinner("Running YOLOv8 segmentation..."):
results = model(img_array)
annotated_img = results[0].plot()
# Display images
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)
# Download annotated image
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")
# -------------------------------
# Elbow Curve (robust)
# -------------------------------
st.subheader("Elbow Curve")
try:
from sklearn.cluster import KMeans
# Extract centroids from masks
masks = getattr(results[0], "masks", None)
centroids = []
if masks is not None:
try:
mask_np = masks.data
except AttributeError:
mask_np = masks
# Convert to numpy
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)])
# Fallback if no masks found
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)
# Only compute elbow if >= 2 points
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_)
# Slider to select k
selected_k = st.slider("Select k to highlight on Elbow Curve", 1, max_k, value=3)
# Plot elbow curve
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}") |