# app.py

import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image

# Load YOLOv8 segmentation model
model = YOLO("yolov8n-seg.pt")  # you can replace with yolov8s-seg.pt for better accuracy

st.title("🖼️ Image Segmentation using YOLOv8")
st.write("Upload an image and see segmentation results using YOLOv8!")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert uploaded file to OpenCV format
    image = Image.open(uploaded_file).convert("RGB")
    img_np = np.array(image)  # PIL → NumPy
    img_cv = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

    st.subheader("Original Image")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Run YOLO segmentation
    results = model(img_cv)

    # Get annotated image
    for r in results:
        annotated_img = r.plot()  # segmentation + boxes + labels

    # Convert back BGR → RGB for display
    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)

    st.subheader("Segmented Image")
    st.image(annotated_img, caption="YOLOv8 Segmentation", use_container_width=True)

    # Option to download result
    result_pil = Image.fromarray(annotated_img)
    st.download_button(
        label="Download Segmented Image",
        data=cv2.imencode('.jpg', cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR))[1].tobytes(),
        file_name="segmented_output.jpg",
        mime="image/jpeg"
    )