| |
|
|
| import streamlit as st |
| import cv2 |
| import numpy as np |
| from ultralytics import YOLO |
| from PIL import Image |
|
|
| |
| |
| |
| st.title("YOLOv8 Pose Detection App") |
|
|
| |
| |
| |
| @st.cache_resource |
| def load_model(model_name="yolov8n-pose.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="Pose Detected", use_column_width=True) |
|
|
| |
| result_bgr = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR) |
| cv2.imwrite("pose_output.jpg", result_bgr) |
| with open("pose_output.jpg", "rb") as f: |
| st.download_button("Download Pose Image", f, "pose_output.jpg") |
