Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import numpy as np | |
| from PIL import Image, ImageDraw, ImageFont | |
| from ultralytics import YOLO | |
| import torch | |
| import infer | |
| 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 = infer.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 = infer.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 = infer.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}') | |
| upload_tab, camera_tab = st.tabs(["Upload Photo", "Webcam Capture"]) | |
| with upload_tab: | |
| uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) | |
| with camera_tab: | |
| img_file_buffer = st.camera_input("Take a picture") | |
| img = None | |
| if img_file_buffer is not None: | |
| img = Image.open(img_file_buffer) | |
| if uploaded_file is not None: | |
| img = Image.open(uploaded_file) | |
| if img is not None: | |
| # predict | |
| with st.spinner('Predicting...'): | |
| model = load_model() | |
| pred = model(img)[0] | |
| keypoints = infer.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 = infer.get_eye_angles(keypoints) | |
| lba, rba = infer.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) | |