File size: 1,578 Bytes
1965b44 | 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 | # streamlit_yolo_app.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 Object Detection in Images")
model = YOLO("yolov8n.pt")
# -------------------------------
# Load YOLO Model
@st.cache_resource
def load_model(model_name):
return YOLO(model_name)
# -------------------------------
# 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
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_container_width=True)
with col2:
st.image(annotated_img, caption="Detected Objects", use_container_width=True)
# Option to download the annotated image
result_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
cv2.imwrite("annotated_output.jpg", result_bgr)
with open("annotated_output.jpg", "rb") as f:
st.download_button("Download Annotated Image", f, "annotated_output.jpg")
|