File size: 1,646 Bytes
0f70bae | 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 | # streamlit_yolo_segmentation.py
import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
# -------------------------------
# Streamlit App Title
# -------------------------------
st.title("YOLOv8 Image Segmentation App")
# -------------------------------
# Load YOLO Segmentation Model
# -------------------------------
@st.cache_resource
def load_model(model_name="yolov8n-seg.pt"):
return YOLO(model_name)
model = load_model()
# -------------------------------
# Upload Image
# -------------------------------
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Load uploaded image
image = Image.open(uploaded_file).convert("RGB")
img_array = np.array(image)
# Perform Inference
results = model(img_array)
# Annotated results with masks
annotated_img = results[0].plot()
# Get shape of annotated image
h, w, c = annotated_img.shape
st.write(f"**Annotated Image Shape:** {h} x {w} x {c}")
# Display Input and Output
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)
# Option to download the annotated image
result_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
cv2.imwrite("segmented_output.jpg", result_bgr)
with open("segmented_output.jpg", "rb") as f:
st.download_button("Download Segmented Image", f, "segmented_output.jpg") |