Spaces:
Sleeping
Sleeping
File size: 3,543 Bytes
c66f90e 67fd17e c66f90e 67fd17e c66f90e 67fd17e c66f90e 67fd17e c66f90e a7a43e9 c66f90e a7a43e9 c66f90e a7a43e9 c66f90e a7a43e9 c66f90e a7a43e9 67fd17e c66f90e 67fd17e 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import streamlit as st
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import torch
import utils
@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):
draw = ImageDraw.Draw(image_pil)
line_width = 10
font = ImageFont.truetype("DejaVuSerif-Bold.ttf", 70)
ear, eye = None, None
if keypoints["left_ear"] and keypoints["left_eye"]:
ear = keypoints["left_ear"]
eye = keypoints["left_eye"]
elif keypoints["right_ear"] and keypoints["right_eye"]:
ear = keypoints["right_ear"]
eye = keypoints["right_eye"]
# draw extended left and right eye lines
if ear and eye:
left_new_point = utils.extend_line(ear, eye, 3)
l1 = [ear, left_new_point]
draw.line(l1, fill='red', width=line_width)
# draw a horizontal line from ear forwards
ear = np.array(ear)
l1 = np.array(l1)
l1_vector = l1[1] - l1[0]
x_s = np.sign(l1_vector)[0]
length_l1 = np.linalg.norm(l1_vector)
p2 = ear + np.array([length_l1*x_s, 0])
ear = tuple(ear.tolist())
l = [ear, tuple(p2.tolist())]
draw.line(l, fill='gray', width=line_width//2)
# draw angle
angle = utils.calculate_angle_to_horizontal(l1_vector)
draw.text(ear, f'{angle:.2f}', fill='red', font=font)
# draw elbow angles
left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
if left_elbow_angle:
draw.text(keypoints['left_elbow'], f'{left_elbow_angle:.2f}', fill='red', font=font)
# draw polyline for left arm
draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=line_width)
if right_elbow_angle:
draw.text(keypoints['right_elbow'], f'{right_elbow_angle:.2f}', fill='red', font=font)
# draw polyline for right arm
draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=line_width)
return image_pil
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()
pred = model(img)[0]
st.markdown('**Results:**')
keypoints = utils.get_keypoints(pred)
if keypoints is not None:
img = draw_output(img, keypoints)
st.image(img, caption='Predicted image', use_column_width=True)
lea, rea = utils.get_eye_angles(keypoints)
lba, rba = utils.get_elbow_angles(keypoints)
st.write('Angles:')
st.json({'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba})
st.write('Raw keypoints:')
st.json(keypoints)
else:
st.error('No keypoints detected!')
st.image(img, caption='Original image', use_column_width=True)
|