{
 "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
}