|
|
|
|
| import streamlit as st
|
| import cv2
|
| import numpy as np
|
| from ultralytics import YOLO
|
| from PIL import Image
|
|
|
|
|
|
|
|
|
| st.title("YOLOv8 Image Segmentation App")
|
|
|
|
|
|
|
|
|
| @st.cache_resource
|
| def load_model(model_name="yolov8n-seg.pt"):
|
| return YOLO(model_name)
|
|
|
| model = load_model()
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
| results = model(img_array)
|
|
|
|
|
| annotated_img = results[0].plot()
|
|
|
|
|
| h, w, c = annotated_img.shape
|
| st.write(f"**Annotated Image Shape:** {h} x {w} x {c}")
|
|
|
|
|
| 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)
|
|
|
|
|
| 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") |
