create_poses_csv.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mediapipe as mp\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp_pose = mp.solutions.pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "landmarks_points = {\n",
    "    \"nose\": 0,\n",
    "    \"left_shoulder\": 11, \"right_shoulder\": 12,\n",
    "    \"left_elbow\": 13, \"right_elbow\": 14,\n",
    "    \"left_wrist\": 15, \"right_wrist\": 16,\n",
    "    \"left_hip\": 23, \"right_hip\": 24,\n",
    "    \"left_knee\": 25, \"right_knee\": 26,\n",
    "    \"left_ankle\": 27, \"right_ankle\": 28,\n",
    "    \"left_heel\": 29, \"right_heel\": 30,\n",
    "    \"left_foot_index\": 31, \"right_foot_index\": 32,\n",
    "}\n",
    "\n",
    "col_names = []\n",
    "for i in range(len(landmarks_points.keys())):\n",
    "    name = list(landmarks_points.keys())[i]\n",
    "    col_names.append(name + \"_x\")\n",
    "    col_names.append(name + \"_y\")\n",
    "    col_names.append(name + \"_z\")\n",
    "    col_names.append(name + \"_v\")\n",
    "\n",
    "pose_name = col_names.copy()\n",
    "\n",
    "pose_name.append(\"pose\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_list = []\n",
    "\n",
    "main_dir = \"./poses_dataset/Images\"\n",
    "pose_dir_list = os.listdir(main_dir)\n",
    "\n",
    "for i in range(0, len(pose_dir_list)):\n",
    "    images_dir_list = os.listdir(f\"{main_dir}/{pose_dir_list[i]}\")\n",
    "    for l in range(0, len(images_dir_list)):\n",
    "        pre_list = []\n",
    "        with mp_pose.Pose(static_image_mode=True, enable_segmentation=True) as pose:\n",
    "            image = cv2.imread(\n",
    "                f\"{main_dir}/{pose_dir_list[i]}/{images_dir_list[l]}\")\n",
    "            result = pose.process(\n",
    "                cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "            try:\n",
    "                predict = True\n",
    "                for landmarks in result.pose_landmarks.landmark:\n",
    "                    pre_list.append(landmarks)\n",
    "            except AttributeError:\n",
    "                print(\n",
    "                    f\"No points {main_dir}/{pose_dir_list[i]}/{images_dir_list[l]}\")\n",
    "                predict = False\n",
    "\n",
    "        if predict == True:\n",
    "            gen1116 = np.array([\n",
    "                [\n",
    "                    pre_list[m].x,\n",
    "                    pre_list[m].y,\n",
    "                    pre_list[m].z,\n",
    "                    pre_list[m].visibility\n",
    "                ] for m in range(11, 17)\n",
    "            ]).flatten().tolist()\n",
    "\n",
    "            gen2333 = np.array([\n",
    "                [\n",
    "                    pre_list[m].x,\n",
    "                    pre_list[m].y,\n",
    "                    pre_list[m].z,\n",
    "                    pre_list[m].visibility\n",
    "                ] for m in range(23, 33)\n",
    "            ]).flatten().tolist()\n",
    "\n",
    "            gen1116.extend(gen2333)\n",
    "\n",
    "            all_list = [\n",
    "                pre_list[0].x,\n",
    "                pre_list[0].y,\n",
    "                pre_list[0].z,\n",
    "                pre_list[0].visibility,\n",
    "            ]\n",
    "\n",
    "            all_list.extend(gen1116)\n",
    "            tpl = all_list.copy()\n",
    "            tpl.append(i)\n",
    "            pose_list.append(tpl)\n",
    "\n",
    "data_pose = pd.DataFrame(pose_list, columns=pose_name)\n",
    "data_pose.to_csv(\"./csv_files/poses_data_pose.csv\", index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}