File size: 1,568 Bytes
a432cb7 | 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 | # 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"
)
|