Srikanth0804/Image_Segmentation_Using_YOLO

1# streamlit_yolo_segmentation.py
2
3import streamlit as st
4import cv2
5import numpy as np
6from ultralytics import YOLO
7from PIL import Image
8
9# -------------------------------
10# Streamlit App Title
11# -------------------------------
12st.title("YOLOv8 Image Segmentation App")
13
14# -------------------------------
15# Load YOLO Segmentation Model
16# -------------------------------
17@st.cache_resource
18def load_model(model_name="yolov8n-seg.pt"):
19    return YOLO(model_name)
20
21model = load_model()
22
23# -------------------------------
24# Upload Image
25# -------------------------------
26uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
27
28if uploaded_file is not None:
29    # Load uploaded image
30    image = Image.open(uploaded_file).convert("RGB")
31    img_array = np.array(image)
32
33    # Perform Inference
34    results = model(img_array)
35
36    # Annotated results with masks
37    annotated_img = results[0].plot()
38
39    # Get shape of annotated image
40    h, w, c = annotated_img.shape
41    st.write(f"**Annotated Image Shape:** {h} x {w} x {c}")
42
43    # Display Input and Output
44    col1, col2 = st.columns(2)
45    with col1:
46        st.image(image, caption="Original Image", use_column_width=True)
47    with col2:
48        st.image(annotated_img, caption="Segmented Image", use_column_width=True)
49
50    # Option to download the annotated image
51    result_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
52    cv2.imwrite("segmented_output.jpg", result_bgr)
53    with open("segmented_output.jpg", "rb") as f:
54        st.download_button("Download Segmented Image", f, "segmented_output.jpg")