# 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}")