Spaces:
Sleeping
Sleeping
File size: 1,841 Bytes
c66f90e fadf95f c66f90e 67fd17e c66f90e fadf95f c66f90e fadf95f c66f90e a7a43e9 c66f90e a7a43e9 c66f90e a7a43e9 c66f90e a7a43e9 c66f90e fadf95f a7a43e9 fadf95f c66f90e fadf95f 67fd17e fadf95f c66f90e fadf95f c66f90e |
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 55 56 57 58 59 60 61 62 63 64 |
import streamlit as st
from PIL import Image
from ultralytics import YOLO
import torch
import utils
import utils
from drawing import draw_keypoints
@st.cache_resource()
def load_model():
print('Loading model...')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_pose = YOLO('yolov8l-pose.pt')
model_pose.to(device)
return model_pose
def draw_output(image_pil: Image.Image, keypoints: dict):
output_image = draw_keypoints(image_pil, keypoints).convert("RGB")
return output_image
st.title('Pose Estimation App')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
st.caption(f'Using device: {device}')
mode = st.radio('Select mode:', ['Upload an Image', 'Webcam Capture'])
if mode == 'Upload an Image':
img_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
elif mode == 'Webcam Capture':
img_file = st.camera_input("Take a picture")
img = None
if img_file is not None:
img = Image.open(img_file)
st.divider()
if img is not None:
# predict
with st.spinner('Predicting...'):
model = load_model()
result = model(img)[0]
st.markdown('**Results:**')
keypoints = utils.get_keypoints(result)
if keypoints is not None:
img = draw_output(img, keypoints)
st.image(img, caption='Predicted image', use_column_width=True)
# calculate angles
lea, rea = utils.get_eye_angles(keypoints)
lba, rba = utils.get_elbow_angles(keypoints)
angles = {'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba}
st.write('Angles:')
st.json(angles)
st.write('Raw keypoints:')
st.json(keypoints)
else:
st.error('No keypoints detected!')
st.image(img, caption='Original image', use_column_width=True)
|