pszmk-model-research.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "KEYPOINTS_NAMES = [\n",
    "    \"0 - nose\",\n",
    "    \"1 - left eye (inner)\",\n",
    "    \"2 - left eye\",\n",
    "    \"3 - left eye (outer)\",\n",
    "    \"4 - right eye (inner)\",\n",
    "    \"5 - right eye\",\n",
    "    \"6 - right eye (outer)\",\n",
    "    \"7 - left ear\",\n",
    "    \"8 - right ear\",\n",
    "    \"9 - mouth (left)\",\n",
    "    \"10 - mouth (right)\",\n",
    "    \"11 - left shoulder\",\n",
    "    \"12 - right shoulder\",\n",
    "    \"13 - left elbow\",\n",
    "    \"14 - right elbow\",\n",
    "    \"15 - left wrist\",\n",
    "    \"16 - right wrist\",\n",
    "    \"17 - left pinky\",\n",
    "    \"18 - right pinky\",\n",
    "    \"19 - left index\",\n",
    "    \"20 - right index\",\n",
    "    \"21 - left thumb\",\n",
    "    \"22 - right thumb\",\n",
    "    \"23 - left hip\",\n",
    "    \"24 - right hip\",\n",
    "    \"25 - left knee\",\n",
    "    \"26 - right knee\",\n",
    "    \"27 - left ankle\",\n",
    "    \"28 - right ankle\",\n",
    "    \"29 - left heel\",\n",
    "    \"30 - right heel\",\n",
    "    \"31 - left foot index\",\n",
    "    \"32 - right foot index\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "import mediapipe as mp\n",
    "import cv2\n",
    "import random\n",
    "from PIL import Image\n",
    "from IPython.display import display\n",
    "from typing import Tuple\n",
    "\n",
    "current_working_dir = os.getcwd()\n",
    "GENERAL_DATA_PATH = Path(current_working_dir) / 'data'\n",
    "DATA_PATH = GENERAL_DATA_PATH / 'yogaposes-aii22-challenge'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_numpy(image_file_name : str) -> np.ndarray:\n",
    "    return np.array(Image.open(DATA_PATH / 'Train' / image_file_name))\n",
    "\n",
    "def get_random_image_numpy() -> np.ndarray:\n",
    "    images = os.listdir(DATA_PATH / 'Train')\n",
    "    random_image_file_name = random.choice(images)\n",
    "    return get_image_numpy(random_image_file_name)\n",
    "\n",
    "class PoseExtractor():\n",
    "    keypoints_names = KEYPOINTS_NAMES\n",
    "    def __init__(self,\n",
    "                 source_data_path : str,\n",
    "                 destination_data_path : str,\n",
    "                 save_file_name : str,\n",
    "                 model_complexity : int,\n",
    "                 n_samples : int\n",
    "                 ):\n",
    "        self.n_samples = n_samples\n",
    "        self.source_data_path = source_data_path\n",
    "        self.destination_data_path = destination_data_path\n",
    "        self.save_file_name = save_file_name\n",
    "        self.pose = mp.solutions.pose.Pose(\n",
    "            static_image_mode=True,\n",
    "            model_complexity=model_complexity,\n",
    "            smooth_landmarks=True,\n",
    "            enable_segmentation=False,\n",
    "            smooth_segmentation=False,\n",
    "            min_detection_confidence=0.5,\n",
    "            min_tracking_confidence=0.5\n",
    "        )\n",
    "    \n",
    "    def extract_pose(self,\n",
    "                     image: np.ndarray\n",
    "                     ) -> Tuple:\n",
    "        return np.hstack([np.array([landmark.x, landmark.y, landmark.z, landmark.visibility], dtype=np.float64) for landmark in self.pose.process(image).pose_landmarks.landmark])\n",
    "    \n",
    "    def extract_poses(self):\n",
    "        for img in os.listdir(self.source_data_path):\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(132,)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pose_extractor = PoseExtractor(\n",
    "    DATA_PATH,\n",
    "    GENERAL_DATA_PATH,\n",
    "    'train-poses.csv',\n",
    "    1\n",
    ")\n",
    "\n",
    "image = get_image_numpy('01300.jpg')\n",
    "landmarks = pose_extractor.extract_pose(image)\n",
    "landmarks.shape\n",
    "# for landmark in landmarks:\n",
    "#     print(landmark)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import mediapipe as mp\n",
    "import numpy as np\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset\n",
    "from PIL import Image\n",
    "from typing import Tuple, List, Union\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data_dir: str):\n",
    "        self.data_dir = data_dir\n",
    "        self.image_files = os.listdir(data_dir)\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "        ])\n",
    "\n",
    "    def get_label(self, img_name: str) -> int:\n",
    "        # Extract the label from the image name (assuming it's the first digit)\n",
    "        label = int(img_name[0])  # Convert the first character to an integer\n",
    "        return label\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.image_files)\n",
    "\n",
    "    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:\n",
    "        img_name = os.path.join(self.data_dir, self.image_files[idx])\n",
    "        image = cv2.imread(img_name)  # Read image using OpenCV\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB\n",
    "\n",
    "        tensor_image = self.transform(image)\n",
    "        label = self.get_label(self.image_files[idx])\n",
    "\n",
    "        # Save the preprocessed image to the \"NewTrain\" folder\n",
    "        save_path = os.path.join(\"NewTrain\", f\"preprocessed_{self.image_files[idx]}\")\n",
    "        if tensor_image.shape[0] == 1:\n",
    "            # If it's a single-channel image, convert it to three channels before saving\n",
    "            tensor_image_rgb = torch.cat([tensor_image] * 3, dim=0)\n",
    "            cv2.imwrite(save_path, cv2.cvtColor(tensor_image_rgb.numpy().transpose(1, 2, 0), cv2.COLOR_RGB2BGR))\n",
    "        else:\n",
    "            # If it's already a three-channel image, save as is\n",
    "            cv2.imwrite(save_path, cv2.cvtColor(tensor_image.numpy().transpose(1, 2, 0), cv2.COLOR_RGB2BGR))\n",
    "\n",
    "\n",
    "        return tensor_image, label\n",
    "\n",
    "# Example usage:\n",
    "data_directory = \"Train\"\n",
    "custom_dataset = CustomDataset(data_directory)\n",
    "\n",
    "# Example loading one image\n",
    "sample_image, label = custom_dataset[2]\n",
    "file = os.listdir(data_directory)\n",
    "print(file[2])\n",
    "print(sample_image.shape)  # Check the shape (should be torch.Size([1, 224, 224]))\n",
    "print(\"Label:\", label)\n",
    "print(sample_image)\n",
    "\n",
    "min_value = torch.min(sample_image)\n",
    "max_value = torch.max(sample_image)\n",
    "\n",
    "print(\"Minimum value:\", min_value.item())\n",
    "print(\"Maximum value:\", max_value.item())\n",
    "\n",
    "# Convert the tensor to a NumPy array\n",
    "numpy_image = sample_image.squeeze().numpy()\n",
    "\n",
    "# Display the image\n",
    "plt.imshow(numpy_image, cmap='gray')  # Use 'gray' colormap for single-channel images\n",
    "plt.title(\"Image Title\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Pose.__init__() got an unexpected keyword argument 'num_poses'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[25], line 3\u001b[0m\n\u001b[0;32m      1\u001b[0m image_np \u001b[38;5;241m=\u001b[39m get_random_image_numpy()\n\u001b[1;32m----> 3\u001b[0m pose \u001b[38;5;241m=\u001b[39m \u001b[43mmp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msolutions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpose\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mPose\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m      4\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_poses\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      5\u001b[0m \u001b[43m            \u001b[49m\u001b[43mmin_detection_confidence\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.5\u001b[39;49m\n\u001b[0;32m      6\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      8\u001b[0m results \u001b[38;5;241m=\u001b[39m pose\u001b[38;5;241m.\u001b[39mprocess(image_np)\u001b[38;5;66;03m#cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\u001b[39;00m\n\u001b[0;32m     10\u001b[0m keypoints \u001b[38;5;241m=\u001b[39m results\u001b[38;5;241m.\u001b[39mpose_landmarks\u001b[38;5;241m.\u001b[39mlandmark  \u001b[38;5;66;03m# List of normalized keypoints (x, y, z)\u001b[39;00m\n",
      "\u001b[1;31mTypeError\u001b[0m: Pose.__init__() got an unexpected keyword argument 'num_poses'"
     ]
    }
   ],
   "source": [
    "image_np = get_random_image_numpy()\n",
    "\n",
    "pose = mp.solutions.pose.Pose(\n",
    "            num_poses=1,\n",
    "            min_detection_confidence=0.5\n",
    "            )\n",
    "\n",
    "results = pose.process(image_np)#cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "keypoints = results.pose_landmarks.landmark  # List of normalized keypoints (x, y, z)\n",
    "\n",
    "feature_vector = np.array([kp.x for kp in keypoints] + [kp.y for kp in keypoints] + [kp.z for kp in keypoints])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "x: 0.2824716\n",
       "y: 0.43324316\n",
       "z: 0.058916435\n",
       "visibility: 0.9987006"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keypoints[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "99"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(feature_vector)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pose-estimation-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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}