ear_shoulder_extract.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ngMAbx2eSVCy",
        "outputId": "dbc6dacc-9e7c-4b98-b6ad-384dcc5fb713"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: mediapipe in /usr/local/lib/python3.10/dist-packages (0.10.9)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from mediapipe) (1.4.0)\n",
            "Requirement already satisfied: attrs>=19.1.0 in /usr/local/lib/python3.10/dist-packages (from mediapipe) (23.1.0)\n",
            "Requirement already satisfied: flatbuffers>=2.0 in /usr/local/lib/python3.10/dist-packages (from mediapipe) (23.5.26)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from mediapipe) (3.7.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from mediapipe) (1.23.5)\n",
            "Requirement already satisfied: opencv-contrib-python in /usr/local/lib/python3.10/dist-packages (from mediapipe) (4.8.0.76)\n",
            "Requirement already satisfied: protobuf<4,>=3.11 in /usr/local/lib/python3.10/dist-packages (from mediapipe) (3.20.3)\n",
            "Requirement already satisfied: sounddevice>=0.4.4 in /usr/local/lib/python3.10/dist-packages (from mediapipe) (0.4.6)\n",
            "Requirement already satisfied: CFFI>=1.0 in /usr/local/lib/python3.10/dist-packages (from sounddevice>=0.4.4->mediapipe) (1.16.0)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (1.2.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (4.46.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (1.4.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (23.2)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (9.4.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (3.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->mediapipe) (2.8.2)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.10/dist-packages (from CFFI>=1.0->sounddevice>=0.4.4->mediapipe) (2.21)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->mediapipe) (1.16.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install mediapipe"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import mediapipe as mp\n",
        "import numpy as np\n",
        "#import csv\n",
        "\n",
        "import json"
      ],
      "metadata": {
        "id": "4lOjpPR0Szfr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mp_pose = mp.solutions.pose\n",
        "pose = mp_pose.Pose()"
      ],
      "metadata": {
        "id": "7ZNAAc86S-A8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def calculate_angle(landmark1, landmark2):\n",
        "    dy = landmark2[1] - landmark1[1]\n",
        "    dx = landmark2[0] - landmark1[0]\n",
        "    angle_radians = np.arctan2(dy, dx)\n",
        "    angle_degrees = np.degrees(angle_radians)\n",
        "    return angle_degrees"
      ],
      "metadata": {
        "id": "gpXIL8OPpWjG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def extract_coordinates(results):\n",
        "    landmarks = results.pose_landmarks.landmark\n",
        "    coordinates = {\n",
        "        \"left_ear\": (landmarks[mp_pose.PoseLandmark.LEFT_EAR.value].x, landmarks[mp_pose.PoseLandmark.LEFT_EAR.value].y),\n",
        "        \"left_shoulder\": (landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x, landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y)\n",
        "    }\n",
        "    return coordinates"
      ],
      "metadata": {
        "id": "b4hcOk1PpZs1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def process_video(video_path, output_path):\n",
        "    cap = cv2.VideoCapture(video_path)\n",
        "    if not cap.isOpened():\n",
        "        print(\"Error: Cannot open video.\")\n",
        "        return\n",
        "\n",
        "    rows = []\n",
        "    while cap.isOpened():\n",
        "        ret, frame = cap.read()\n",
        "        if not ret:\n",
        "            break\n",
        "\n",
        "        image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
        "        results = pose.process(image_rgb)\n",
        "\n",
        "        if results.pose_landmarks:\n",
        "            coordinates = extract_coordinates(results)\n",
        "\n",
        "            # 귀, 어깨 거리 계산\n",
        "            acromion_distance = abs(coordinates['left_ear'][0] - coordinates['left_shoulder'][0])\n",
        "            angle = calculate_angle(coordinates['left_ear'], coordinates['left_shoulder'])\n",
        "\n",
        "            # 거북이 판별\n",
        "            turtleneck_status = \"거북이\" if acromion_distance >= 2.5 or angle >= 50 else \"정상\"\n",
        "            coordinates['turtleneck_status'] = turtleneck_status\n",
        "            rows.append(coordinates)\n",
        "\n",
        "            coordinates['turtleneck_status'] = turtleneck_status\n",
        "            rows.append(coordinates)\n",
        "\n",
        "    cap.release()\n",
        "    save_to_json(output_path, rows)"
      ],
      "metadata": {
        "id": "rI0kfQewTA1J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def extract_coordinates(results):\n",
        "    landmarks = results.pose_landmarks.landmark\n",
        "\n",
        "# left 기준\n",
        "    coordinates = {\n",
        "        \"left_ear\": (landmarks[mp_pose.PoseLandmark.LEFT_EAR.value].x, landmarks[mp_pose.PoseLandmark.LEFT_EAR.value].y),\n",
        "        \"left_shoulder\": (landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x, landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y)\n",
        "    }\n",
        "\n",
        "    return coordinates"
      ],
      "metadata": {
        "id": "4W74-bGPTCUv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def save_to_json(output_path, rows):\n",
        "    with open(output_path, 'w') as file:\n",
        "        json.dump(rows, file, indent=4)"
      ],
      "metadata": {
        "id": "mQxtQnU9TbdO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if __name__ == \"__main__\":\n",
        "    video_path = \"/content/drive/MyDrive/GDSC_AI/SC_test/test_video/turtle_neck_test_video.mp4\"\n",
        "    output_path = \"turtle_neck_results.json\"\n",
        "    process_video(video_path, output_path)"
      ],
      "metadata": {
        "id": "g4mR1y7PTd5b"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}