main code

.gitignore

@@ -0,0 +1,3 @@
+*.pt
+/images
+__pycache__

Dockerfile

@@ -1,3 +1,12 @@
 FROM nvcr.io/nvidia/pytorch:23.12-py3
 
-RUN pip install ultralytics
+RUN pip install ultralytics
+RUN pip install streamlit
+RUN pip install opencv-python==4.6.0.66
+RUN pip install Pillow==10.3.0
+
+RUN apt update && apt install fonts-dejavu
+
+EXPOSE 8501
+CMD streamlit run app.py \
+    --server.headless true

app.py

@@ -0,0 +1,101 @@
+import streamlit as st
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from ultralytics import YOLO
+import torch
+import infer
+
+
+@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 = 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)
+

infer.py

@@ -0,0 +1,74 @@
+import numpy as np
+from ultralytics.engine.results import Results
+
+
+KEYPOINT_NAMES = ["nose","left_eye","right_eye","left_ear","right_ear","left_shoulder",
+                "right_shoulder","left_elbow","right_elbow","left_wrist","right_wrist",
+                "left_hip","right_hip","left_knee","right_knee","left_ankle","right_ankle"]
+
+
+def get_keypoints(result: Results):
+    keypoints = None
+    for i, box in enumerate(result.boxes):
+        if box.cls != 0.:  # Only consider the person class
+            continue
+        person_conf = box.conf.item()
+        k = result.keypoints.data[i]
+        x = k[:, 0].tolist()
+        y = k[:, 1].tolist()
+        score = k[:, 2]
+        visible = (score > 0.5).tolist()
+        # keypoints = {'x': x, 'y': y, 'visible': visible}
+        keypoints = {key_name: (x_, y_) if v_ else None for key_name,x_,y_,v_ in zip(KEYPOINT_NAMES, x, y, visible)}
+        break
+    return keypoints
+
+
+def calculate_angle(p1, p2, p3):
+    v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
+    v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
+    angle_rad = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
+    angle_deg = np.degrees(angle_rad)
+    return angle_deg
+
+
+def calculate_angle_to_horizontal(vector):
+    angle_rad = np.arctan2(vector[1], vector[0])
+    angle_deg = np.degrees(angle_rad)
+    # Adjust the angle to be within -90 to +90 degrees
+    if angle_deg > 90:
+        angle_deg = 180 - angle_deg
+    elif angle_deg < -90:
+        angle_deg = -180 - angle_deg
+    return -angle_deg
+
+
+def extend_line(start, end, extend_factor=3):
+    vector = np.array(end) - np.array(start)
+    length = np.linalg.norm(vector)
+    unit_vector = vector / np.linalg.norm(vector)
+    new_point = end + unit_vector * length * extend_factor
+    new_point = new_point.tolist()
+    return (new_point[0], new_point[1])
+
+
+def get_elbow_angles(keypoints: dict):
+    left_elbow_angle = None
+    right_elbow_angle = None
+    if keypoints['left_shoulder'] and keypoints['left_elbow'] and keypoints['left_wrist']:
+        left_elbow_angle = calculate_angle(keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist'])
+    if keypoints['right_shoulder'] and keypoints['right_elbow'] and keypoints['right_wrist']:
+        right_elbow_angle = calculate_angle(keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist'])
+    return left_elbow_angle, right_elbow_angle
+
+
+def get_eye_angles(keypoints: dict):
+    left_eye_angle = None
+    right_eye_angle = None
+    if keypoints['left_ear'] and keypoints['left_eye']:
+        left_vector = (keypoints['left_eye'][0] - keypoints['left_ear'][0], keypoints['left_eye'][1] - keypoints['left_ear'][1])
+        left_eye_angle = calculate_angle_to_horizontal(left_vector)
+    if keypoints['right_ear'] and keypoints['right_eye']:
+        right_vector = (keypoints['right_eye'][0] - keypoints['right_ear'][0], keypoints['right_eye'][1] - keypoints['right_ear'][1])
+        right_eye_angle = calculate_angle_to_horizontal(right_vector)
+    return left_eye_angle, right_eye_angle

yolov8-test.ipynb

@@ -0,0 +1,146 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ultralytics import YOLO\n",
+    "import torch\n",
+    "from PIL import Image, ImageDraw, ImageFont\n",
+    "import numpy as np\n",
+    "import infer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from importlib import reload\n",
+    "reload(infer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "model_pose = YOLO('yolov8l-pose.pt')\n",
+    "model_pose.to(device)\n",
+    "\n",
+    "model_det = YOLO('yolov8m.pt')\n",
+    "model_det.to(device);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "url = \"image.jpg\"\n",
+    "results = model_pose(url)\n",
+    "results_det = model_det(url)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def draw_output(image_pil: Image.Image, keypoints: dict):    \n",
+    "    draw = ImageDraw.Draw(image_pil)\n",
+    "    line_width = 10\n",
+    "    font = ImageFont.truetype(\"DejaVuSerif-Bold.ttf\", 70)\n",
+    "    \n",
+    "    ear, eye = None, None\n",
+    "    if keypoints[\"left_ear\"] and keypoints[\"left_eye\"]:\n",
+    "        ear = keypoints[\"left_ear\"]\n",
+    "        eye = keypoints[\"left_eye\"]\n",
+    "    elif keypoints[\"right_ear\"] and keypoints[\"right_eye\"]:\n",
+    "        ear = keypoints[\"right_ear\"]\n",
+    "        eye = keypoints[\"right_eye\"]\n",
+    "    \n",
+    "    # draw extended left and right eye lines\n",
+    "    if ear and eye:\n",
+    "        left_new_point = infer.extend_line(ear, eye, 3)\n",
+    "        l1 = [ear, left_new_point]\n",
+    "        draw.line(l1, fill='red', width=line_width)\n",
+    "        # draw a horizontal line from ear forwards\n",
+    "        ear = np.array(ear)\n",
+    "        l1 = np.array(l1)\n",
+    "        l1_vector = l1[1] - l1[0]\n",
+    "        x_s = np.sign(l1_vector)[0]\n",
+    "        length_l1 = np.linalg.norm(l1_vector)\n",
+    "        p2 = ear + np.array([length_l1*x_s, 0])\n",
+    "        ear = tuple(ear.tolist())\n",
+    "        l = [ear, tuple(p2.tolist())]\n",
+    "        draw.line(l, fill='gray', width=line_width//2)\n",
+    "        # draw angle\n",
+    "        angle = infer.calculate_angle_to_horizontal(l1_vector)\n",
+    "        draw.text(ear, f'{angle:.2f}', fill='red', font=font)\n",
+    "        print(infer.get_eye_angles(keypoints))\n",
+    "\n",
+    "\n",
+    "    # draw elbow angles\n",
+    "    left_elbow_angle, right_elbow_angle = infer.get_elbow_angles(keypoints)\n",
+    "    if left_elbow_angle:\n",
+    "        draw.text(keypoints['left_elbow'], f'{left_elbow_angle:.2f}', fill='red', font=font)\n",
+    "        # draw polyline for left arm\n",
+    "        draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=line_width)\n",
+    "    if right_elbow_angle:\n",
+    "        draw.text(keypoints['right_elbow'], f'{right_elbow_angle:.2f}', fill='red', font=font)\n",
+    "        # draw polyline for right arm\n",
+    "        draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=line_width)\n",
+    "\n",
+    "    return image_pil"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "keypoints = infer.get_keypoints(results[0])\n",
+    "img = Image.open(url)\n",
+    "img = draw_output(img, keypoints)\n",
+    "img.resize((800, 800))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}