main.py

"""
Motion Learn

User can use their right hand to point to 3 distinct landmarks
(Right ear for head, Right shoulder for shoulder, near heart for heart)

After that, the information about the body part will show up in the window

Reference:
https://www.youtube.com/watch?v=06TE_U21FK4

Initialize git later:
git init
git add .
git commit -m "Initial commit"
git remote add origin <repo url>
git pull origin main
git push -u origin main

//stage changes
git add .
git commit -m "fesfusoie"
git push origin main
"""

import time
import numpy as np
import cv2
import mediapipe as mp

from body_info import BODY_PART_INFO

mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

def calculate_distance(p1, p2):
    """
    Calculates the distance between 2 pose landmarks
    :input: 2 arrays with [x,y] coordinates
    :process: Pythagoras theorem to calculate the distance
            sqrt[(x1-x2)^2 + (y1-y2)^2]
    :output: The distance in float
    """
    return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)

def draw_info_box(image, info_text, side_image=None):
    """
    Draws a text box with an optional image on the side.
    
    :param image: The main video frame (BGR numpy array).
    :param info_text: A list of strings to display (Title, Line 1, Line 2).
    :param side_image: (Optional) The image to display on the right (BGR numpy array).
    """
    # Setup Box Dimensions
    h, w, _ = image.shape
    box_height = 160
    # Create a semi-transparent overlay
    overlay = image.copy()
    
    # Draw a black rectangle at the bottom of the screen
    cv2.rectangle(overlay, (0, h - box_height), (w, h), (0, 0, 0), -1)
    
    # Blend the overlay with the original image (0.7 opacity)
    alpha = 0.65
    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)

    # Display text
    start_y = h - box_height + 50
    for i, line in enumerate(info_text):
        # The title (first line) is Yellow, others are White
        color = (0, 204, 204) if i == 0 else (255, 255, 255)
        font_scale = 0.9 if i == 0 else 0.65
        thickness = 2 if i == 0 else 1
        font = cv2.FONT_HERSHEY_COMPLEX #if i == 0 else cv2.FONT_HERSHEY_COMPLEX_SMALL

        cv2.putText(image, line, (20, start_y + (i * 35)), 
            font, font_scale, color, thickness, cv2.LINE_AA)
        
    # Display image
    if side_image is not None:
        try:
            # 1. Resize side_image to fit the box height (square aspect ratio)
            # Subtract 20px for padding so it doesn't touch the edges
            img_size = box_height - 20
            resized_icon = cv2.resize(side_image, (img_size, img_size))
            
            # 2. Calculate coordinates (Bottom Right with padding)
            y_offset = h - box_height + 10 # 10px padding from bottom/top of box
            x_offset = w - img_size - 10   # 10px padding from right side of screen

            # 3. Overlay the image
            # We replace the pixels in the main image with the icon pixels
            image[y_offset:y_offset+img_size, x_offset:x_offset+img_size] = resized_icon

        except Exception as e:
            print(f"Could not draw side image: {e}")

# Video Feed setup
webcam = cv2.VideoCapture(2)    # Default webcam index is 0, 2 is for IV cam (1280x960)
with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
    # Look up on config options for confidence, can try to play around with confidence number
    last_check_time = 0
    check_interval = 1.0  # Seconds
    active_body_part = None

    loaded_images = {}
    for key in BODY_PART_INFO:
        # Get images & use OpenCV to load the image file into memory
        image_path = BODY_PART_INFO[key]["image"]
        img = cv2.imread(image_path)
        
        if img is not None:
            loaded_images[key] = img
        else:
            print(f"ERROR: Could not find image at {image_path}")

    while webcam.isOpened():
        ret, frame = webcam.read()
        current_time = time.time()

        # Detect posture and do calculations in the future
        # Recolor image from opencv (BGR) -> RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image.flags.writeable = False

        # Detection of pose
        results = pose.process(image)

        # Change back image color back to BGR for opencv
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        if current_time - last_check_time > check_interval:
            active_body_part = None

            # Extract landmarks
            try:
                landmarks = results.pose_landmarks.landmark

                # 1. Get Coordinates (normalized 0.0 to 1.0)
                # Right Index (The "Pointer")
                right_index = [
                    landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].x,
                    landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].y,
                ]

                # Right Ear (Head Target)
                right_ear = [landmarks[mp_pose.PoseLandmark.RIGHT_EAR.value].x,
                    landmarks[mp_pose.PoseLandmark.RIGHT_EAR.value].y,]

                # Right Shoulder (Right Shoulder Target)
                rs = [landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y,]

                # Left Shoulder (For Heart Midpoint)
                ls = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y,]

                # 2. Calculate Heart Midpoint (Average of shoulders)
                heart = [(rs[0] + ls[0]) / 2, (rs[1] + ls[1]) / 2]

                # 3. Calculate Distances
                dist_head = calculate_distance(right_index, right_ear)
                dist_shoulder = calculate_distance(right_index, rs)
                dist_heart = calculate_distance(right_index, heart)

                # Debugging use
                # print(f"Distance to Ear: {dist_head:.2f}")
                # print(f"Distance to Shoulder: {dist_shoulder:.4f}")
                # print(f"Distance to HeaRT: {dist_heart:.3f}")

                if dist_head < 0.1:
                    # print("TOUCHING: HEAD")
                    active_body_part = "Head"

                elif dist_shoulder < 0.17:
                    # print("TOUCHING: SHOULDER")
                    active_body_part = "Shoulder"

                elif dist_heart < 0.2:
                    # print("TOUCHING: HEART")
                    active_body_part = "Heart"

                else:
                    # print("==System Ready==")
                    active_body_part = None

                # Update the timer after a check
                last_check_time = current_time

            except Exception as e:
                pass

        # Render out detections
        mp_drawing.draw_landmarks(
            image,
            results.pose_landmarks,
            mp_pose.POSE_CONNECTIONS,
            mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
            mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2, circle_radius=2),
            # Color is in BGR Code
        )

        if active_body_part:
            # Get text
            text_data = BODY_PART_INFO[active_body_part]["text"]
            image_data = BODY_PART_INFO[active_body_part]["image"]
            
            # Get image
            img_data = loaded_images.get(active_body_part)
            
            # Call draw box function
            draw_info_box(image, text_data, img_data)

        # Exit the video feed by pressing 'q'
        cv2.imshow("Motion Learn live feed, Press 'q' to quit", image)
        if cv2.waitKey(10) & 0xFF == ord("q"):
            break

    webcam.release()
    cv2.destroyAllWindows()